agent.py

"""
Agent for the Agents Course GAIA final assignment.
This file must be named agent.py and placed in the root of your Space,
next to app.py.
"""

import os
import re
import requests
from smolagents import (
    CodeAgent,
    DuckDuckGoSearchTool,
    VisitWebpageTool,
    InferenceClientModel,
    tool,
)

DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"


@tool
def download_task_file(task_id: str) -> str:
    """
    Downloads the file attached to a GAIA task (if any) and saves it locally.
    Use this when a question references an attached file (image, spreadsheet,
    audio, etc). Returns the local file path, or an error message if there is
    no file for this task.

    Args:
        task_id: The task_id string for the current question.
    """
    url = f"{DEFAULT_API_URL}/files/{task_id}"
    try:
        resp = requests.get(url, timeout=30)
        if resp.status_code != 200:
            return f"No file available for task_id {task_id} (status {resp.status_code})"
        cd = resp.headers.get("content-disposition", "")
        match = re.search(r'filename="?([^"]+)"?', cd)
        filename = match.group(1) if match else f"{task_id}.bin"
        local_path = os.path.join("/tmp", filename)
        with open(local_path, "wb") as f:
            f.write(resp.content)
        return local_path
    except Exception as e:
        return f"Error downloading file: {e}"


@tool
def read_text_file(path: str) -> str:
    """
    Reads and returns the text content of a local file (e.g. .txt, .csv, .py).
    Use this after download_task_file when the file is plain text or CSV.

    Args:
        path: Local filesystem path to the file.
    """
    try:
        with open(path, "r", errors="ignore") as f:
            return f.read()[:20000]
    except Exception as e:
        return f"Error reading file: {e}"


class BasicAgent:
    def __init__(self):
        print("BasicAgent initialized.")

        model = InferenceClientModel(
            model_id="Qwen/Qwen2.5-Coder-32B-Instruct"
        )

        self.agent = CodeAgent(
            tools=[
                DuckDuckGoSearchTool(),
                VisitWebpageTool(),
                download_task_file,
                read_text_file,
            ],
            model=model,
            max_steps=12,
            additional_authorized_imports=["re", "itertools", "math", "statistics"],
        )

    def __call__(self, question: str, task_id: str = None) -> str:
        print(f"Agent received question (first 80 chars): {question[:80]}...")

        instructions = (
            "You are answering a benchmark question that will be graded by EXACT "
            "STRING MATCH. Follow these rules strictly:\n"
            "1. Reply with ONLY the final answer — no explanation, no restating "
            "the question, no leading phrases like 'The answer is'.\n"
            "2. Do not include the words 'FINAL ANSWER' in your reply.\n"
            "3. Match the exact format requested (e.g. a number with no commas "
            "unless asked, a comma-separated list with no extra spaces unless "
            "shown otherwise, exact capitalization/pluralization as requested).\n"
            "4. If the question references an attached file, call "
            f"download_task_file(task_id='{task_id}') first to fetch it, then "
            "read/inspect it before answering.\n"
            "5. If you are not fully sure, still give your best single-guess "
            "answer in the correct format — never leave it blank or explain "
            "why you're unsure.\n\n"
            f"Question:\n{question}"
        )

        try:
            answer = self.agent.run(instructions)
        except Exception as e:
            print(f"Agent error: {e}")
            answer = ""

        answer = str(answer).strip()
        print(f"Agent returning answer: {answer}")
        return answer