src/agent.py

"""Core agent logic.

Given one GAIA question object, produce a single submittable answer string.
The agent:

1. inspects the question and any attached file
2. downloads the attached file (via ``task_id``) when present
3. lets the model gather evidence with tool use (web search) and the file
4. produces a final answer
5. normalizes that answer for exact-match scoring

The returned value is **only** the answer text — no explanation, no markdown,
no "FINAL ANSWER" label. Output discipline is enforced twice: once by the
system prompt, and once by ``normalizer.finalize`` as a safety net.

Two entry points:

* :func:`run_agent` — returns just the normalized answer (used by the app).
* :func:`run_agent_detailed` — returns an :class:`AgentResult` carrying the raw
  model output alongside the normalized answer, for local debugging.

Model selection is delegated to ``llm.get_runner`` so the model can be swapped
(Anthropic / OpenAI) via configuration without touching this file. File handling
lives in ``files.py``; this module only maps a ``FileContext`` to model blocks.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from . import config, files, normalizer, prompts
from .llm import get_runner


@dataclass
class AgentResult:
    """One question's run, with both raw and normalized output for debugging."""

    task_id: str
    question: str
    raw_output: str
    submitted_answer: str
    file_kind: Optional[str] = None
    file_summary: Optional[str] = None


def run_agent(question: Dict[str, Any]) -> str:
    """Answer one GAIA question and return the normalized answer string.

    ``question`` is a question dict from ``scoring_api.get_questions`` /
    ``get_random_question``. It must contain ``task_id`` and ``question``; a
    non-empty ``file_name`` indicates an attachment to fetch.
    """
    return run_agent_detailed(question).submitted_answer


def run_agent_detailed(question: Dict[str, Any]) -> AgentResult:
    """Answer one question, returning raw + normalized output and file info.

    The raw model output is preserved verbatim (never mutated) so a debug
    workflow can compare it against the normalized submission.
    """
    task_id = str(question.get("task_id", ""))
    question_text = question.get("question", "")
    file_name = question.get("file_name") or ""

    # --- Steps 1 & 2: download/inspect the file and assemble model input ---
    ctx = files.load_file(task_id, file_name) if file_name else None
    has_file = ctx is not None and ctx.is_readable

    task_prompt = prompts.build_task_prompt(question_text, has_file=has_file)
    user_content: List[Dict[str, Any]] = [{"type": "text", "text": task_prompt}]
    if ctx is not None:
        user_content.extend(_blocks_from_file_context(ctx))

    # --- Steps 3-5: run the agent, then enforce exact-match formatting ---
    # >>> The model call (and its tool loop) lives behind the runner. <<<
    system_prompt = prompts.get_system_prompt(short=config.SHORT_PROMPT)
    runner = get_runner()
    raw_output = runner.run(system_prompt, user_content)

    # finalize() never mutates raw_output; it returns a cleaned copy.
    submitted_answer = normalizer.finalize(raw_output)

    return AgentResult(
        task_id=task_id,
        question=question_text,
        raw_output=raw_output,
        submitted_answer=submitted_answer,
        file_kind=ctx.kind if ctx else None,
        file_summary=ctx.summary if ctx else None,
    )


def _blocks_from_file_context(ctx: files.FileContext) -> List[Dict[str, Any]]:
    """Map a provider-neutral FileContext to Anthropic content blocks.

    Images become vision blocks, PDFs become document blocks, and text-like
    files (txt/csv/json) become a text block prefixed with the file summary.
    Unreadable files contribute a short note so the model knows what's missing.
    """
    if ctx.kind == "image" and ctx.data_b64:
        return [{
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": ctx.media_type or "image/png",
                "data": ctx.data_b64,
            },
        }]

    if ctx.kind == "pdf" and ctx.data_b64:
        return [{
            "type": "document",
            "source": {
                "type": "base64",
                "media_type": "application/pdf",
                "data": ctx.data_b64,
            },
        }]

    if ctx.text is not None:  # text / csv / json
        return [{"type": "text", "text": f"{ctx.summary}\n\n{ctx.text}"}]

    # Missing / unsupported: a note keeps the model informed without crashing.
    return [{"type": "text", "text": f"[{ctx.summary}]"}]