agent.py

"""Groq-powered GAIA Level-1 agent for the HF Agents Course Unit 4 assignment."""

from __future__ import annotations

import os
import re

from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain_groq import ChatGroq

from tools import TOOLS

load_dotenv()

SYSTEM_PROMPT = """You are a careful GAIA evaluation agent. Scoring is exact string match.

Tool rules:
1. If the question contains a YouTube URL, call youtube_transcript on that URL FIRST. Do not
   rely on web_search for video contents.
2. If file_name is provided, call download_task_file(task_id) first, then the right follow-up
   tool (analyze_excel, transcribe_audio, analyze_image, run_python_file).
3. For reversed/backwards text, use reverse_text first.
4. Use wikipedia_search or web_search for facts; try a second, reworded search before giving up.

Answer rules:
5. Think step by step privately, but your FINAL reply must be ONLY the answer string.
6. Never apologise, never explain, and never say a value is unavailable or unspecified.
   If you are unsure, state your single best guess in the required format.
7. No "FINAL ANSWER" prefix, no quotes, no units unless the question asks for them.
8. Match the requested format exactly (plain number, comma-separated list, alphabetical
   order, IOC code, etc.).
"""

REFUSAL_HINTS = (
    "not specified",
    "not available",
    "unable to",
    "unfortunately",
    "i cannot",
    "i could not",
    "i don't have",
    "no file",
    "does not have",
    "not provided",
    "sorry",
    "search results",
)

EXTRACT_PROMPT = """Question:
{question}

Draft response:
{draft}

Reply as <answer>THE ANSWER</answer> and nothing else. The answer must be a short string
(a number, a word, a name, or a comma-separated list) with no sentence, explanation or
apology. If the draft contains no answer, put your single best guess inside the tag."""


def _clean_answer(text: str) -> str:
    """Keep only the final exact-match answer string."""
    if not text:
        return ""
    text = str(text).strip()

    # Prefer content after common markers if the model still adds them
    for marker in ("FINAL ANSWER:", "Final Answer:", "Answer:"):
        if marker in text:
            text = text.split(marker)[-1].strip()

    # Prefer last non-empty line (often the concise answer)
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    if lines:
        text = lines[-1]

    # Strip surrounding quotes / boxed latex
    text = text.strip().strip('"').strip("'")
    boxed = re.search(r"\\boxed\{([^{}]+)\}", text)
    if boxed:
        text = boxed.group(1).strip()
    return text.rstrip(".")


def _is_verbose(text: str) -> bool:
    """Sentence-like or apologetic answers never match the ground truth."""
    lowered = text.lower()
    return len(text) > 100 or any(hint in lowered for hint in REFUSAL_HINTS)


def _salvage(*candidates: str) -> str:
    """Last resort when the model keeps writing prose: prefer a number, else a short clause."""
    texts = [re.sub(r"https?://\S+", " ", c) for c in candidates]
    for text in texts:
        number = re.search(r"-?\d+(?:,\d{3})*(?:\.\d+)?", text)
        if number:
            return number.group(0)
    for text in texts:
        for clause in re.split(r"[.;\n]", text):
            clause = clause.strip()
            if clause and not _is_verbose(clause):
                return clause[:60]
    return ""


class GaiaAgent:
    """Agent that answers one GAIA question using Groq + tools."""

    def __init__(self) -> None:
        api_key = os.getenv("GROQ_API_KEY")
        if not api_key:
            raise RuntimeError("GROQ_API_KEY is missing in .env")

        self._api_key = api_key
        self._fallback_model = os.getenv("GROQ_FALLBACK_MODEL", "llama-3.1-8b-instant")
        self._on_fallback = False
        self._build(os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile"))
        print("GaiaAgent initialized (Groq + tools).")

    def _build(self, model: str) -> None:
        self.model_name = model
        self.llm = ChatGroq(model=model, api_key=self._api_key, temperature=0)
        self.agent = create_agent(
            model=self.llm,
            tools=TOOLS,
            system_prompt=SYSTEM_PROMPT,
        )

    def _switch_to_fallback(self, error: Exception) -> bool:
        """Groq quotas are per model, so a second model keeps the run alive."""
        if self._on_fallback or "rate_limit" not in str(error).lower():
            return False
        self._on_fallback = True
        print(f"Rate limited on {self.model_name}; switching to {self._fallback_model}.")
        self._build(self._fallback_model)
        return True

    def __call__(
        self,
        question: str,
        task_id: str | None = None,
        file_name: str | None = None,
    ) -> str:
        print(f"Agent question: {question[:80]}...")
        extras = []
        if task_id:
            extras.append(f"task_id: {task_id}")
        if file_name:
            extras.append(f"file_name: {file_name}")
        prompt = question
        if extras:
            prompt = question + "\n\n" + "\n".join(extras)

        payload = {"messages": [{"role": "user", "content": prompt}]}
        # Capped so one runaway question cannot eat the daily Groq token budget.
        config = {"recursion_limit": int(os.getenv("AGENT_MAX_STEPS", "12"))}

        raw = None
        for attempt in (1, 2):
            try:
                result = self.agent.invoke(payload, config)
                raw = result["messages"][-1].content
                break
            except Exception as e:  # noqa: BLE001
                if attempt == 1 and self._switch_to_fallback(e):
                    continue
                print(f"Tool run failed ({type(e).__name__}); answering without tools.")
                raw = self._answer_without_tools(question)
                break

        answer = self._finalize(question, raw)
        print(f"Agent answer: {answer}")
        return answer

    def _answer_without_tools(self, question: str) -> str:
        """A blank submission always scores zero, so fall back to the model's own guess."""
        try:
            reply = self.llm.invoke(
                f"{question}\n\nReply as <answer>THE ANSWER</answer> with a short exact "
                "answer and nothing else. Guess if you are unsure."
            )
            return str(reply.content)
        except Exception:  # noqa: BLE001
            return ""

    def _finalize(self, question: str, raw: object) -> str:
        if isinstance(raw, list):
            # multimodal-style content blocks
            raw = " ".join(
                part.get("text", str(part)) if isinstance(part, dict) else str(part)
                for part in raw
            )
        tagged = re.search(r"<answer>(.*?)</answer>", str(raw), re.S)
        if tagged:
            raw = tagged.group(1)
        answer = _clean_answer(str(raw))
        return self._compress(question, raw) if _is_verbose(answer) else answer

    def _compress(self, question: str, draft: str) -> str:
        """Second pass that turns a sentence or a refusal into a bare answer."""
        text = str(draft)
        try:
            reply = self.llm.invoke(
                EXTRACT_PROMPT.format(question=question, draft=text[:3000])
            )
            text = str(reply.content)
        except Exception:  # noqa: BLE001
            pass

        tagged = re.search(r"<answer>(.*?)</answer>", text, re.S)
        answer = _clean_answer(tagged.group(1) if tagged else text)
        return _salvage(answer, str(draft)) if _is_verbose(answer) else answer