update toolkit

agent.py

@@ -29,7 +29,12 @@ import tempfile
 import requests
 from urllib.parse import urlparse
 from typing import Optional
-
+import io
+import contextlib
+import base64
+import subprocess
+import sqlite3
+import traceback
 
 load_dotenv()
 
@@ -290,9 +295,137 @@ def analyze_csv_file(file_path: str, query: str) -> str:
 
     except Exception as e:
         return f"Error analyzing CSV file: {str(e)}"
+
+
+def execute_code_multilang(code: str, language: str = "python") -> str:
+    """
+    Execute code in Python, Bash, SQL, C, or Java and return formatted results.
+
+    Args:
+        code (str): Source code.
+        language (str): Language of the code. One of: 'python', 'bash', 'sql', 'c', 'java'.
+
+    Returns:
+        str: Human-readable execution result.
+    """
+    language = language.lower()
+    exec_id = str(uuid.uuid4())
+    result = {
+        "stdout": "",
+        "stderr": "",
+        "status": "error",
+        "plots": [],
+        "dataframes": [],
+    }
+
+    try:
+        if language == "python":
+            plt.switch_backend("Agg")
+            stdout_buffer = io.StringIO()
+            stderr_buffer = io.StringIO()
+            globals_dict = {"pd": pd, "plt": plt, "Image": Image}
+
+            with contextlib.redirect_stdout(stdout_buffer), contextlib.redirect_stderr(stderr_buffer):
+                exec(code, globals_dict)
+
+            # Save plots
+            if plt.get_fignums():
+                for i, fig_num in enumerate(plt.get_fignums()):
+                    fig = plt.figure(fig_num)
+                    img_path = os.path.join(tempfile.gettempdir(), f"{exec_id}_plot_{i}.png")
+                    fig.savefig(img_path)
+                    with open(img_path, "rb") as f:
+                        img_data = base64.b64encode(f.read()).decode()
+                        result["plots"].append(img_data)
+
+            # Check for dataframes
+            for var_name, var_val in globals_dict.items():
+                if isinstance(var_val, pd.DataFrame):
+                    result["dataframes"].append((var_name, var_val.head().to_string()))
+
+            result["stdout"] = stdout_buffer.getvalue()
+            result["stderr"] = stderr_buffer.getvalue()
+            result["status"] = "success"
+
+        elif language == "bash":
+            completed = subprocess.run(code, shell=True, capture_output=True, text=True, timeout=30)
+            result["stdout"] = completed.stdout
+            result["stderr"] = completed.stderr
+            result["status"] = "success" if completed.returncode == 0 else "error"
+
+        elif language == "sql":
+            conn = sqlite3.connect(":memory:")
+            cur = conn.cursor()
+            cur.execute(code)
+            if code.strip().lower().startswith("select"):
+                cols = [desc[0] for desc in cur.description]
+                rows = cur.fetchall()
+                df = pd.DataFrame(rows, columns=cols)
+                result["dataframes"].append(("query_result", df.head().to_string()))
+            conn.commit()
+            conn.close()
+            result["status"] = "success"
+            result["stdout"] = "SQL executed successfully."
+
+        elif language == "c":
+            with tempfile.TemporaryDirectory() as tmp:
+                src = os.path.join(tmp, "main.c")
+                bin_path = os.path.join(tmp, "main")
+                with open(src, "w") as f:
+                    f.write(code)
+                comp = subprocess.run(["gcc", src, "-o", bin_path], capture_output=True, text=True)
+                if comp.returncode != 0:
+                    result["stderr"] = comp.stderr
+                else:
+                    run = subprocess.run([bin_path], capture_output=True, text=True, timeout=30)
+                    result["stdout"] = run.stdout
+                    result["stderr"] = run.stderr
+                    result["status"] = "success" if run.returncode == 0 else "error"
+
+        elif language == "java":
+            with tempfile.TemporaryDirectory() as tmp:
+                src = os.path.join(tmp, "Main.java")
+                with open(src, "w") as f:
+                    f.write(code)
+                comp = subprocess.run(["javac", src], capture_output=True, text=True)
+                if comp.returncode != 0:
+                    result["stderr"] = comp.stderr
+                else:
+                    run = subprocess.run(["java", "-cp", tmp, "Main"], capture_output=True, text=True, timeout=30)
+                    result["stdout"] = run.stdout
+                    result["stderr"] = run.stderr
+                    result["status"] = "success" if run.returncode == 0 else "error"
+
+        else:
+            return f"❌ Unsupported language: {language}."
+
+    except Exception as e:
+        result["stderr"] = traceback.format_exc()
+
+    # Format response
+    summary = []
+    if result["status"] == "success":
+        summary.append(f"✅ Code executed successfully in **{language.upper()}**")
+        if result["stdout"]:
+            summary.append(f"\n**Output:**\n```\n{result['stdout'].strip()}\n```")
+        if result["stderr"]:
+            summary.append(f"\n**Warnings/Errors:**\n```\n{result['stderr'].strip()}\n```")
+        for name, df in result["dataframes"]:
+            summary.append(f"\n**DataFrame `{name}` Preview:**\n```\n{df}\n```")
+        if result["plots"]:
+            summary.append(f"\n📊 {len(result['plots'])} plot(s) generated (base64-encoded).")
+    else:
+        summary.append(f"❌ Execution failed for **{language.upper()}**")
+        if result["stderr"]:
+            summary.append(f"\n**Error:**\n```\n{result['stderr'].strip()}\n```")
+
+    return "\n".join(summary)
+
+
 tools = [multiply, add, subtract, divide, modulus, 
          wiki_search, web_search, arvix_search, read_excel_file, extract_text_from_pdf,
-         blip_image_caption, save_and_read_file, download_file_from_url, analyze_csv_file]
+         blip_image_caption, save_and_read_file, download_file_from_url, analyze_csv_file,
+         execute_code_multilang]
 
 # ------------------- SYSTEM PROMPT -------------------
 system_prompt_path = "system_prompt.txt"