agent.py

import os
import random
import re
import tempfile
import threading
import time
from functools import lru_cache
from pathlib import Path

import requests

from smolagents import (
    CodeAgent,
    InferenceClientModel,
    LiteLLMModel,
    Tool,
    VisitWebpageTool,
    WebSearchTool,
    WikipediaSearchTool,
)


DEFAULT_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
ANTHROPIC_MODELS_URL = "https://api.anthropic.com/v1/models"
ANTHROPIC_FAMILIES = ("sonnet", "opus", "haiku")
# Used when the models endpoint cannot be reached (no key, offline, API change).
FALLBACK_ANTHROPIC_MODELS = (
    "anthropic/claude-sonnet-4-5",
    "anthropic/claude-opus-4-1",
    "anthropic/claude-haiku-4-5",
)
# smolagents defaults to 120s, which slow reasoning models exceed on long steps.
DEFAULT_MODEL_TIMEOUT = 600
SCORING_API_URL = "https://agents-course-unit4-scoring.hf.space"

AUTHORIZED_IMPORTS = [
    "bs4",
    "collections",
    "csv",
    "datetime",
    "itertools",
    "json",
    "math",
    "pandas",
    "re",
    "statistics",
    "unicodedata",
]

ANSWER_INSTRUCTIONS = (
    "Answer the following benchmark question. Research it with the available tools when you "
    "are not certain, and compute counts, arithmetic, sorting and date maths in code instead "
    "of doing them in your head. When a question is about a Wikipedia list or table, fetch the "
    "page and parse the table with pandas rather than reading the prose. Your final answer must "
    "follow the requested format exactly and contain nothing else: no explanation, reasoning, "
    "labels, units that were not asked for, quotation marks, or bullet points."
    "\n\nPick the single most appropriate tool for each step and call only that one:\n"
    "- Reasoning, parsing, arithmetic, sorting, string work, or anything about a file you were "
    "given: write Python. Never search the web for something you can compute or already know.\n"
    "- A fact that lives on Wikipedia: wikipedia_search.\n"
    "- A page whose URL you already have: visit_webpage.\n"
    "- Only when you need to find a URL you do not have: web_search, once. If the results are "
    "usable, open the best one with visit_webpage instead of searching again; only rephrase the "
    "query if the results were genuinely irrelevant.\n"
    "Never call two tools for the same piece of information, and stop searching as soon as you "
    "have what the question asks for."
)

# Used when the agent hits max_steps: an imperfect guess can still score, an apology never does.
BEST_GUESS_INSTRUCTIONS = (
    "\n\nThe agent ran out of steps, so you must answer from the evidence it already gathered. "
    "Commit to the single most likely answer, even if the evidence is incomplete: infer it from "
    "partial results, and if you truly have nothing, fall back on your own knowledge. Never reply "
    "that the answer could not be determined, and never describe what is missing. "
    "Reply with one line only, in the exact format the question asked for, prefixed by "
    "'FINAL ANSWER: '."
)

SEARCH_ENGINES = ("duckduckgo", "bing")
SEARCH_ATTEMPTS = 3
# Free search endpoints throttle hard, and parallel workers share this process.
SEARCH_MIN_INTERVAL = 2.0
_search_lock = threading.Lock()
_last_search_at = 0.0


class ResilientWebSearchTool(Tool):
    """Web search that retries and falls back between engines instead of failing a whole run."""

    name = "web_search"
    description = (
        "Performs a web search (think a Google search) and returns the top results as markdown "
        "links with short descriptions. Retries automatically across several search engines."
    )
    inputs = {"query": {"type": "string", "description": "The search query to perform."}}
    output_type = "string"

    def __init__(self, max_results: int = 10):
        super().__init__()
        self.engines = [WebSearchTool(max_results=max_results, engine=name) for name in SEARCH_ENGINES]

    def _throttle(self) -> None:
        global _last_search_at
        with _search_lock:
            wait = SEARCH_MIN_INTERVAL - (time.monotonic() - _last_search_at)
            if wait > 0:
                time.sleep(wait)
            _last_search_at = time.monotonic()

    def forward(self, query: str) -> str:
        failures = []
        for attempt in range(SEARCH_ATTEMPTS):
            for engine, tool in zip(SEARCH_ENGINES, self.engines):
                self._throttle()
                try:
                    return tool.forward(query)
                except Exception as error:
                    failures.append(f"{engine}: {error}")
            # Jitter so parallel workers do not retry in lockstep.
            time.sleep(2**attempt + random.uniform(0, 1))
        raise Exception(
            f"Web search failed for {query!r} after {SEARCH_ATTEMPTS} rounds ({'; '.join(failures)}). "
            "Try a shorter query, or use wikipedia_search / visit_webpage instead."
        )

_ANSWER_PREFIX = re.compile(r"^\s*(final\s*answer|answer)\s*[:\-]\s*", re.IGNORECASE)
_FINAL_ANSWER_LINE = re.compile(r"final\s*answer\s*[:\-]\s*(.+)", re.IGNORECASE)
_GROUPED_NUMBER = re.compile(r"-?\d{1,3}(,\d{3})+(\.\d+)?")


def normalize_answer(answer: object) -> str:
    """Grading is exact string match, so strip the decorations models wrap answers in."""
    text = str(answer).strip()
    # The max-steps fallback answers in prose ending with a FINAL ANSWER line.
    matches = _FINAL_ANSWER_LINE.findall(text)
    if matches:
        text = matches[-1].strip()
    elif "\n" in text:
        lines = [line.strip() for line in text.splitlines() if line.strip()]
        text = lines[-1] if lines else text
    text = _ANSWER_PREFIX.sub("", text.strip()).strip(" `*").strip()
    if len(text) >= 2 and text[0] == text[-1] and text[0] in "\"'":
        text = text[1:-1].strip()
    text = text.rstrip(".").strip()
    if _GROUPED_NUMBER.fullmatch(text):
        return text.replace(",", "")
    if "," in text:
        return ", ".join(part.strip() for part in text.split(","))
    return text


def download_attachment(task_id: str, file_name: str) -> Path | None:
    try:
        response = requests.get(f"{SCORING_API_URL}/files/{task_id}", timeout=30)
        response.raise_for_status()
    except requests.RequestException as error:
        print(f"Could not download attachment for {task_id}: {error}")
        return None

    directory = Path(tempfile.gettempdir()) / "gaia_attachments" / task_id
    directory.mkdir(parents=True, exist_ok=True)
    # file_name comes from the remote API, so keep only the basename.
    target = directory / Path(file_name).name
    target.write_bytes(response.content)
    return target


@lru_cache(maxsize=1)
def anthropic_model_choices() -> tuple[str, ...]:
    """Latest Sonnet, Opus and Haiku, read from the Anthropic models endpoint."""
    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
    if not anthropic_key:
        return FALLBACK_ANTHROPIC_MODELS
    try:
        response = requests.get(
            ANTHROPIC_MODELS_URL,
            headers={"x-api-key": anthropic_key, "anthropic-version": "2023-06-01"},
            params={"limit": 100},
            timeout=15,
        )
        response.raise_for_status()
        models = response.json().get("data", [])
    except (requests.RequestException, ValueError) as error:
        print(f"Could not list Anthropic models, using defaults: {error}")
        return FALLBACK_ANTHROPIC_MODELS

    models.sort(key=lambda model: model.get("created_at", ""), reverse=True)
    latest = []
    for family in ANTHROPIC_FAMILIES:
        newest = next((m["id"] for m in models if family in m.get("id", "")), None)
        if newest:
            latest.append(f"anthropic/{newest}")
    return tuple(latest) or FALLBACK_ANTHROPIC_MODELS


def default_anthropic_model() -> str:
    return anthropic_model_choices()[0]


def build_model(timeout: int, model_id: str | None = None, requests_per_minute: int = 30):
    """Use Anthropic when an API key is set, otherwise Hugging Face inference."""
    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
    if anthropic_key:
        model_id = model_id or default_anthropic_model()
        # LiteLLM needs the provider prefix to route to Anthropic.
        if "/" not in model_id:
            model_id = f"anthropic/{model_id}"
        return LiteLLMModel(
            model_id=model_id,
            api_key=anthropic_key,
            timeout=timeout,
            max_tokens=8192,
            requests_per_minute=requests_per_minute,
        )
    model_id = model_id or os.getenv("HF_MODEL_ID") or DEFAULT_MODEL_ID
    return InferenceClientModel(
        model_id=model_id,
        provider=os.getenv("HF_PROVIDER") or None,
        timeout=timeout,
        requests_per_minute=requests_per_minute,
    )


class BasicAgent:
    def __init__(
        self,
        verbose: bool = False,
        model_id: str | None = None,
        requests_per_minute: int = 30,
    ):
        self.timeout = int(os.getenv("HF_MODEL_TIMEOUT", DEFAULT_MODEL_TIMEOUT))
        model = build_model(self.timeout, model_id, requests_per_minute)
        self.model_id = model.model_id
        self.agent = CodeAgent(
            model=model,
            tools=[
                ResilientWebSearchTool(max_results=10),
                WikipediaSearchTool(extract_format="HTML"),
                VisitWebpageTool(),
            ],
            additional_authorized_imports=AUTHORIZED_IMPORTS,
            max_steps=25,
            planning_interval=3,
            verbosity_level=2 if verbose else 0,
        )
        self.agent.prompt_templates["final_answer"]["post_messages"] += BEST_GUESS_INSTRUCTIONS
        print(
            f"BasicAgent initialized with model: {self.model_id} "
            f"(request timeout: {self.timeout}s)"
        )

    def __call__(
        self,
        question: str,
        task_id: str | None = None,
        file_name: str | None = None,
    ) -> str:
        prompt = f"{ANSWER_INSTRUCTIONS}\n\nQuestion: {question}"
        if task_id and file_name:
            path = download_attachment(task_id, file_name)
            if path is not None:
                prompt += f"\n\nThe file referenced by this question is available locally at: {path}"
            else:
                prompt += (
                    f"\n\nThis question references the file {file_name}, but it could not be "
                    "downloaded. Answer from other sources if possible."
                )
        return normalize_answer(self.agent.run(prompt))