solution_smolagents.py

"""
solution_smolagents.py
======================
Unit 4 hands-on solution built on the **smolagents** framework (Unit 2.1).

Strategy
--------
A single `CodeAgent` equipped with research / multimodal / file tools. The agent
writes its own Python (native smolagents capability) so it can compute the
algebra / set-theory / Excel / attached-Python questions itself, while the tools
cover web research, page reading, audio transcription and image understanding.

The GAIA system prompt (paper Figure 2) is injected so that the model obeys the
strict final-answer format required for exact-match scoring.

Usage
-----
    from solution_smolagents import GAIAAgent
    agent = GAIAAgent()
    print(agent.answer_question("8e867cd7-...", "How many studio albums ..."))
"""

from __future__ import annotations

import importlib.resources

import yaml
from smolagents import CodeAgent, Tool

import gaia_common as gc


def _default_prompt_templates() -> dict:
    """Load smolagents' default CodeAgent prompt templates so we can override
    only the system prompt without losing the rest."""
    return yaml.safe_load(
        importlib.resources.files("smolagents.prompts").joinpath("code_agent.yaml").read_text()
    )


# ----------------------------------------------------------------------------
# Tools (smolagents `Tool` subclasses - robust, no source introspection needed)
# ----------------------------------------------------------------------------

class WebSearchTool(Tool):
    name = "web_search"
    description = (
        "Search the web (DuckDuckGo with fallbacks) and return the top results as "
        "numbered text with titles, URLs and snippets. Use for fact-finding and to "
        "discover pages to open with fetch_page."
    )
    inputs = {
        "query": {"type": "string", "description": "The search query."},
        "max_results": {"type": "integer", "description": "Number of results to return (default 5).", "nullable": True},
    }
    output_type = "string"

    def forward(self, query: str, max_results: int = 5) -> str:
        return gc.web_search(query, max_results=max_results)


class FetchPageTool(Tool):
    name = "fetch_page"
    description = (
        "Fetch a URL (HTML or PDF) and return its human-readable text content. "
        "Use after web_search to read the actual source of an answer."
    )
    inputs = {
        "url": {"type": "string", "description": "The full URL to fetch."},
    }
    output_type = "string"

    def forward(self, url: str) -> str:
        return gc.fetch_page(url)


class DownloadTaskFileTool(Tool):
    name = "download_task_file"
    description = (
        "Download the file attached to a GAIA question (image, audio, spreadsheet, "
        "Python file, PDF, ...). Returns the local path to the downloaded file, or "
        "'NO FILE' if the question has no attachment."
    )
    inputs = {
        "task_id": {"type": "string", "description": "The GAIA task id of the question."},
    }
    output_type = "string"

    def forward(self, task_id: str) -> str:
        path = gc.download_task_file(task_id)
        return path if path else "NO FILE"


class ReadFileTool(Tool):
    name = "read_file"
    description = (
        "Read a local file and return its content as text. Dispatches automatically on the "
        "extension: spreadsheets -> markdown table, PDF -> extracted text, images -> vision "
        "description, audio -> transcription, text/Python -> raw text."
    )
    inputs = {
        "path": {"type": "string", "description": "Local path to the file."},
        "question": {"type": "string", "description": "Optional: the original question, used to guide image analysis.", "nullable": True},
    }
    output_type = "string"

    def forward(self, path: str, question: str = "") -> str:
        return gc.read_any_file(path, question)


class TranscribeAudioTool(Tool):
    name = "transcribe_audio"
    description = (
        "Transcribe an audio file (local path) or a YouTube video (URL) to text using "
        "Whisper. Use for questions that reference .mp3 attachments or YouTube videos."
    )
    inputs = {
        "target": {"type": "string", "description": "Local audio file path OR a YouTube URL."},
    }
    output_type = "string"

    def forward(self, target: str) -> str:
        if target.startswith("http"):
            return gc.youtube_transcript(target)
        return gc.transcribe_audio(target)


class AnalyzeImageTool(Tool):
    name = "analyze_image"
    description = (
        "Analyze an image file with a vision-language model and return a detailed text "
        "description. Use for images (e.g. chess positions, figures, screenshots)."
    )
    inputs = {
        "path": {"type": "string", "description": "Local path to the image file."},
        "question": {"type": "string", "description": "The question or specific instruction for the image."},
    }
    output_type = "string"

    def forward(self, path: str, question: str) -> str:
        return gc.analyze_image(path, question)


class ExecutePythonTool(Tool):
    name = "execute_python"
    description = (
        "Run a self-contained Python program in a fresh subprocess and return its stdout. "
        "Use for exact arithmetic, data munging, set/group theory checks, and for running "
        "an attached .py file (read it first with read_file, then execute its code)."
    )
    inputs = {
        "code": {"type": "string", "description": "The complete Python code to run. It must be self-contained (imports inside)."},
    }
    output_type = "string"

    def forward(self, code: str) -> str:
        return gc.execute_python(code)


# ----------------------------------------------------------------------------
# The agent
# ----------------------------------------------------------------------------

DEFAULT_TOOLS = [
    WebSearchTool(),
    FetchPageTool(),
    DownloadTaskFileTool(),
    ReadFileTool(),
    TranscribeAudioTool(),
    AnalyzeImageTool(),
    ExecutePythonTool(),
]

EXTRA_AUTHORIZED_IMPORTS = [
    "pandas", "numpy", "math", "statistics", "json", "re", "datetime",
    "collections", "itertools", "fractions", "urllib", "requests", "csv", "html",
]


class GAIAAgent:
    """smolagents CodeAgent specialised for the GAIA level-1 leaderboard."""

    def __init__(
        self,
        model=None,
        tools: list[Tool] | None = None,
        max_steps: int = 14,
        additional_authorized_imports: list[str] | None = None,
    ):
        self.model = model or gc.make_smolagents_model()
        templates = _default_prompt_templates()
        templates["system_prompt"] = gc.GAIA_SYSTEM_PROMPT
        self.agent = CodeAgent(
            tools=tools or DEFAULT_TOOLS,
            model=self.model,
            max_steps=max_steps,
            additional_authorized_imports=additional_authorized_imports or EXTRA_AUTHORIZED_IMPORTS,
            prompt_templates=templates,
            verbosity_level=1,
        )

    # -- public API ---------------------------------------------------------
    def answer_question(self, task_id: str, question: str) -> str:
        """Full pipeline for one question: download attachment, run the agent,
        extract and normalize the final answer."""
        file_path = gc.download_task_file(task_id)
        prompt = self._build_prompt(question, file_path)
        try:
            output = self.agent.run(prompt)
        finally:
            self.agent.memory.reset()  # fresh memory per question
        answer = gc.extract_final_answer(str(output))
        return gc.normalize_final_answer(answer)

    def __call__(self, task_id: str, question: str) -> str:
        return self.answer_question(task_id, question)

    # -- helpers ------------------------------------------------------------
    @staticmethod
    def _build_prompt(question: str, file_path: str | None) -> str:
        file_hint = ""
        if file_path:
            file_hint = (
                f"\n\nAn attachment for this question has already been downloaded to: {file_path}\n"
                f"Use the read_file tool on it if you need it. The question asks: {question}"
            )
        return f"{question}{file_hint}{gc.ANSWER_ONLY_PROMPT}"


def run_demo(subset: int | None = 3):
    """Quick smoke test on the first `subset` questions."""
    agent = GAIAAgent()
    for item in gc.fetch_questions()[:subset]:
        print("\n---", item["task_id"], "---")
        print(item["question"][:150])
        print("ANSWER:", agent.answer_question(item["task_id"], item["question"]))


if __name__ == "__main__":
    run_demo()