agent.py

"""
agent.py – Claude-smolagents based solution for GAIA challenge
-----------------------------------------------------------
Environment
-----------
ANTHROPIC_API_KEY   – API key from Anthropic (set in Hugging Face space secrets)
GAIA_API_URL     – (optional) override for the GAIA scoring endpoint
"""

from __future__ import annotations

import base64
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, 
    DuckDuckGoSearchTool, 
    PythonInterpreterTool,
    tool,
)

# --------------------------------------------------------------------------- #
# constants & helpers
# --------------------------------------------------------------------------- #
DEFAULT_API_URL = os.getenv(
    "GAIA_API_URL", "https://agents-course-unit4-scoring.hf.space"
)
FILE_TAG = re.compile(r"<file:([^>]+)>")  # <file:xyz>

def _download_file(file_id: str) -> bytes:
    """Download the attachment for a GAIA task."""
    url = f"{DEFAULT_API_URL}/files/{file_id}"
    resp = requests.get(url, timeout=30)
    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

# --------------------------------------------------------------------------- #
# Fixed LiteLLM model for Anthropic
# --------------------------------------------------------------------------- #
class FixedAnthropicModel:
    """
    A wrapper around LiteLLM that properly handles Anthropic API calls
    and avoids the "system_instruction: Extra inputs are not permitted" error
    """
    
    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,
        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
        """
        # Get API key from env if not provided
        if api_key is None:
            api_key = os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError("No Anthropic API key provided. Set ANTHROPIC_API_KEY env var.")
                
        self.model_id = model_id
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        
        # Store the system prompt
        self.system_prompt = system_prompt or """You are a concise, highly accurate assistant specialized in solving challenges. 
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}")
    
    def __call__(self, prompt: str, **kwargs) -> str:
        """
        Call the model with appropriate handling of system prompts for Anthropic
        
        Args:
            prompt: The prompt to send to the model
            **kwargs: Additional arguments to pass to LiteLLM
            
        Returns:
            The model's response as a string
        """
        # Wait according to rate limiter
        RATE_LIMITER.wait()
        
        try:
            # For Anthropic models, we need to modify how system prompts are handled
            # We do this by using the 'messages' parameter directly with the system content
            
            # Extract system_instruction from kwargs if it exists and remove it
            # (to avoid the "Extra inputs are not permitted" error)
            if 'system_instruction' in kwargs:
                # We'll ignore it and use our stored system prompt instead
                del kwargs['system_instruction']
                
            # Create our messages array with the system message and user prompt
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": prompt}
            ]
            
            # Call LiteLLM with the proper message format for Anthropic
            from litellm import completion
            response = completion(
                model=self.model_id,
                messages=messages,
                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:
            if "rate_limit" in str(e).lower():
                # Specific handling for rate limit errors
                print(f"Rate limit error: {e}")
                print("Waiting 60 seconds before retrying...")
                time.sleep(60)
                # Recursive retry after waiting
                return self.__call__(prompt, **kwargs)
            else:
                # Re-raise other errors
                print(f"Error calling Anthropic API: {e}")
                raise

# --------------------------------------------------------------------------- #
# custom tool: fetch GAIA attachments
# --------------------------------------------------------------------------- #
@tool
def gaia_file_reader(file_id: str) -> str:
    """
    Download a GAIA attachment and return its contents.
    Args:
        file_id: identifier that appears inside a <file:...> placeholder.
    Returns:
        base64-encoded string for binary files (images, PDFs, …) or decoded
        UTF-8 text for textual files.
    """
    try:
        raw = _download_file(file_id)
        mime = mimetypes.guess_type(file_id)[0] or "application/octet-stream"
        if mime.startswith("text") or mime in ("application/json",):
            return raw.decode(errors="ignore")
        return base64.b64encode(raw).decode()
    except Exception as exc:
        return f"ERROR downloading {file_id}: {exc}"

# --------------------------------------------------------------------------- #
# additional tool functions
# --------------------------------------------------------------------------- #
@tool
def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
    """
    Save content to a temporary file and return the path.
    Useful for processing files from the GAIA API.
    
    Args:
        content: The content to save to the file
        filename: Optional filename, will generate a random name if not provided
        
    Returns:
        Path to the saved file
    """
    temp_dir = tempfile.gettempdir()
    if filename is None:
        temp_file = tempfile.NamedTemporaryFile(delete=False)
        filepath = temp_file.name
    else:
        filepath = os.path.join(temp_dir, filename)
    
    # Write content to the file
    with open(filepath, 'w') as f:
        f.write(content)
    
    return f"File saved to {filepath}. You can read this file to process its contents."

@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """
    Download a file from a URL and save it to a temporary location.
    
    Args:
        url: The URL to download from
        filename: Optional filename, will generate one based on URL if not provided
        
    Returns:
        Path to the downloaded file
    """
    try:
        # Parse URL to get filename if not provided
        if not filename:
            path = urlparse(url).path
            filename = os.path.basename(path)
            if not filename:
                # Generate a random name if we couldn't extract one
                import uuid
                filename = f"downloaded_{uuid.uuid4().hex[:8]}"
        
        # Create temporary file
        temp_dir = tempfile.gettempdir()
        filepath = os.path.join(temp_dir, filename)
        
        # Download the file
        response = requests.get(url, stream=True)
        response.raise_for_status()
        
        # Save the file
        with open(filepath, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        return f"File downloaded to {filepath}. You can now process this file."
    except Exception as e:
        return f"Error downloading file: {str(e)}"

@tool
def extract_text_from_image(image_path: str) -> str:
    """
    Extract text from an image using pytesseract (if available).
    
    Args:
        image_path: Path to the image file
        
    Returns:
        Extracted text or error message
    """
    try:
        # Try to import pytesseract
        import pytesseract
        from PIL import Image
        
        # Open the image
        image = Image.open(image_path)
        
        # Extract text
        text = pytesseract.image_to_string(image)
        
        return f"Extracted text from image:\n\n{text}"
    except ImportError:
        return "Error: pytesseract is not installed. Please install it with 'pip install pytesseract' and ensure Tesseract OCR is installed on your system."
    except Exception as e:
        return f"Error extracting text from image: {str(e)}"

@tool
def analyze_csv_file(file_path: str, query: str) -> str:
    """
    Analyze a CSV file using pandas and answer a question about it.
    
    Args:
        file_path: Path to the CSV file
        query: Question about the data
        
    Returns:
        Analysis result or error message
    """
    try:
        import pandas as pd
        
        # Read the CSV file
        df = pd.read_csv(file_path)
        
        # Run various analyses based on the query
        result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"
        
        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())
        
        return result
    except ImportError:
        return "Error: pandas is not installed. Please install it with 'pip install pandas'."
    except Exception as e:
        return f"Error analyzing CSV file: {str(e)}"

@tool
def analyze_excel_file(file_path: str, query: str) -> str:
    """
    Analyze an Excel file using pandas and answer a question about it.
    
    Args:
        file_path: Path to the Excel file
        query: Question about the data
        
    Returns:
        Analysis result or error message
    """
    try:
        import pandas as pd
        
        # Read the Excel file
        df = pd.read_excel(file_path)
        
        # Run various analyses based on the query
        result = f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
        result += f"Columns: {', '.join(df.columns)}\n\n"
        
        # Add summary statistics
        result += "Summary statistics:\n"
        result += str(df.describe())
        
        return result
    except ImportError:
        return "Error: pandas and openpyxl are not installed. Please install them with 'pip install pandas openpyxl'."
    except Exception as e:
        return f"Error analyzing Excel file: {str(e)}"

# --------------------------------------------------------------------------- #
# GAIAAgent class
# --------------------------------------------------------------------------- #
class GAIAAgent:
    def __init__(
        self, 
        api_key: Optional[str] = None,
        temperature: float = 0.1,
        verbose: bool = False,
        max_tokens: int = 1024,
    ):
        """
        Initialize a GAIAAgent with Claude model
        
        Args:
            api_key: Anthropic API key (fetched from environment if not provided)
            temperature: Temperature for text generation
            verbose: Enable verbose logging
            max_tokens: Maximum number of tokens to generate per response
        """
        # Set verbosity
        self.verbose = verbose
        
        # System prompt for all Claude interactions
        self.system_prompt = """You are a concise, highly accurate assistant specialized in solving challenges for the GAIA benchmark. 
Unless explicitly required, reply with ONE short sentence.
Your answers should be precise, direct, and exactly match the expected format.
All answers are graded by exact string match, so format carefully!"""
        
        # 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.")
            
        if self.verbose:
            print(f"Using Anthropic token: {api_key[:5]}...")
                
        # Initialize Claude model with our fixed wrapper
        self.model = FixedAnthropicModel(
            model_id="anthropic/claude-3-5-sonnet-20240620",  # Use Claude 3.5 Sonnet
            api_key=api_key,
            temperature=temperature,
            max_tokens=max_tokens,
            system_prompt=self.system_prompt,
        )
            
        if self.verbose:
            print(f"Initialized model: FixedAnthropicModel - claude-3-5-sonnet-20240620")
        
        # Initialize default tools
        self.tools = [
            DuckDuckGoSearchTool(),
            PythonInterpreterTool(),
            save_and_read_file,
            download_file_from_url,
            analyze_csv_file,
            analyze_excel_file,
            gaia_file_reader
        ]
        
        # Add extract_text_from_image if PIL and pytesseract are available
        try:
            import pytesseract
            from PIL import Image
            self.tools.append(extract_text_from_image)
            if self.verbose:
                print("Added image processing tool")
        except ImportError:
            if self.verbose:
                print("Image processing libraries not available")
            
        if self.verbose:
            print(f"Initialized with {len(self.tools)} tools")
        
        # Setup imports allowed
        self.imports = ["pandas", "numpy", "datetime", "json", "re", "math", "os", "requests", "csv", "urllib"]
            
        # Initialize the CodeAgent
        self.agent = CodeAgent(
            tools=self.tools,
            model=self.model,
            additional_authorized_imports=self.imports,
            executor_type="local",
            verbosity_level=2 if self.verbose else 0
        )
        
        if self.verbose:
            print("Agent initialized and ready")
    
    def answer_question(self, question: str, task_file_path: Optional[str] = None) -> str:
        """
        Process a GAIA benchmark question and return the answer
        
        Args:
            question: The question to answer
            task_file_path: Optional path to a file associated with the question
            
        Returns:
            The answer to the question
        """
        try:
            if self.verbose:
                print(f"Processing question: {question}")
                if task_file_path:
                    print(f"With associated file: {task_file_path}")
            
            # Create a context with file information if available
            context = question
            file_content = None
            
            # 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 = 8000  # Characters - reduced further to help with token limits
                    with open(task_file_path, 'r', errors='ignore') as f:
                        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
                    file_ext = os.path.splitext(task_file_path)[1].lower()
                    
                    context = f"""
Question: {question}
This question has an associated file. Here is the file content (it may be truncated):
```{file_ext}
{file_content}
```
Analyze the available file content to answer the question.
"""
                except Exception as file_e:
                    try:
                        # Try to read in binary mode
                        with open(task_file_path, 'rb') as f:
                            binary_content = f.read()
                        
                        # For image files
                        if file_ext.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']:
                            context = f"""
Question: {question}
This question has an associated image file. Please use the extract_text_from_image tool to process it.
File path: {task_file_path}
"""
                        else:
                            context = f"""
Question: {question}
This question has an associated file at path: {task_file_path}
This is a binary file. Use appropriate tools to analyze it.
"""
                    except Exception as binary_e:
                        context = f"""
Question: {question}
This question has an associated file at path: {task_file_path}
However, there was an error reading the file: {file_e}
You can still try to answer the question based on the information provided.
"""
            
            # Check for special cases that need specific formatting
            # Reversed text questions
            if question.startswith(".") or ".rewsna eht sa" in question:
                context = f"""
This question appears to be in reversed text. Here's the reversed version:
{question[::-1]}
Now answer the question above. Remember to format your answer exactly as requested.
"""
            
            # 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.
Example: If asked "What is the capital of France?", respond just with "Paris".
"""
            
            # Run the agent with the question
            answer = self.agent.run(full_prompt)
            
            # Clean up the answer to ensure it's in the expected format
            # Remove common prefixes that models often add
            answer = self._clean_answer(answer)
            
            if self.verbose:
                print(f"Generated answer: {answer}")
                
            return answer
        except Exception as e:
            error_msg = f"Error answering question: {e}"
            if self.verbose:
                print(error_msg)
            return error_msg
    
    def _clean_answer(self, answer: any) -> str:
        """
        Clean up the answer to remove common prefixes and formatting
        that models often add but that can cause exact match failures.
        
        Args:
            answer: The raw answer from the model
            
        Returns:
            The cleaned answer as a string
        """
        # Convert non-string types to strings
        if not isinstance(answer, str):
            # Handle numeric types (float, int)
            if isinstance(answer, float):
                # Format floating point numbers properly
                # Check if it's an integer value in float form (e.g., 12.0)
                if answer.is_integer():
                    formatted_answer = str(int(answer))
                else:
                    # For currency values that might need formatting
                    if abs(answer) >= 1000:
                        formatted_answer = f"${answer:,.2f}"
                    else:
                        formatted_answer = str(answer)
                return formatted_answer
            elif isinstance(answer, int):
                return str(answer)
            else:
                # For any other type
                return str(answer)
        
        # Now we know answer is a string, so we can safely use string methods
        # Normalize whitespace
        answer = answer.strip()
        
        # Remove common prefixes and formatting that models add
        prefixes_to_remove = [
            "The answer is ", 
            "Answer: ",
            "Final answer: ",
            "The result is ",
            "To answer this question: ",
            "Based on the information provided, ",
            "According to the information: ",
        ]
        
        for prefix in prefixes_to_remove:
            if answer.startswith(prefix):
                answer = answer[len(prefix):].strip()
        
        # Remove quotes if they wrap the entire answer
        if (answer.startswith('"') and answer.endswith('"')) or (answer.startswith("'") and answer.endswith("'")):
            answer = answer[1:-1].strip()
        
        return answer

# --------------------------------------------------------------------------- #
# ClaudeAgent class - Wrapper around GAIAAgent
# --------------------------------------------------------------------------- #
class ClaudeAgent:
    """Claude-enhanced agent for GAIA challenge"""
    
    def __init__(self):
        # Try to initialize GAIAAgent with Claude
        try:
            # Get API key
            api_key = os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError("ANTHROPIC_API_KEY environment variable not found")
                
            print("✅ Initializing GAIAAgent with Claude")
            
            # Create GAIAAgent instance
            self.agent = GAIAAgent(
                api_key=api_key,
                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}")
            raise
    
    def __call__(self, question: str) -> str:
        """
        Process a GAIA question and return the answer
        
        Args:
            question: The question to answer
            
        Returns:
            The answer to the question
        """
        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")
                # GAIAAgent handles reversed text internally
            
            # Detect if there's a file
            file_match = re.search(r"<file:([^>]+)>", question)
            if file_match:
                file_id = file_match.group(1)
                print(f"Detected file reference: {file_id}")
                
                # Download the file
                try:
                    file_content = _download_file(file_id)
                    
                    # Create temporary file for the file
                    temp_dir = tempfile.gettempdir()
                    file_path = os.path.join(temp_dir, file_id)
                    
                    # Save file content
                    with open(file_path, 'wb') as f:
                        f.write(file_content)
                    
                    print(f"File downloaded to: {file_path}")
                    
                    # Remove file tag from question
                    clean_question = re.sub(r"<file:[^>]+>", "", question).strip()
                    
                    # Process question with file path
                    answer = self.agent.answer_question(clean_question, file_path)
                    return self._clean_answer(answer)
                except Exception as e:
                    print(f"Error processing file: {e}")
                    # Fall back to processing without file
            
            # Process standard question
            answer = self.agent.answer_question(question)
            return self._clean_answer(answer)
        except Exception as e:
            print(f"Error processing question: {e}")
            error_msg = f"Unable to process question: {str(e)}"
            return error_msg
    
    def _clean_answer(self, answer: str) -> str:
        """
        Final cleanup of answer to ensure correct format
        Reuses GAIAAgent's cleaning method
        """
        # Already cleaned in GAIAAgent, but do additional checks
        if isinstance(answer, str):
            # Remove any trailing periods and whitespace
            answer = answer.rstrip(". \t\n\r")
            
            # Ensure it's not too long an answer - GAIA usually needs concise responses
            if len(answer) > 1000:
                # Try to find the first sentence or statement of the answer
                sentences = answer.split('. ')
                if len(sentences) > 1:
                    return sentences[0].strip()
        
        return answer