Implementing agent tools and logic

agent.py

@@ -1,4 +1,13 @@
 from llama_index.llms.google_genai import GoogleGenAI
+from llama_index.tools.arxiv import ArxivToolSpec
+from llama_index.tools.wikipedia import WikipediaToolSpec
+from llama_index.tools.duckduckgo import DuckDuckGoSearchResultsToolSpec
+from llama_index.core.tools import FunctionTool
+from llama_index.core.agent.workflow import AgentWorkflow
+
+from tools import interpret_python_math_code
+from gaia_system_prompt import GAIA_SYSTEM_PROMPT
+
 import os
 
 GEMINI_API_KEY = os.getenv("GEMINI_TOKEN")
@@ -6,11 +15,36 @@ GEMINI_MODEL_NAME = "gemini-2.5-flash-preview-04-17"
 
 class FinalAgent:
     def __init__(self):
+        # LLM Initialization
         self.llm = GoogleGenAI(model=GEMINI_MODEL_NAME, api_key=GEMINI_API_KEY)
 
+        # Tool Initialization
+        self.tools = [
+            FunctionTool.from_defaults(
+                func=interpret_python_math_code,
+                name="InterpretPythonMathCode",
+                description="Interprets Python code for mathematical expressions."
+            ),
+            DuckDuckGoSearchResultsToolSpec(),
+            WikipediaToolSpec(),
+            ArxivToolSpec()
+        ]
+
+        # Agent Workflow Initialization
+        self.agent = AgentWorkflow(
+            llm=self.llm,
+            tools=self.tools,
+            system_prompt=GAIA_SYSTEM_PROMPT
+        )
+
         print("FinalAgent initialized.")
     def __call__(self, question: str) -> str:
+        # Example
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         fixed_answer = "This is a default answer."
         print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer
+
+        # Implement agent logic here
+        response = self.agent.run(question)
+        
+        return response

requirements.txt

@@ -2,7 +2,11 @@ gradio[oauth]
 requests
 numpy
 pandas
+scipy
 llama-index
 llama-index-llms-gemini
 llama-index-llms-google-genai
-llama-index-utils-workflow
+llama-index-utils-workflow
+llama-index-tools-duckduckgo
+llama-index-tools-arxiv
+llama-index-tools-wikipedia

tools.py

@@ -0,0 +1,124 @@
+import os
+import ast
+import io
+import sys
+import numpy as np
+import pandas as pd
+import scipy
+
+ALLOWED_MODULES = {"numpy", "pandas", "scipy"}
+
+def interpret_python_math_code(python_code: str) -> str:
+    """
+    Interprets a string of Python code to perform math calculations.
+
+    Security Note: This function uses exec(). While it attempts to restrict
+    imports to numpy, pandas, and scipy, and runs with a restricted
+    global scope, executing arbitrary code always carries risks. Ensure
+    that input code is from a trusted source or properly sanitized.
+
+    The code must only import modules from the allowed list: numpy, pandas, scipy.
+    Submodules of these (e.g., numpy.linalg, scipy.stats) are permitted.
+    For example:
+    'import numpy as np' is allowed.
+    'from scipy.stats import norm' is allowed.
+    'import os' is NOT allowed.
+
+    To return a result, the code should either:
+    1. End with an expression (e.g., '1 + 1' or 'np.array([1,2,3]).sum()').
+    2. Assign the result to a variable named '_result' (e.g., '_result = my_calculation').
+    
+    Print statements will also be captured and returned along with the result.
+    """
+    # 1. Validate imports using AST
+    try:
+        tree = ast.parse(python_code)
+        for node in tree.body:
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    root_module = alias.name.split('.')[0]
+                    if root_module not in ALLOWED_MODULES:
+                        return (f"Error: Import of '{alias.name}' is not allowed. "
+                                f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
+            elif isinstance(node, ast.ImportFrom):
+                if node.module:  # Handles cases like 'from . import something' where module is None
+                    root_module = node.module.split('.')[0]
+                    if root_module not in ALLOWED_MODULES:
+                        return (f"Error: Import from '{node.module}' is not allowed. "
+                                f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
+    except SyntaxError as e:
+        return f"Syntax Error in input code: {e}"
+
+    # 2. Prepare execution environment
+    restricted_globals = {
+        "__builtins__": {
+            "print": print,
+            "abs": abs, "round": round, "min": min, "max": max, "sum": sum, "len": len,
+            "range": range, "zip": zip, "enumerate": enumerate,
+            "int": int, "float": float, "str": str, "list": list, "dict": dict, "tuple": tuple, "set": set,
+            "True": True, "False": False, "None": None,
+            "__import__": __import__, # Add this line
+        }
+        # numpy, pandas, scipy are NOT pre-loaded here.
+        # The user's code `import numpy` will use Python's import mechanism.
+        # The AST check above is the primary guard.
+    }
+    local_vars = {}
+
+    # 3. Capture stdout
+    old_stdout = sys.stdout
+    redirected_output = io.StringIO()
+    sys.stdout = redirected_output
+
+    # 4. Execute code and retrieve result
+    calculated_value = None
+    result_source = ""
+    output_str = ""
+
+    try:
+        compiled_code = compile(python_code, '<string>', 'exec')
+        exec(compiled_code, restricted_globals, local_vars)
+        
+        # Priority 1: Check for '_result' variable
+        if "_result" in local_vars:
+            calculated_value = local_vars["_result"]
+            result_source = "variable '_result'"
+        # Priority 2: If no _result, and the last AST node was an expression, evaluate it.
+        elif tree.body and isinstance(tree.body[-1], ast.Expr):
+            # Ensure the expression node's value is a valid AST object for ast.Expression
+            if isinstance(tree.body[-1].value, ast.AST):
+                last_expr_ast = ast.Expression(body=tree.body[-1].value)
+                # Compile the expression in 'eval' mode
+                compiled_expr = compile(last_expr_ast, '<string>', 'eval')
+                # Evaluate in the context of restricted_globals and local_vars (which holds state from exec)
+                calculated_value = eval(compiled_expr, restricted_globals, local_vars)
+                result_source = "last expression"
+            
+        sys.stdout = old_stdout # Restore stdout before getting its value
+        output_str = redirected_output.getvalue()
+
+        if calculated_value is not None:
+            return f"Result (from {result_source}):\n{calculated_value}\n\nCaptured Output:\n{output_str}".strip()
+        else:
+            return f"Executed successfully.\n\nCaptured Output:\n{output_str}\n(No specific result value found via '_result' variable or last expression evaluation.)".strip()
+
+    except Exception as e:
+        if sys.stdout == redirected_output: # Ensure stdout is restored on error too
+            sys.stdout = old_stdout
+        output_str = redirected_output.getvalue() # Get any output captured before the error
+        return f"Execution Error: {type(e).__name__}: {e}\n\nCaptured Output:\n{output_str}".strip()
+    finally:
+        # Ensure stdout is always restored
+        if sys.stdout == redirected_output:
+            sys.stdout = old_stdout
+
+
+# Example usage:
+if __name__ == "__main__":
+    code = """
+import numpy as np
+# import os # This should trigger an error since 'os' is not allowed
+arr = np.array([1, 2, 3, 4, 5])
+_result = arr.mean()
+"""
+    print(interpret_python_math_code(code))