Update agent.py

agent.py

@@ -7,7 +7,8 @@ from pathlib import Path
 from typing import Optional
 
 from dotenv import load_dotenv
-from smolagents import ToolCallingAgent, tool, LiteLLMModel
+# --- MODIFIED: Switched ToolCallingAgent to CodeAgent ---
+from smolagents import CodeAgent, tool, LiteLLMModel
 
 # Import your custom logic
 from tools import (
@@ -44,8 +45,6 @@ class RateLimitedModel(LiteLLMModel):
         )
 
 # ----------- 2. smolagents Tool Definitions -----------
-# smolagents relies heavily on type hints and docstrings to build the system prompt.
-
 @tool
 def enhanced_web_search(query: str) -> str:
     """Enhanced web search with intelligent query processing. Use for recent/broad web info.
@@ -105,7 +104,6 @@ def extract_youtube(url: str) -> str:
     return extract_youtube_info(url)
 
 # ----------- 3. Enhanced File Processing -----------
-# (Kept exactly as your original logic, as it accurately handles metadata and caching)
 def detect_file_type(file_path: str) -> Optional[str]:
     ext = Path(file_path).suffix.lower()
     file_type_mapping = {
@@ -163,19 +161,16 @@ def process_file(task_id: str, question_text: str) -> str:
 
 # ----------- 4. Agent Class -----------
 class GaiaAgent:
-    """GAIA Agent powered by smolagents"""
+    """GAIA Agent powered by smolagents CodeAgent"""
     
     def __init__(self):
-        # Using LiteLLM format. 
-        # Note: If you want to use your free NVIDIA endpoint, change model_id to:
-        # "openrouter/nvidia/nemotron-3-ultra-550b-a55b:free"
         self.model = RateLimitedModel(
             model_id=os.getenv("GEMINI_MODEL", "gemini/gemini-2.5-flash"),
-            api_key=os.getenv("GEMINI_API_KEY") # Or os.getenv("OPENROUTER_API_KEY")
+            api_key=os.getenv("GEMINI_API_KEY")
         )
         
-        # ToolCallingAgent natively loops through tools until it solves the prompt
-        self.agent = ToolCallingAgent(
+        # --- UPGRADED: Instantiated as CodeAgent with dependency authorization ---
+        self.agent = CodeAgent(
             tools=[
                 enhanced_web_search, 
                 enhanced_wikipedia, 
@@ -186,10 +181,14 @@ class GaiaAgent:
                 extract_youtube
             ],
             model=self.model,
-            max_steps=6 # Prevents infinite agent loops if tools fail
+            # Increased max_steps to 12. Code-writing agents require a few extra internal 
+            # iterations to verify, execute, and format complex multi-modal answers.
+            max_steps=12, 
+            # Crucial for GAIA spreadsheet parsing and data manipulation questions
+            additional_authorized_imports=["pandas", "numpy", "re", "math", "json", "collections"]
         )
         
-        print("✓ smolagents GAIA Agent initialized")
+        print("✓ smolagents CodeAgent Architecture initialized")
         print("✓ 4-Second Rate Limit Protection Active")
 
     def __call__(self, task_id: str, question: str) -> str:
@@ -200,7 +199,7 @@ class GaiaAgent:
         processed_question = process_file(task_id, question)
         
         try:
-            # 2. smolagents executes the ReAct loop automatically
+            # 2. Executing code loop sequence
             result = self.agent.run(processed_question)
             
             print(f"[{task_id}] FINAL ANSWER: {result}")
@@ -211,10 +210,10 @@ class GaiaAgent:
             error_msg = f"Critical error in execution: {str(e)}"
             print(f"[{task_id}] {error_msg}")
             
-            # Fallback
+            # Fallback direct query if agent framework loop experiences structural issues
             try:
                 print("Attempting fallback direct response...")
-                return self.model([{"role": "user", "content": question}]).content
+                return self.model(messages=[{"role": "user", "content": question}]).content
             except:
                 return error_msg