Update agent.py

agent.py

@@ -2,6 +2,7 @@ import os
 import time
 from pathlib import Path
 from typing import Optional, Union
+from itertools import cycle
 
 import pandas as pd
 from dotenv import load_dotenv
@@ -20,11 +21,59 @@ from smolagents.tools import Tool
 # Load environment variables
 load_dotenv()
 
-# Initialize the model
-model = LiteLLMModel(
-    model_id="gemini/gemini-2.5-flash-preview-09-2025",
-    api_key=os.getenv("GEMINI_API_KEY"),
-)
+
+class MultiModelManager:
+    """Manages multiple Groq models with rotation and fallback."""
+    
+    def __init__(self):
+        # Your selected Groq models with function calling / tool use support
+        self.models = [
+            "openai/gpt-oss-120b",                          # GPT OSS 120B - Most powerful, 500 tok/s
+            "openai/gpt-oss-20b",                           # GPT OSS 20B - Fast, 1000 tok/s
+            "meta-llama/llama-4-scout-17b-16e-instruct",    # Llama 4 Scout - Multimodal
+            "qwen/qwen3-32b",                               # Qwen 3 32B - Advanced reasoning
+            "moonshotai/kimi-k2-instruct",                  # Kimi K2 - 1T params, agentic
+        ]
+        
+        self.api_key = os.getenv("GROQ_API_KEY")
+        self.model_cycle = cycle(self.models)
+        self.current_model_name = self.models[0]
+    
+    def get_next_model(self):
+        """Get the next model in rotation."""
+        self.current_model_name = next(self.model_cycle)
+        return LiteLLMModel(
+            model_id=self.current_model_name,
+            api_key=self.api_key,
+        )
+    
+    def get_model_by_complexity(self, complexity: str = "high"):
+        """
+        Get a model based on task complexity.
+        
+        Args:
+            complexity: "high", "medium", or "low"
+        """
+        if complexity == "high":
+            model_id = self.models[0]  # llama-3.3-70b
+        elif complexity == "medium":
+            model_id = self.models[2]  # mixtral-8x7b
+        else:  # low
+            model_id = self.models[3]  # llama-3.1-8b
+        
+        self.current_model_name = model_id
+        return LiteLLMModel(
+            model_id=model_id,
+            api_key=self.api_key,
+        )
+    
+    def get_primary_model(self):
+        """Get the primary (best) model."""
+        self.current_model_name = self.models[0]
+        return LiteLLMModel(
+            model_id=self.models[0],
+            api_key=self.api_key,
+        )
 
 
 class ExcelToTextTool(Tool):
@@ -74,12 +123,25 @@ class ExcelToTextTool(Tool):
 
 
 class GaiaAgent:
-    """An agent capable of using tools to answer general questions with rate limiting."""
+    """An agent with multiple Groq models for better performance."""
     
-    def __init__(self):
-        print("✅ GaiaAgent initialized with tools.")
+    def __init__(self, strategy: str = "primary"):
+        """
+        Initialize agent with model strategy.
         
-        tools = [
+        Args:
+            strategy: "primary" (use best model), "rotate" (cycle through models),
+                     or "adaptive" (choose based on task complexity)
+        """
+        print(f"✅ GaiaAgent initialized with '{strategy}' strategy.")
+        
+        self.strategy = strategy
+        self.model_manager = MultiModelManager()
+        self.retry_count = 0
+        self.max_retries = 2
+        
+        # Initialize tools
+        self.tools = [
             DuckDuckGoSearchTool(),
             WikipediaSearchTool(),
             ExcelToTextTool(),
@@ -87,16 +149,61 @@ class GaiaAgent:
             FinalAnswerTool(),
         ]
         
+        # System prompt for better performance
+        self.system_prompt = (
+            "You are a helpful assistant that answers questions accurately. "
+            "When given a question:\n"
+            "1. Analyze it carefully and break it down step by step\n"
+            "2. Use tools when needed (web search, Python code, etc.)\n"
+            "3. For math/logic problems, write Python code to verify your answer\n"
+            "4. For factual questions, search the web if needed\n"
+            "5. Double-check your work before providing the final answer\n"
+            "6. Provide concise, direct answers without unnecessary explanation\n"
+            "7. If dealing with tables or data, use Python/pandas to analyze them accurately"
+        )
+        
+        # Rate limiting
+        self.last_call_time = 0
+        self.min_delay = 1  # Groq is very fast, minimal delay needed
+        
+        # Initialize agent with primary model
+        self._reinitialize_agent()
+    
+    def _reinitialize_agent(self):
+        """Reinitialize the agent with a new model."""
+        if self.strategy == "primary":
+            model = self.model_manager.get_primary_model()
+        elif self.strategy == "rotate":
+            model = self.model_manager.get_next_model()
+        else:  # adaptive
+            model = self.model_manager.get_model_by_complexity("high")
+        
+        print(f"🤖 Using model: {self.model_manager.current_model_name}")
+        
         self.agent = CodeAgent(
             model=model,
-            tools=tools,
+            tools=self.tools,
             add_base_tools=True,
             additional_authorized_imports=["pandas", "numpy", "csv", "subprocess"],
+            system_prompt=self.system_prompt,
         )
+    
+    def _detect_complexity(self, question: str) -> str:
+        """Detect question complexity based on keywords."""
+        question_lower = question.lower()
         
-        # Rate limiting configuration
-        self.last_call_time = 0
-        self.min_delay = 60  # 60 seconds between calls (10 per minute = 6s spacing)
+        # High complexity indicators
+        high_keywords = ["analyze", "complex", "multiple", "calculate", "prove", 
+                        "demonstrate", "derive", "algorithm"]
+        if any(keyword in question_lower for keyword in high_keywords):
+            return "high"
+        
+        # Low complexity indicators
+        low_keywords = ["what is", "who is", "when", "define", "list"]
+        if any(keyword in question_lower for keyword in low_keywords):
+            return "low"
+        
+        return "medium"
     
     def __call__(self, task_id: str, question: str) -> str:
         # Apply rate limiting
@@ -106,17 +213,51 @@ class GaiaAgent:
             print(f"⏳ Rate limiting: waiting {wait_time:.1f}s...")
             time.sleep(wait_time)
         
-        print(f"🔹 Agent received task_id='{task_id}' | question='{question[:50]}...'")
+        print(f"🔹 Task ID: {task_id}")
+        print(f"🔹 Question: {question[:100]}...")
         
-        try:
-            answer = self.agent.run(question)
-            if not answer:
-                answer = "⚠️ Sorry, I could not generate a valid response."
-        except Exception as e:
-            answer = f"⚠️ Agent failed with error: {e}"
+        # Adaptive strategy: choose model based on complexity
+        if self.strategy == "adaptive":
+            complexity = self._detect_complexity(question)
+            model = self.model_manager.get_model_by_complexity(complexity)
+            print(f"🎯 Detected complexity: {complexity}")
+            self._reinitialize_agent()
+        elif self.strategy == "rotate":
+            self._reinitialize_agent()
+        
+        # Try to get answer with retry logic
+        answer = None
+        for attempt in range(self.max_retries + 1):
+            try:
+                answer = self.agent.run(question)
+                if answer:
+                    break
+            except Exception as e:
+                print(f"⚠️ Attempt {attempt + 1} failed: {str(e)[:100]}")
+                if attempt < self.max_retries:
+                    print(f"🔄 Retrying with next model...")
+                    self._reinitialize_agent()
+                    time.sleep(2)
+                else:
+                    answer = f"⚠️ Agent failed after {self.max_retries + 1} attempts: {e}"
+        
+        if not answer:
+            answer = "⚠️ Sorry, I could not generate a valid response."
         
         # Update last call time
         self.last_call_time = time.time()
         
-        print(f"🔹 Agent returning answer: {answer}")
-        return answer
+        print(f"✅ Answer: {str(answer)[:100]}...")
+        return answer
+
+
+# Example usage configurations:
+
+# Strategy 1: Use primary (best) model for all tasks
+# agent = GaiaAgent(strategy="primary")
+
+# Strategy 2: Rotate between models to distribute load
+# agent = GaiaAgent(strategy="rotate")
+
+# Strategy 3: Adaptive - choose model based on question complexity
+# agent = GaiaAgent(strategy="adaptive")