Update agent.py

agent.py

@@ -46,26 +46,39 @@ class ExcelToTextTool(Tool):
 
 class GaiaAgent:
     """
-    Multi-model agent with intelligent routing and fallback.
+    Single-model agent using Llama 4 Scout exclusively.
     
-    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
-    """
+    Why Llama 4 Scout:
+    - 30K TPM (highest available - 5x more than llama-3.1-8b)
+    - 500K context window
+    - Multimodal support (images, chess)
+    - 1K RPM
     
-    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
-    }
+    This avoids the 6K TPM bottleneck of llama-3.1-8b-instant.
+    """
 
     def __init__(self):
-        print("✅ GaiaAgent initialized with multi-model routing")
+        print("="*70)
+        print("✅ GaiaAgent initialized with Llama 4 Scout (30K TPM)")
+        print("="*70)
+        
         self.api_key = os.getenv("GROQ_API_KEY")
         if not self.api_key:
             raise ValueError("GROQ_API_KEY not found in environment variables")
-
+        
+        # Single model configuration - Llama 4 Scout for all tasks
+        self.model_id = "groq/meta-llama/llama-4-scout-17b-16e-instruct"
+        
+        print(f"🤖 Model: {self.model_id}")
+        print(f"📊 Limits: 30K TPM | 1K RPM | 500K context | Multimodal")
+        print("="*70 + "\n")
+        
+        # Initialize model
+        self.model = LiteLLMModel(
+            model_id=self.model_id,
+            api_key=self.api_key,
+        )
+        
         # Tools
         self.tools = [
             DuckDuckGoSearchTool(),
@@ -74,160 +87,147 @@ class GaiaAgent:
             PythonInterpreterTool(),
             FinalAnswerTool(),
         ]
-
-        # Rate limiting - conservative for CodeAgent's multiple API calls
+        
+        # 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 - 30K TPM is generous but agents make multiple calls
         self.last_call_time = 0
-        self.min_delay = 15  # 15s between tasks (agents make 5-10+ calls internally)
-        self.max_retries = 2  # Reduced to 2 to avoid long waits
-
+        self.min_delay = 10  # 10s between tasks (reasonable with 30K TPM)
+        self.max_retries = 3  # More retries since we have higher TPM
+        
         # Stats
         self.total_tasks = 0
         self.successful_tasks = 0
         self.failed_tasks = 0
         self.rate_limit_hits = 0
-        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()
-        
-        # Multimodal: images, chess, visual content
-        if image or any(kw in q for kw in ["chess", "image", "picture", "photo", "visual", "diagram"]):
-            return "multimodal"
-        
-        # 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"
-        
-        # 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}")
-        
-        model = LiteLLMModel(model_id=model_id, api_key=self.api_key)
+    def _extract_wait_time(self, error_str: str) -> float:
+        """Extract wait time from rate limit error message."""
+        patterns = [
+            r'try again in (\d+\.?\d*)\s*s',
+            r'retry in (\d+\.?\d*)\s*s',
+            r'(\d+\.?\d*)\s*s',
+        ]
         
-        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
+        for pattern in patterns:
+            match = re.search(pattern, error_str)
+            if match:
+                return float(match.group(1)) + 5  # Add 5s buffer
         
-        return None
+        return 30  # Default fallback
 
-    def __call__(self, task_id: str, question: str, image: Optional[str] = None) -> str:
-        """Process a task with intelligent routing and fallback."""
+    def __call__(self, task_id: str, question: str) -> str:
+        """Process a task with automatic rate limiting and retry."""
         self.total_tasks += 1
         
         # 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...")
+            print(f"⏳ 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}")
+        print(f"📋 Task #{self.total_tasks} | ID: {task_id}")
+        print(f"❓ Question: {question[:150]}{'...' if len(question) > 150 else ''}")
+        print(f"{'='*70}\n")
 
-        # Route to primary model
-        primary_model = self._route_model(question, image)
-        print(f"🎯 Routing to: {primary_model} model")
-        
-        # Define fallback chain
-        fallback_chain = {
-            "reasoning": ["fast"],
-            "multimodal": ["reasoning", "fast"],
-            "fast": [],  # No fallback for fast model
-        }
-        
-        # Try primary model first
-        models_to_try = [primary_model] + fallback_chain.get(primary_model, [])
-        
         answer = None
-        for model_key in models_to_try:
+        
+        # Retry loop with exponential backoff
+        for attempt in range(self.max_retries + 1):
             try:
-                agent = self._create_agent(model_key)
-                answer = self._try_model(agent, question, model_key)
+                print(f"🚀 Attempt {attempt + 1}/{self.max_retries + 1}")
+                answer = self.agent.run(question)
                 
-                if answer:
+                if answer and len(str(answer).strip()) > 0:
                     self.successful_tasks += 1
+                    print(f"✅ Success!")
                     break
                 else:
-                    print(f"🔄 Trying fallback model...")
-                    time.sleep(3)  # Brief pause before fallback
+                    print(f"⚠️ Empty answer received")
+                    if attempt < self.max_retries:
+                        time.sleep(5)
+                        continue
                     
             except Exception as e:
-                print(f"❌ Failed to initialize {model_key}: {e}")
-                continue
+                error_str = str(e)
+                
+                # Show condensed error
+                if len(error_str) > 300:
+                    print(f"❌ Error: {error_str[:300]}...")
+                else:
+                    print(f"❌ Error: {error_str}")
+                
+                # Check if it's a rate limit error
+                if "rate limit" in error_str.lower() or "rate_limit" in error_str.lower():
+                    self.rate_limit_hits += 1
+                    wait_time = self._extract_wait_time(error_str)
+                    
+                    if attempt < self.max_retries:
+                        print(f"⏳ Rate limit hit. Waiting {wait_time:.1f}s before retry...")
+                        
+                        # Show countdown for long waits
+                        if wait_time > 10:
+                            for remaining in range(int(wait_time), 0, -5):
+                                print(f"   ⏱️  {remaining}s remaining...", flush=True)
+                                time.sleep(5)
+                        else:
+                            time.sleep(wait_time)
+                        
+                        print(f"🔄 Retrying...")
+                        continue
+                    else:
+                        answer = "⚠️ Rate limit exceeded after all retries."
+                        self.failed_tasks += 1
+                        
+                # Authentication error
+                elif "authentication" in error_str.lower() or "api key" in error_str.lower():
+                    answer = "⚠️ Authentication failed. Check your GROQ_API_KEY."
+                    self.failed_tasks += 1
+                    break
+                    
+                # Other errors
+                else:
+                    if attempt < self.max_retries:
+                        print(f"🔄 Retrying in 5s...")
+                        time.sleep(5)
+                        continue
+                    else:
+                        answer = f"⚠️ Failed after {self.max_retries + 1} attempts."
+                        self.failed_tasks += 1
 
+        # Fallback
         if not answer:
-            answer = "⚠️ All models failed to generate a valid response."
+            answer = "⚠️ Could not generate a valid response."
             self.failed_tasks += 1
 
+        # Update timing
         self.last_call_time = time.time()
         
+        # Print result
         print(f"\n{'='*70}")
-        print(f"📝 Answer: {str(answer)[:200]}{'...' if len(str(answer)) > 200 else ''}")
+        answer_preview = str(answer)[:250] + ('...' if len(str(answer)) > 250 else '')
+        print(f"✍️  Answer: {answer_preview}")
         print(f"{'='*70}\n")
         
-        return answer
+        return str(answer)
 
     def get_stats(self) -> dict:
         """Get agent performance statistics."""
+        success_rate = (self.successful_tasks / self.total_tasks * 100) if self.total_tasks > 0 else 0
         return {
             "total_tasks": self.total_tasks,
             "successful_tasks": self.successful_tasks,
             "failed_tasks": self.failed_tasks,
-            "success_rate": f"{(self.successful_tasks / self.total_tasks * 100):.1f}%" if self.total_tasks > 0 else "0%",
+            "success_rate": f"{success_rate:.1f}%",
             "rate_limit_hits": self.rate_limit_hits,
-            "model_usage": self.model_usage,
         }
 
     def print_stats(self):
@@ -237,25 +237,21 @@ class GaiaAgent:
         print(f"📊 AGENT STATISTICS")
         print(f"{'='*70}")
         print(f"Total Tasks:       {stats['total_tasks']}")
-        print(f"Successful:        {stats['successful_tasks']}")
-        print(f"Failed:            {stats['failed_tasks']}")
+        print(f"Successful:        {stats['successful_tasks']} ✅")
+        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"Rate Limit Hits:   {stats['rate_limit_hits']} 🚫")
         print(f"{'='*70}\n")
 
 
 # Example usage
 if __name__ == "__main__":
     agent = GaiaAgent()
-
-    # Test with a simple question
+    
+    # Test
     answer = agent(
         task_id="test-001",
-        question="What is 2+2? Explain your reasoning."
+        question="What is 2+2? Show your calculation."
     )
-
-    # Print statistics
+    
     agent.print_stats()