Update agent.py

agent.py

@@ -24,82 +24,49 @@ load_dotenv()
 
 class ExcelToTextTool(Tool):
     """Render an Excel worksheet as a Markdown table."""
-    
     name = "excel_to_text"
-    description = (
-        "Read an Excel file and return a Markdown table of the requested sheet. "
-        "Accepts either the sheet name or a zero-based index (as a string)."
-    )
+    description = "Read an Excel file and return a Markdown table of the requested sheet."
     inputs = {
-        "excel_path": {
-            "type": "string",
-            "description": "Path to the Excel file (.xlsx or .xls).",
-        },
-        "sheet_name": {
-            "type": "string",
-            "description": (
-                "Worksheet name or zero-based index (as a string). "
-                "Optional; defaults to the first sheet."
-            ),
-            "nullable": True,
-        },
+        "excel_path": {"type": "string", "description": "Path to the Excel file."},
+        "sheet_name": {"type": "string", "description": "Worksheet name or index. Optional.", "nullable": True},
     }
     output_type = "string"
 
     def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
-        """Load the Excel file and return the sheet as a Markdown table."""
         file_path = Path(excel_path).expanduser().resolve()
-        
         if not file_path.is_file():
             return f"Error: Excel file not found at {file_path}"
-        
         try:
-            sheet: Union[str, int] = (
-                int(sheet_name) if sheet_name and sheet_name.isdigit() else sheet_name or 0
-            )
+            sheet: Union[str, int] = int(sheet_name) if sheet_name and sheet_name.isdigit() else sheet_name or 0
             df = pd.read_excel(file_path, sheet_name=sheet)
-            
-            if hasattr(df, "to_markdown"):
-                return df.to_markdown(index=False)
-            
-            return tabulate(df, headers="keys", tablefmt="github", showindex=False)
-        
+            return df.to_markdown(index=False) if hasattr(df, "to_markdown") else tabulate(df, headers="keys", tablefmt="github", showindex=False)
         except Exception as e:
             return f"Error reading Excel file: {e}"
 
 
 class GaiaAgent:
     """
-    An agent optimized for Llama 4 Scout with multimodal capabilities.
+    Multi-model agent with intelligent routing and fallback.
     
-    Features:
-    - Uses Llama 4 Scout (30K TPM, 500K context, multimodal)
-    - Intelligent rate limiting with exponential backoff
-    - Automatic retry logic for rate limit errors
-    - Support for text and image inputs
+    Models:
+    - Fast (14.4K RPM): Simple Q&A, factual questions
+    - Reasoning (1K RPM, 12K TPM): Logic, math, complex reasoning
+    - Multimodal (1K RPM, 30K TPM): Images, chess, visual tasks
     """
     
+    MODELS = {
+        "fast": "groq/llama-3.1-8b-instant",  # 14.4K RPM, 6K TPM
+        "reasoning": "groq/llama-3.3-70b-versatile",  # 1K RPM, 12K TPM
+        "multimodal": "groq/meta-llama/llama-4-scout-17b-16e-instruct",  # 1K RPM, 30K TPM
+    }
+
     def __init__(self):
-        """Initialize agent with Llama 4 Scout model."""
-        print("✅ GaiaAgent initialized with Llama 4 Scout (30K TPM, Multimodal)")
-        
-        # Model configuration
-        self.model_id = "groq/meta-llama/llama-4-scout-17b-16e-instruct"
+        print("✅ GaiaAgent initialized with multi-model routing")
         self.api_key = os.getenv("GROQ_API_KEY")
-        
         if not self.api_key:
             raise ValueError("GROQ_API_KEY not found in environment variables")
-        
-        # Initialize model
-        self.model = LiteLLMModel(
-            model_id=self.model_id,
-            api_key=self.api_key,
-        )
-        
-        print(f"🤖 Using model: {self.model_id}")
-        print(f"📊 Limits: 30K TPM | 1K RPM | 500K context")
-        
-        # Initialize tools
+
+        # Tools
         self.tools = [
             DuckDuckGoSearchTool(),
             WikipediaSearchTool(),
@@ -107,146 +74,151 @@ class GaiaAgent:
             PythonInterpreterTool(),
             FinalAnswerTool(),
         ]
-        
-        # Create agent
-        self.agent = CodeAgent(
-            model=self.model,
-            tools=self.tools,
-            add_base_tools=True,
-            additional_authorized_imports=["pandas", "numpy", "csv", "subprocess", "PIL", "requests"],
-        )
-        
-        # Rate limiting configuration (optimized for 30K TPM)
+
+        # Rate limiting - conservative for CodeAgent's multiple API calls
         self.last_call_time = 0
-        self.min_delay = 3  # 3 seconds between tasks (generous with 30K TPM)
-        self.max_retries = 3
-        
-        # Statistics tracking
+        self.min_delay = 15  # 15s between tasks (agents make 5-10+ calls internally)
+        self.max_retries = 2  # Reduced to 2 to avoid long waits
+
+        # Stats
         self.total_tasks = 0
         self.successful_tasks = 0
         self.failed_tasks = 0
         self.rate_limit_hits = 0
-    
-    def _wait_for_rate_limit(self, wait_time: float):
-        """Wait for rate limit with progress indicator."""
-        print(f"⏳ Rate limit: waiting {wait_time:.1f}s...", end="", flush=True)
-        
-        # Show countdown
-        for remaining in range(int(wait_time), 0, -1):
-            print(f"\r⏳ Rate limit: waiting {remaining}s...  ", end="", flush=True)
-            time.sleep(1)
+        self.model_usage = {"fast": 0, "reasoning": 0, "multimodal": 0}
+
+    def _route_model(self, question: str, image: Optional[str] = None) -> str:
+        """Intelligently route to appropriate model based on question type."""
+        q = question.lower()
         
-        print("\r✓ Ready to retry                    ")
-    
-    def _extract_wait_time(self, error_str: str) -> float:
-        """Extract wait time from rate limit error message."""
-        # Look for patterns like "3.675499999s" or "try again in 3.6s"
-        patterns = [
-            r'(\d+\.?\d*)\s*s',
-            r'try again in (\d+\.?\d*)',
-            r'retry in (\d+\.?\d*)',
-        ]
+        # Multimodal: images, chess, visual content
+        if image or any(kw in q for kw in ["chess", "image", "picture", "photo", "visual", "diagram"]):
+            return "multimodal"
         
-        for pattern in patterns:
-            match = re.search(pattern, error_str)
-            if match:
-                return float(match.group(1)) + 2  # Add 2s buffer
+        # Reasoning: complex logic, math, proofs
+        if any(kw in q for kw in [
+            "prove", "counterexample", "logic", "reasoning", "algebra",
+            "calculate", "derive", "theorem", "equation", "formula",
+            "commutative", "associative", "distributive"
+        ]):
+            return "reasoning"
         
-        return 15  # Default fallback
-    
-    def _handle_rate_limit_error(self, error_str: str, attempt: int) -> float:
-        """Handle rate limit error and return wait time."""
-        self.rate_limit_hits += 1
-        
-        # Try to extract wait time from error
-        wait_time = self._extract_wait_time(error_str)
+        # Fast: simple factual questions
+        return "fast"
+
+    def _create_agent(self, model_key: str) -> CodeAgent:
+        """Create a new agent with the specified model."""
+        model_id = self.MODELS[model_key]
+        print(f"🤖 Initializing {model_key} model: {model_id}")
         
-        # Apply exponential backoff if extraction failed
-        if wait_time == 15:
-            wait_time = 15 * (attempt + 1)  # 15s, 30s, 45s, 60s
+        model = LiteLLMModel(model_id=model_id, api_key=self.api_key)
         
-        return min(wait_time, 60)  # Cap at 60s
-    
-    def __call__(self, task_id: str, question: str) -> str:
-        """
-        Process a task with automatic rate limiting and retry logic.
+        return CodeAgent(
+            model=model,
+            tools=self.tools,
+            add_base_tools=True,
+            additional_authorized_imports=["pandas", "numpy", "csv", "subprocess", "PIL", "requests"],
+        )
+
+    def _try_model(self, agent: CodeAgent, question: str, model_key: str) -> Optional[str]:
+        """Try to get answer from a model with retries."""
+        for attempt in range(self.max_retries + 1):
+            try:
+                print(f"🚀 Attempt {attempt + 1}/{self.max_retries + 1} with {model_key} model")
+                answer = agent.run(question)
+                if answer:
+                    self.model_usage[model_key] += 1
+                    return answer
+                    
+            except Exception as e:
+                err = str(e)
+                print(f"⚠️ Error: {err[:200]}{'...' if len(err) > 200 else ''}")
+                
+                if "rate limit" in err.lower() or "rate_limit" in err.lower():
+                    self.rate_limit_hits += 1
+                    
+                    # Extract wait time from error
+                    wait_match = re.search(r'(\d+\.?\d*)\s*s', err)
+                    wait_time = float(wait_match.group(1)) + 3 if wait_match else 20
+                    
+                    if attempt < self.max_retries:
+                        print(f"⏳ Rate limit hit. Waiting {wait_time:.1f}s...")
+                        time.sleep(wait_time)
+                        continue
+                    else:
+                        print(f"❌ Rate limit exhausted for {model_key}")
+                        return None
+                else:
+                    # Non-rate-limit error
+                    if attempt < self.max_retries:
+                        print(f"🔄 Retrying in 5s...")
+                        time.sleep(5)
+                        continue
+                    else:
+                        print(f"❌ Failed after {self.max_retries + 1} attempts")
+                        return None
         
-        Args:
-            task_id: Unique identifier for the task
-            question: The question to answer
-            
-        Returns:
-            The answer string
-        """
+        return None
+
+    def __call__(self, task_id: str, question: str, image: Optional[str] = None) -> str:
+        """Process a task with intelligent routing and fallback."""
         self.total_tasks += 1
         
-        # Apply base rate limiting
+        # Rate limiting
         elapsed = time.time() - self.last_call_time
         if elapsed < self.min_delay:
             wait_time = self.min_delay - elapsed
             print(f"⏳ Base rate limit: waiting {wait_time:.1f}s...")
             time.sleep(wait_time)
-        
+
         print(f"\n{'='*70}")
         print(f"🔹 Task #{self.total_tasks} | ID: {task_id}")
         print(f"🔹 Question: {question[:120]}{'...' if len(question) > 120 else ''}")
         print(f"{'='*70}")
+
+        # Route to primary model
+        primary_model = self._route_model(question, image)
+        print(f"🎯 Routing to: {primary_model} model")
         
-        answer = None
+        # Define fallback chain
+        fallback_chain = {
+            "reasoning": ["fast"],
+            "multimodal": ["reasoning", "fast"],
+            "fast": [],  # No fallback for fast model
+        }
         
-        # Retry loop with exponential backoff
-        for attempt in range(self.max_retries + 1):
+        # Try primary model first
+        models_to_try = [primary_model] + fallback_chain.get(primary_model, [])
+        
+        answer = None
+        for model_key in models_to_try:
             try:
-                print(f"\n🚀 Attempt {attempt + 1}/{self.max_retries + 1}...")
-                answer = self.agent.run(question)
+                agent = self._create_agent(model_key)
+                answer = self._try_model(agent, question, model_key)
                 
                 if answer:
                     self.successful_tasks += 1
-                    print(f"✅ Success!")
                     break
+                else:
+                    print(f"🔄 Trying fallback model...")
+                    time.sleep(3)  # Brief pause before fallback
                     
             except Exception as e:
-                error_str = str(e)
-                print(f"\n⚠️ Error on attempt {attempt + 1}:")
-                print(f"   {error_str[:200]}{'...' if len(error_str) > 200 else ''}")
-                
-                # Check if it's a rate limit error
-                if any(keyword in error_str.lower() for keyword in ['rate_limit', 'rate limit', 'quota']):
-                    wait_time = self._handle_rate_limit_error(error_str, attempt)
-                    
-                    if attempt < self.max_retries:
-                        print(f"🔄 Rate limit detected. Retrying after wait...")
-                        self._wait_for_rate_limit(wait_time)
-                        continue
-                    else:
-                        answer = f"⚠️ Rate limit exceeded after {self.max_retries + 1} attempts. Please try again later."
-                        self.failed_tasks += 1
-                        
-                # Non-rate-limit error
-                else:
-                    if attempt < self.max_retries:
-                        print(f"🔄 Retrying in 5s...")
-                        time.sleep(5)
-                        continue
-                    else:
-                        answer = f"⚠️ Agent failed: {error_str[:300]}"
-                        self.failed_tasks += 1
-        
-        # Fallback if no answer generated
+                print(f"❌ Failed to initialize {model_key}: {e}")
+                continue
+
         if not answer:
-            answer = "⚠️ Sorry, I could not generate a valid response."
+            answer = "⚠️ All models failed to generate a valid response."
             self.failed_tasks += 1
-        
-        # Update timing
+
         self.last_call_time = time.time()
         
-        # Print results
         print(f"\n{'='*70}")
         print(f"📝 Answer: {str(answer)[:200]}{'...' if len(str(answer)) > 200 else ''}")
-        print(f"{'='*70}")
+        print(f"{'='*70}\n")
         
         return answer
-    
+
     def get_stats(self) -> dict:
         """Get agent performance statistics."""
         return {
@@ -255,8 +227,9 @@ class GaiaAgent:
             "failed_tasks": self.failed_tasks,
             "success_rate": f"{(self.successful_tasks / self.total_tasks * 100):.1f}%" if self.total_tasks > 0 else "0%",
             "rate_limit_hits": self.rate_limit_hits,
+            "model_usage": self.model_usage,
         }
-    
+
     def print_stats(self):
         """Print agent performance statistics."""
         stats = self.get_stats()
@@ -268,19 +241,21 @@ class GaiaAgent:
         print(f"Failed:            {stats['failed_tasks']}")
         print(f"Success Rate:      {stats['success_rate']}")
         print(f"Rate Limit Hits:   {stats['rate_limit_hits']}")
+        print(f"\nModel Usage:")
+        for model, count in stats['model_usage'].items():
+            print(f"  {model.capitalize():12} {count} tasks")
         print(f"{'='*70}\n")
 
 
-# Example usage:
+# Example usage
 if __name__ == "__main__":
-    # Initialize agent
     agent = GaiaAgent()
-    
+
     # Test with a simple question
     answer = agent(
         task_id="test-001",
         question="What is 2+2? Explain your reasoning."
     )
-    
+
     # Print statistics
     agent.print_stats()