Update agent.py

agent.py

@@ -1,9 +1,9 @@
 """
-agent.py – Gemini-smolagents baseline using google-genai SDK
+agent.py – Claude-smolagents based solution for GAIA challenge
 -----------------------------------------------------------
 Environment
 -----------
-GOOGLE_API_KEY   – API key from Google AI Studio
+ANTHROPIC_API_KEY   – API key from Anthropic (set in Hugging Face space secrets)
 GAIA_API_URL     – (optional) override for the GAIA scoring endpoint
 """
 
@@ -13,15 +13,17 @@ import base64
 import mimetypes
 import os
 import re
-from typing import List
-
-import google.genai as genai
+import tempfile
+from typing import List, Dict, Any, Optional
+import json
 import requests
-from google.genai import types as gtypes
+from urllib.parse import urlparse
+
 from smolagents import (
-    CodeAgent,
-    DuckDuckGoSearchTool,
+    CodeAgent, 
+    DuckDuckGoSearchTool, 
     PythonInterpreterTool,
+    LiteLLMModel,
     tool,
 )
 
@@ -40,61 +42,6 @@ def _download_file(file_id: str) -> bytes:
     resp.raise_for_status()
     return resp.content
 
-# --------------------------------------------------------------------------- #
-# model wrapper
-# --------------------------------------------------------------------------- #
-class GeminiModel:
-    """
-    Minimal adapter around google-genai.Client so the instance itself is
-    callable (required by smolagents).
-    """
-
-    def __init__(
-        self,
-        model_name: str = "gemini-2.0-flash",
-        temperature: float = 0.1,
-        max_tokens: int = 128,
-    ):
-        api_key = os.getenv("GOOGLE_API_KEY")
-        if not api_key:
-            raise EnvironmentError("GOOGLE_API_KEY is not set.")
-        self.client = genai.Client(api_key=api_key)
-
-        self.model_name = model_name
-        self.temperature = temperature
-        self.max_tokens = max_tokens
-
-    # internal helper -------------------------------------------------------- #
-    def _genai_call(
-        self,
-        contents,
-        system_instruction: str,
-    ) -> str:
-        resp = self.client.models.generate_content(
-            model=self.model_name,
-            contents=contents,
-            config=gtypes.GenerateContentConfig(
-                system_instruction=system_instruction,
-                temperature=self.temperature,
-                max_output_tokens=self.max_tokens,
-            ),
-        )
-        return resp.text.strip()
-
-    # public helpers --------------------------------------------------------- #
-    def __call__(self, prompt: str, system_instruction: str, **__) -> str:
-        """Used by CodeAgent for plain-text questions."""
-        return self._genai_call(prompt, system_instruction)
-
-    def call_parts(
-        self,
-        parts: List[gtypes.Part],
-        system_instruction: str,
-    ) -> str:
-        """Multimodal path used by GeminiAgent for <file:…> questions."""
-        user_content = gtypes.Content(role="user", parts=parts)
-        return self._genai_call([user_content], system_instruction)
-
 # --------------------------------------------------------------------------- #
 # custom tool: fetch GAIA attachments
 # --------------------------------------------------------------------------- #
@@ -102,10 +49,8 @@ class GeminiModel:
 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.
@@ -116,65 +61,509 @@ def gaia_file_reader(file_id: str) -> str:
         if mime.startswith("text") or mime in ("application/json",):
             return raw.decode(errors="ignore")
         return base64.b64encode(raw).decode()
-    except Exception as exc:  # pragma: no cover
+    except Exception as exc:
         return f"ERROR downloading {file_id}: {exc}"
 
 # --------------------------------------------------------------------------- #
-# final agent
+# additional tool functions
 # --------------------------------------------------------------------------- #
-class GeminiAgent:
-    """Instantiated once in app.py; called once per question."""
+@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)}"
 
-    SYSTEM_PROMPT = (
-        "You are a concise, highly accurate assistant. "
-        "Unless explicitly required, reply with ONE short sentence. "
-        "Use the provided tools if needed. "
-        "All answers are graded by exact string match."
-    )
+@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)}"
 
-    def __init__(self):
-        self.model = GeminiModel()
+@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,
+        system_prompt: Optional[str] = None
+    ):
+        """
+        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
+            system_prompt: Custom system prompt (optional)
+        """
+        # Set verbosity
+        self.verbose = verbose
+        self.system_prompt = system_prompt or """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 or pass api_key parameter.")
+            
+        if self.verbose:
+            print(f"Using Anthropic token: {api_key[:5]}...")
+                
+        # Initialize Claude model
+        self.model = LiteLLMModel(
+            model_id="anthropic/claude-3-5-sonnet-20240620",  # Use Claude 3.5 Sonnet
+            api_key=api_key,
+            temperature=temperature
+        )
+            
+        if self.verbose:
+            print(f"Initialized model: LiteLLMModel - anthropic/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,
-            tools=[
-                PythonInterpreterTool(),
-                DuckDuckGoSearchTool(),
-                gaia_file_reader,
-            ],
-            verbosity_level=0,
+            additional_authorized_imports=self.imports,
+            executor_type="local",
+            verbosity_level=2 if self.verbose else 0
         )
-        print("✅ GeminiAgent initialised.")
+        
+        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:
+                    with open(task_file_path, 'r', errors='ignore') as f:
+                        file_content = f.read()
+                    
+                    # 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:
+```{file_ext}
+{file_content}
+```
+Analyze the file content above 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
+            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".
+"""
+            
+            # 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
 
-    # --------------------------------------------------------------------- #
-    # main entry point
-    # --------------------------------------------------------------------- #
+# --------------------------------------------------------------------------- #
+# GeminiAgent 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
+            )
+        except Exception as e:
+            print(f"Error initializing GAIAAgent: {e}")
+            raise
+    
     def __call__(self, question: str) -> str:
-        file_ids = FILE_TAG.findall(question)
-
-        # ---------- multimodal branch (images / files) -------------------- #
-        if file_ids:
-            parts: List[gtypes.Part] = []
-
-            text_part = FILE_TAG.sub("", question).strip()
-            if text_part:
-                parts.append(gtypes.Part.from_text(text_part))
-
-            for fid in file_ids:
+        """
+        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}")
+            
+            # 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:
-                    img_bytes = _download_file(fid)
-                    mime = mimetypes.guess_type(fid)[0] or "image/png"
-                    parts.append(
-                        gtypes.Part.from_bytes(data=img_bytes, mime_type=mime)
-                    )
-                except Exception as exc:
-                    parts.append(gtypes.Part.from_text(f"[FILE {fid} ERROR: {exc}]"))
-
-            answer = self.model.call_parts(parts, system_instruction=self.SYSTEM_PROMPT)
-
-        # ---------- plain-text branch ------------------------------------- #
-        else:
-            # Prepend system prompt to make sure CodeAgent->model sees it.
-            prompt = f"{self.SYSTEM_PROMPT}\n\n{question}"
-            answer = self.agent.model(prompt, system_instruction=self.SYSTEM_PROMPT)
-
-        return answer.rstrip(" .\n\r\t")
+                    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