Update agent.py

agent.py

@@ -18,11 +18,15 @@ from supabase.client import Client, create_client
 from langchain_openai import ChatOpenAI
 
 from langchain.tools import Tool
+from code_interpreter import CodeInterpreter
+
+interpreter_instance = CodeInterpreter()
 
 
 load_dotenv()
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
+### ======================================== MATHEMATICAL TOOLS ======================================== ###
 def multiply(a: int, b: int) -> int:
     return a * b
 
@@ -70,6 +74,29 @@ modulus_tool = Tool(
     description="Modulus two numbers. Args (a: first int, b: second int)"
 )
 
+def power(a: float, b: float) -> float:
+    return a**b
+
+power_tool = Tool(
+    name="power",
+    func=power,
+    description="Power two numbers. Args (a: first float, b: second float)"
+)
+
+def square_root(a: float) -> float | complex:
+    if a >= 0:
+        return a**0.5
+    return cmath.sqrt(a)
+
+square_root_power = Tool(
+    name="square_root",
+    func=square_root,
+    description="Square two numbers. Args (a: float)"
+)
+
+
+### ======================================== BROWSER TOOLS ======================================== ###
+
 def wiki_search(query: str) -> str:
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
     formatted_search_docs = "\n\n---\n\n".join(
@@ -120,6 +147,81 @@ arvix_search_tool = Tool(
 )
 
 
+### ======================================== CODE INTERPRETER TOOLS ======================================== ###
+def execute_code_multilang(code: str, language: str = "python") -> str:
+    """Execute code in multiple languages (Python, Bash, SQL, C, Java) and return results.
+    Args:
+        code (str): The source code to execute.
+        language (str): The language of the code. Supported: "python", "bash", "sql", "c", "java".
+    Returns:
+        A string summarizing the execution results (stdout, stderr, errors, plots, dataframes if any).
+    """
+    supported_languages = ["python", "bash", "sql", "c", "java"]
+    language = language.lower()
+
+    if language not in supported_languages:
+        return f"❌ Unsupported language: {language}. Supported languages are: {', '.join(supported_languages)}"
+
+    result = interpreter_instance.execute_code(code, language=language)
+
+    response = []
+
+    if result["status"] == "success":
+        response.append(f"✅ Code executed successfully in **{language.upper()}**")
+
+        if result.get("stdout"):
+            response.append(
+                "\n**Standard Output:**\n```\n" + result["stdout"].strip() + "\n```"
+            )
+
+        if result.get("stderr"):
+            response.append(
+                "\n**Standard Error (if any):**\n```\n"
+                + result["stderr"].strip()
+                + "\n```"
+            )
+
+        if result.get("result") is not None:
+            response.append(
+                "\n**Execution Result:**\n```\n"
+                + str(result["result"]).strip()
+                + "\n```"
+            )
+
+        if result.get("dataframes"):
+            for df_info in result["dataframes"]:
+                response.append(
+                    f"\n**DataFrame `{df_info['name']}` (Shape: {df_info['shape']})**"
+                )
+                df_preview = pd.DataFrame(df_info["head"])
+                response.append("First 5 rows:\n```\n" + str(df_preview) + "\n```")
+
+        if result.get("plots"):
+            response.append(
+                f"\n**Generated {len(result['plots'])} plot(s)** (Image data returned separately)"
+            )
+
+    else:
+        response.append(f"❌ Code execution failed in **{language.upper()}**")
+        if result.get("stderr"):
+            response.append(
+                "\n**Error Log:**\n```\n" + result["stderr"].strip() + "\n```"
+            )
+
+    return "\n".join(response)
+
+execute_code_multilang_tool = Tool(
+    name="execute_code_multilang",
+    func=execute_code_multilang,
+    description="""Execute code in multiple languages (Python, Bash, SQL, C, Java) and return results.
+        Args:
+            code (str): The source code to execute.
+            language (str): The language of the code. Supported: "python", "bash", "sql", "c", "java".
+    """
+)
+
+
+
 # load the system prompt from the file
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
@@ -133,6 +235,8 @@ tools = [
     substract_tool,
     divide_tool,
     modulus_tool,
+    power_tool,
+    square_root,
     wiki_search_tool,
     web_search_tool,
     arvix_search_tool,