Update agent.py

agent.py

@@ -1,5 +1,5 @@
 """
-agent.py – Claude-smolagents based solution for GAIA challenge
+agent.py – Simplified Claude implementation for GAIA challenge
 -----------------------------------------------------------
 Environment
 -----------
@@ -7,8 +7,6 @@ ANTHROPIC_API_KEY   – API key from Anthropic (set in Hugging Face space secret
 GAIA_API_URL     – (optional) override for the GAIA scoring endpoint
 """
 
-from __future__ import annotations
-
 import base64
 import mimetypes
 import os
@@ -16,20 +14,13 @@ import re
 import tempfile
 import time
 from typing import List, Dict, Any, Optional
-import json
+import random
 import requests
 from urllib.parse import urlparse
-import random
-
-from smolagents import (
-    CodeAgent, 
-    DuckDuckGoSearchTool, 
-    PythonInterpreterTool,
-    tool,
-)
+from smolagents import CodeAgent, DuckDuckGoSearchTool, PythonInterpreterTool, tool
 
 # --------------------------------------------------------------------------- #
-# constants & helpers
+# Constants & helpers
 # --------------------------------------------------------------------------- #
 DEFAULT_API_URL = os.getenv(
     "GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space"
@@ -74,15 +65,74 @@ class RateLimiter:
             time.sleep(random.uniform(0.2, 1.0))
 
 # Global rate limiter instance
-RATE_LIMITER = RateLimiter(requests_per_minute=25)  # Keep below 40 for safety
+RATE_LIMITER = RateLimiter(requests_per_minute=15)  # Reduced to be extra cautious
 
 # --------------------------------------------------------------------------- #
-# Fixed LiteLLM model for Anthropic
+# Direct function to call Claude via LiteLLM
 # --------------------------------------------------------------------------- #
-class FixedAnthropicModel:
+def call_claude(
+    prompt: str, 
+    system_prompt: Optional[str] = None, 
+    temperature: float = 0.1,
+    max_tokens: int = 1024,
+    model_name: str = "anthropic/claude-3-5-sonnet-20240620"
+) -> str:
     """
-    A wrapper around LiteLLM that properly handles Anthropic API calls
-    and avoids the "system_instruction: Extra inputs are not permitted" error
+    Call Claude through LiteLLM directly, following official LiteLLM documentation
+    
+    Args:
+        prompt: The user's question
+        system_prompt: Optional system prompt
+        temperature: Temperature for generation
+        max_tokens: Max tokens to generate
+        model_name: Claude model to use
+    
+    Returns:
+        The response text from Claude
+    """
+    from litellm import completion
+    
+    # Respect rate limits
+    RATE_LIMITER.wait()
+    
+    try:
+        # Build messages following exactly LiteLLM's documented format
+        messages = []
+        
+        # Add system message if provided
+        if system_prompt:
+            messages.append({"role": "system", "content": system_prompt})
+        
+        # Add user message - this is simple text only format
+        messages.append({"role": "user", "content": prompt})
+        
+        # Make the API call exactly as documented
+        response = completion(
+            model=model_name,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens
+        )
+        
+        # Extract just the text content from the response
+        return response.choices[0].message.content
+        
+    except Exception as e:
+        if "rate_limit" in str(e).lower():
+            print(f"Rate limit hit: {e}")
+            # Wait 60 seconds and try again
+            time.sleep(60)
+            return call_claude(prompt, system_prompt, temperature, max_tokens, model_name)
+        else:
+            print(f"Error calling Claude API: {e}")
+            raise
+
+# --------------------------------------------------------------------------- #
+# Simple Claude Model wrapper for smolagents
+# --------------------------------------------------------------------------- #
+class SimpleClaudeModel:
+    """
+    A minimal wrapper around LiteLLM's direct call to Anthropic that works with smolagents
     """
     
     def __init__(
@@ -93,16 +143,7 @@ class FixedAnthropicModel:
         max_tokens: int = 1024,
         system_prompt: Optional[str] = None,
     ):
-        """
-        Initialize a model that properly handles system prompts for Anthropic via LiteLLM
-        
-        Args:
-            model_id: Claude model ID to use
-            api_key: API key (will use ANTHROPIC_API_KEY env var if not provided)
-            temperature: Temperature for text generation
-            max_tokens: Maximum tokens to generate
-            system_prompt: System prompt to use
-        """
+        """Initialize a minimal Claude model wrapper"""
         # Get API key from env if not provided
         if api_key is None:
             api_key = os.getenv("ANTHROPIC_API_KEY")
@@ -119,49 +160,18 @@ class FixedAnthropicModel:
 Your answers should be precise, direct, and exactly match the expected format.
 All answers are graded by exact string match, so format carefully!"""
         
-        print(f"Initialized FixedAnthropicModel with {model_id}")
+        print(f"Initialized SimpleClaudeModel with {model_id}")
     
     def __call__(self, prompt: str, **kwargs) -> str:
-        """
-        Call the model with appropriate handling of system prompts for Anthropic
-        """
-        # Wait according to rate limiter
-        RATE_LIMITER.wait()
-        
-        try:
-            # Remove system_instruction if present in kwargs
-            if 'system_instruction' in kwargs:
-                del kwargs['system_instruction']
-                
-            # For Anthropic via LiteLLM, use the direct completion method
-            from litellm import completion
-            
-            # Create a simple prompt with system instructions at the beginning
-            # This avoids the nested message structure issue
-            complete_prompt = f"{self.system_prompt}\n\n{prompt}"
-            
-            response = completion(
-                model=self.model_id,
-                messages=[{"role": "user", "content": complete_prompt}],
-                api_key=self.api_key,
-                temperature=self.temperature,
-                max_tokens=self.max_tokens,
-                **kwargs
-            )
-            
-            # Extract the content from the response
-            return response.choices[0].message.content
-            
-        except Exception as e:
-            # Handle rate limit errors
-            if "rate_limit" in str(e).lower():
-                print(f"Rate limit error: {e}")
-                print("Waiting 60 seconds before retrying...")
-                time.sleep(60)
-                return self.__call__(prompt, **kwargs)
-            else:
-                print(f"Error calling Anthropic API: {e}")
-                raise
+        """Call method to make this class callable by smolagents CodeAgent"""
+        # Directly use the call_claude function
+        return call_claude(
+            prompt=prompt,
+            system_prompt=self.system_prompt,
+            temperature=self.temperature,
+            max_tokens=self.max_tokens,
+            model_name=self.model_id
+        )
 
 # --------------------------------------------------------------------------- #
 # custom tool: fetch GAIA attachments
@@ -383,8 +393,8 @@ 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 with our fixed wrapper
-        self.model = FixedAnthropicModel(
+        # Initialize Claude model with our simplified wrapper
+        self.model = SimpleClaudeModel(
             model_id="anthropic/claude-3-5-sonnet-20240620",  # Use Claude 3.5 Sonnet
             api_key=api_key,
             temperature=temperature,
@@ -393,7 +403,7 @@ All answers are graded by exact string match, so format carefully!"""
         )
             
         if self.verbose:
-            print(f"Initialized model: FixedAnthropicModel - claude-3-5-sonnet-20240620")
+            print(f"Initialized model: SimpleClaudeModel - claude-3-5-sonnet-20240620")
         
         # Initialize default tools
         self.tools = [