Update agent.py

agent.py

@@ -14,10 +14,12 @@ import mimetypes
 import os
 import re
 import tempfile
+import time
 from typing import List, Dict, Any, Optional
 import json
 import requests
 from urllib.parse import urlparse
+import random
 
 from smolagents import (
     CodeAgent, 
@@ -42,6 +44,39 @@ def _download_file(file_id: str) -> bytes:
     resp.raise_for_status()
     return resp.content
 
+# --------------------------------------------------------------------------- #
+# Rate limiting helper
+# --------------------------------------------------------------------------- #
+class RateLimiter:
+    """Simple rate limiter to prevent Anthropic API rate limit errors"""
+    def __init__(self, requests_per_minute=20, burst=3):
+        self.requests_per_minute = requests_per_minute
+        self.burst = burst
+        self.request_times = []
+        
+    def wait(self):
+        """Wait if needed to avoid exceeding rate limits"""
+        now = time.time()
+        # Remove timestamps older than 1 minute
+        self.request_times = [t for t in self.request_times if now - t < 60]
+        
+        # If we've made too many requests in the last minute, wait
+        if len(self.request_times) >= self.requests_per_minute:
+            oldest = min(self.request_times)
+            sleep_time = 60 - (now - oldest) + 1  # +1 for safety
+            print(f"Rate limit approaching. Waiting {sleep_time:.2f} seconds before next request...")
+            time.sleep(sleep_time)
+        
+        # Add current timestamp to the list
+        self.request_times.append(time.time())
+        
+        # Add a small random delay to avoid bursts of requests
+        if len(self.request_times) > self.burst:
+            time.sleep(random.uniform(0.2, 1.0))
+
+# Global rate limiter instance
+RATE_LIMITER = RateLimiter(requests_per_minute=25)  # Keep below 40 for safety
+
 # --------------------------------------------------------------------------- #
 # custom tool: fetch GAIA attachments
 # --------------------------------------------------------------------------- #
@@ -224,6 +259,81 @@ def analyze_excel_file(file_path: str, query: str) -> str:
     except Exception as e:
         return f"Error analyzing Excel file: {str(e)}"
 
+# --------------------------------------------------------------------------- #
+# Custom LiteLLM model with rate limiting and error handling
+# --------------------------------------------------------------------------- #
+class RateLimitedClaudeModel:
+    def __init__(
+        self,
+        model_id: str = "anthropic/claude-3-5-sonnet-20240620",
+        api_key: Optional[str] = None,
+        temperature: float = 0.1,
+        max_tokens: int = 1024,
+        max_retries: int = 3,
+        retry_delay: int = 5,
+    ):
+        """
+        Initialize a Claude model with rate limiting and error handling
+        
+        Args:
+            model_id: The model ID to use
+            api_key: The API key to use
+            temperature: The temperature to use
+            max_tokens: The maximum number of tokens to generate
+            max_retries: The maximum number of retries on rate limit errors
+            retry_delay: The initial delay between retries (will increase exponentially)
+        """
+        # Get API key
+        if api_key is None:
+            api_key = os.getenv("ANTHROPIC_API_KEY")
+            if not api_key:
+                raise ValueError("No Anthropic token provided. Please set ANTHROPIC_API_KEY environment variable or pass api_key parameter.")
+        
+        self.model_id = model_id
+        self.api_key = api_key
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        
+        # Create the underlying LiteLLM model
+        self.model = LiteLLMModel(
+            model_id=model_id,
+            api_key=api_key,
+            temperature=temperature
+        )
+    
+    def __call__(self, prompt: str, system_instruction: str, **kwargs) -> str:
+        """
+        Call the model with rate limiting and error handling
+        
+        Args:
+            prompt: The prompt to generate from
+            system_instruction: The system instruction to use
+            
+        Returns:
+            The generated text
+        """
+        retries = 0
+        while True:
+            try:
+                # Wait according to rate limiter
+                RATE_LIMITER.wait()
+                
+                # Call the model
+                return self.model(prompt, system_instruction=system_instruction, **kwargs)
+            
+            except Exception as e:
+                # Check if it's a rate limit error
+                if "rate_limit_error" in str(e) and retries < self.max_retries:
+                    retries += 1
+                    sleep_time = self.retry_delay * (2 ** (retries - 1))  # Exponential backoff
+                    print(f"Rate limit exceeded, retrying in {sleep_time} seconds (attempt {retries}/{self.max_retries})...")
+                    time.sleep(sleep_time)
+                else:
+                    # If it's not a rate limit error or we've exceeded max retries, raise
+                    raise
+
 # --------------------------------------------------------------------------- #
 # GAIAAgent class
 # --------------------------------------------------------------------------- #
@@ -233,7 +343,8 @@ class GAIAAgent:
         api_key: Optional[str] = None,
         temperature: float = 0.1,
         verbose: bool = False,
-        system_prompt: Optional[str] = None
+        system_prompt: Optional[str] = None,
+        max_tokens: int = 1024,
     ):
         """
         Initialize a GAIAAgent with Claude model
@@ -243,6 +354,7 @@ class GAIAAgent:
             temperature: Temperature for text generation
             verbose: Enable verbose logging
             system_prompt: Custom system prompt (optional)
+            max_tokens: Maximum number of tokens to generate per response
         """
         # Set verbosity
         self.verbose = verbose
@@ -260,15 +372,16 @@ All answers are graded by exact string match, so format carefully!"""
         if self.verbose:
             print(f"Using Anthropic token: {api_key[:5]}...")
                 
-        # Initialize Claude model
-        self.model = LiteLLMModel(
+        # Initialize Claude model with rate limiting
+        self.model = RateLimitedClaudeModel(
             model_id="anthropic/claude-3-5-sonnet-20240620",  # Use Claude 3.5 Sonnet
             api_key=api_key,
-            temperature=temperature
+            temperature=temperature,
+            max_tokens=max_tokens,
         )
             
         if self.verbose:
-            print(f"Initialized model: LiteLLMModel - anthropic/claude-3-5-sonnet-20240620")
+            print(f"Initialized model: RateLimitedClaudeModel - anthropic/claude-3-5-sonnet-20240620")
         
         # Initialize default tools
         self.tools = [
@@ -334,8 +447,12 @@ All answers are graded by exact string match, so format carefully!"""
             # If there's a file, read it and include its content in the context
             if task_file_path:
                 try:
+                    # Limit file content size to avoid token limits
+                    max_file_size = 10000  # Characters
                     with open(task_file_path, 'r', errors='ignore') as f:
-                        file_content = f.read()
+                        file_content = f.read(max_file_size)
+                        if len(file_content) >= max_file_size:
+                            file_content = file_content[:max_file_size] + "... [content truncated to prevent exceeding token limits]"
                     
                     # Determine file type from extension
                     import os
@@ -343,11 +460,11 @@ All answers are graded by exact string match, so format carefully!"""
                     
                     context = f"""
 Question: {question}
-This question has an associated file. Here is the file content:
+This question has an associated file. Here is the file content (it may be truncated):
 ```{file_ext}
 {file_content}
 ```
-Analyze the file content above to answer the question.
+Analyze the available file content to answer the question.
 """
                 except Exception as file_e:
                     try:
@@ -385,12 +502,12 @@ This question appears to be in reversed text. Here's the reversed version:
 Now answer the question above. Remember to format your answer exactly as requested.
 """
             
-            # Add a prompt to ensure precise answers
+            # Add a prompt to ensure precise answers but keep it concise
             full_prompt = f"""{context}
 When answering, provide ONLY the precise answer requested. 
 Do not include explanations, steps, reasoning, or additional text.
 Be direct and specific. GAIA benchmark requires exact matching answers.
-For example, if asked "What is the capital of France?", respond simply with "Paris".
+Example: If asked "What is the capital of France?", respond just with "Paris".
 """
             
             # Run the agent with the question
@@ -486,8 +603,9 @@ class ClaudeAgent:
             # Create GAIAAgent instance
             self.agent = GAIAAgent(
                 api_key=api_key,
-                temperature=0.1,  # Use low temperature for precise answers
-                verbose=True,     # Enable verbose logging
+                temperature=0.1,      # Use low temperature for precise answers
+                verbose=True,         # Enable verbose logging
+                max_tokens=1024,      # Reduce max tokens to avoid hitting rate limits
             )
         except Exception as e:
             print(f"Error initializing GAIAAgent: {e}")
@@ -506,6 +624,9 @@ class ClaudeAgent:
         try:
             print(f"Received question: {question[:100]}..." if len(question) > 100 else f"Received question: {question}")
             
+            # Add delay between questions to respect rate limits
+            time.sleep(random.uniform(0.5, 2.0))
+            
             # Detect reversed text
             if question.startswith(".") or ".rewsna eht sa" in question:
                 print("Detected reversed text question")