agent.py

"""GAIA assignment agent — Wikipedia, web, files, audio, vision, YouTube."""

from __future__ import annotations

import os
import re
from pathlib import Path

from smolagents import CodeAgent, InferenceClientModel, VisitWebpageTool, WebSearchTool, WikipediaSearchTool

from files import file_context_block
from tools import build_custom_tools

ANSWER_INSTRUCTIONS = """
You are a precise GAIA benchmark agent. Solve the question using tools when needed.

Available capabilities:
- wikipedia_search: English Wikipedia articles
- web_search: DuckDuckGo web search
- visit_webpage: Read a specific URL as markdown
- read_spreadsheet: Parse Excel/CSV attachments
- execute_python_file: Run an attached .py file
- transcribe_audio: Speech-to-text for .mp3 attachments
- analyze_image: Vision Q&A for image attachments
- get_youtube_transcript: Captions from YouTube URLs in the question

Strategy:
1. Read the question carefully — note required OUTPUT FORMAT (number, list, names only, etc.).
2. If an attached file path is provided, use the matching file tool first.
3. For YouTube links in the question, use get_youtube_transcript.
4. Prefer wikipedia_search for Wikipedia-specific questions; otherwise web_search + visit_webpage.
5. For logic/reversal/math in the question text, reason directly — no tool needed.
6. Call final_answer(...) as soon as you are confident.

final_answer rules (CRITICAL — exact match grading):
- Return ONLY the answer string — no explanation, no preamble.
- Do NOT write "FINAL ANSWER" or "The answer is".
- Follow the question's format exactly (comma-separated, two decimals, last names only, etc.).
- Use the minimal text the question asks for.
"""


def maybe_unreverse(question: str) -> str:
    """Un-reverse GAIA-style reversed English prompts."""
    stripped = question.strip()
    if not stripped:
        return question
    reversed_q = stripped[::-1]
    starters = ("if you", "what ", "how ", "who ", "when ", "where ", "which ", "give ")
    if reversed_q.lower().startswith(starters):
        return reversed_q
    return question


def normalize_answer(text) -> str:
    """Strip common LLM formatting so exact-match grading has a better chance."""
    if text is None:
        return ""

    answer = str(text).strip()
    if answer.lower() in ("none", "null", "n/a", "unknown"):
        return ""

    # CodeAgent sometimes returns final_answer(...) as text instead of executing it.
    final_call = re.search(
        r"final_answer\s*\(\s*(['\"])(.*?)\1\s*\)",
        answer,
        flags=re.IGNORECASE | re.DOTALL,
    )
    if final_call:
        answer = final_call.group(2).strip()
    else:
        final_call = re.search(
            r"final_answer\s*\(\s*([^)]+)\s*\)",
            answer,
            flags=re.IGNORECASE,
        )
        if final_call:
            answer = final_call.group(1).strip().strip("'\"")

    # Drop leading "Thought:" blocks if a cleaner tail exists.
    if "Thought:" in answer and "\n" in answer:
        lines = [line.strip() for line in answer.splitlines() if line.strip()]
        if lines:
            answer = lines[-1]

    answer = re.sub(r"^```(?:\w+)?\s*|\s*```$", "", answer, flags=re.MULTILINE).strip()

    for prefix in (
        "FINAL ANSWER:",
        "Final answer:",
        "Final_answer",
        "FINAL_ANSWER",
        "Answer:",
        "The answer is:",
        "The answer is",
    ):
        if answer.lower().startswith(prefix.lower()):
            answer = answer[len(prefix) :].strip(" :")
            break

    if len(answer) >= 2 and answer[0] == answer[-1] and answer[0] in "\"'":
        answer = answer[1:-1].strip()

    # Reject obvious failure messages.
    lowered = answer.lower()
    if lowered.startswith("unfortunately") or lowered.startswith("i was unable"):
        return ""

    return answer


class BasicAgent:
    """Full GAIA agent with search, web, and multimodal file tools."""

    def __init__(self):
        model_id = os.getenv("HF_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
        token = os.getenv("HF_TOKEN")

        model = InferenceClientModel(model_id=model_id, token=token)
        wiki_tool = WikipediaSearchTool(
            user_agent="HF-Agents-Course-Student (https://huggingface.co/agents-course)",
            language="en",
            content_type="text",
        )

        toolset = [
            wiki_tool,
            WebSearchTool(max_results=8),
            VisitWebpageTool(max_output_length=30000),
            *build_custom_tools(),
        ]

        self.agent = CodeAgent(
            tools=toolset,
            model=model,
            max_steps=int(os.getenv("AGENT_MAX_STEPS", "25")),
            verbosity_level=int(os.getenv("AGENT_VERBOSITY", "0")),
            additional_authorized_imports=["re", "json", "math"],
        )
        tool_names = [getattr(t, "name", str(t)) for t in toolset]
        print(f"BasicAgent initialized (model={model_id}, tools={tool_names}).")

    def __call__(
        self,
        question: str,
        task_id: str | None = None,
        file_path: str | Path | None = None,
    ) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        if file_path:
            print(f"Attached file: {file_path}")

        question = maybe_unreverse(question)
        path = Path(file_path) if file_path else None
        task = f"{ANSWER_INSTRUCTIONS.strip()}\n\nQuestion:\n{question}"
        task += file_context_block(path)

        raw = self.agent.run(task)
        answer = normalize_answer(raw)
        if not answer:
            print("Agent did not produce a final answer (max steps or empty result).")
            answer = "UNKNOWN"
        print(f"Agent returning answer: {answer[:120]}...")
        return answer