agent.py

from __future__ import annotations

import ast
import os
import re
from collections.abc import Iterable
from dataclasses import dataclass, field
from difflib import SequenceMatcher
from functools import lru_cache
from io import BytesIO
from typing import Any
from urllib.parse import parse_qsl, unquote, urlencode, urljoin, urlparse, urlunparse
from uuid import uuid4

import requests
import trafilatura
import wikipedia
import wikipedia.wikipedia as wiki_module
from bs4 import BeautifulSoup
from bs4.element import Tag
from langchain.agents import create_agent
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
from langchain.tools import tool
from langchain_chroma import Chroma
from langchain_community.tools import DuckDuckGoSearchResults, WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_core.messages import SystemMessage, ToolMessage
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from pydantic import BaseModel, Field
from pypdf import PdfReader

# wiki patches
wikipedia.set_lang("en")
wiki_module.API_URL = "https://en.wikipedia.org/w/api.php"

wikipedia.set_user_agent("GAIA-LangChain-Agent/1.0 (contact: placeholder-email@example.com)")


URL_PATTERN = re.compile(r"https?://[^\s<>\"']+", re.IGNORECASE)

MAX_TOOL_CALLS = 5
MAX_SEARCH_CALLS = 2
MAX_WEBPAGE_READS = 2
MAX_AGENT_ITERATIONS = 6
MAX_TOOL_OUTPUT_CHARS = 10_000

_DROP_QUERY_KEYS = {
    "download",
    "display",
    "print",
    "output",
    "view",
    "format",
    "share",
    "source",
    "ref",
    "referrer",
    "fbclid",
    "gclid",
    "mc_cid",
    "mc_eid",
}


def canonicalize_url(url: str) -> str:
    """Return a stable source identity, collapsing common wrappers and variants."""
    value = unquote(url.strip())
    for _ in range(3):
        parsed = urlparse(value)
        host = (parsed.hostname or "").lower()
        if host in {"r.jina.ai", "r.jina.ai/http", "r.jina.ai/https"}:
            wrapped = parsed.path.lstrip("/")
            if wrapped.startswith(("http://", "https://")):
                value = wrapped
                continue
        query = dict(parse_qsl(parsed.query, keep_blank_values=True))
        if host in {"www.google.com", "google.com", "webcache.googleusercontent.com"}:
            wrapped = query.get("url") or query.get("q")
            if wrapped and wrapped.startswith(("http://", "https://")):
                value = wrapped
                continue
        break

    parsed = urlparse(value)
    host = (parsed.hostname or "").lower()
    port = parsed.port
    netloc = host if port in {None, 80, 443} else f"{host}:{port}"
    path = re.sub(r"/(?:print|download|amp)/?$", "", parsed.path, flags=re.I) or "/"
    kept = []
    for key, val in parse_qsl(parsed.query, keep_blank_values=True):
        lowered = key.casefold()
        if lowered.startswith("utm_") or lowered in _DROP_QUERY_KEYS:
            continue
        kept.append((key, val))
    return urlunparse(
        (
            parsed.scheme.lower() or "https",
            netloc,
            path.rstrip("/") or "/",
            "",
            urlencode(sorted(kept)),
            "",
        )
    )


def _normalized_query(query: str) -> str:
    words = re.findall(r"[a-z0-9]+", query.casefold())
    stop = {"a", "an", "and", "for", "in", "of", "on", "the", "to", "what", "who", "is"}
    return " ".join(sorted(word for word in words if word not in stop))


def _queries_similar(left: str, right: str) -> bool:
    a, b = _normalized_query(left), _normalized_query(right)
    if not a or not b:
        return a == b
    aset, bset = set(a.split()), set(b.split())
    jaccard = len(aset & bset) / len(aset | bset)
    return jaccard >= 0.8 or SequenceMatcher(None, a, b).ratio() >= 0.88


@dataclass
class ResearchBudget:
    """Per-agent-run deterministic retrieval accounting."""

    total_calls: int = 0
    search_calls: int = 0
    webpage_reads: int = 0
    model_iterations: int = 0
    output_chars: int = 0
    queries: list[str] = field(default_factory=list)
    urls: set[str] = field(default_factory=set)
    focused_failures: set[str] = field(default_factory=set)

    def reject_tool_call(self, name: str, args: dict[str, Any]) -> str | None:
        if self.total_calls >= MAX_TOOL_CALLS:
            return "LIMIT_REACHED: maximum total tool calls reached; answer from existing evidence."
        self.total_calls += 1
        if name in {"duckduckgo_search", "wikipedia_search"}:
            query = str(args.get("query", ""))
            if self.search_calls >= MAX_SEARCH_CALLS:
                return "LIMIT_REACHED: maximum search calls reached; answer from existing evidence."
            if any(_queries_similar(query, previous) for previous in self.queries):
                return "DUPLICATE: equivalent search already attempted; use existing evidence."
            self.search_calls += 1
            self.queries.append(query)
        elif name == "read_webpage":
            url = canonicalize_url(str(args.get("url", "")))
            focus = str(args.get("focus", "")).strip()
            if self.webpage_reads >= MAX_WEBPAGE_READS:
                return (
                    "LIMIT_REACHED: maximum webpage reads reached; answer from existing evidence."
                )
            if url in self.urls:
                if not focus and url in self.focused_failures:
                    return "DUPLICATE: empty-focus retry forbidden after focused failure."
                return "DUPLICATE: canonical URL already attempted; use existing evidence."
            self.webpage_reads += 1
            self.urls.add(url)
        return None

    def compact_output(self, name: str, args: dict[str, Any], content: str) -> str:
        status = content.split(":", 1)[0]
        if name == "read_webpage" and status in {"NOT_FOUND", "LOW_QUALITY", "ERROR"}:
            canonical = canonicalize_url(str(args.get("url", "")))
            if str(args.get("focus", "")).strip():
                self.focused_failures.add(canonical)
            label = (urlparse(canonical).hostname or "Webpage").removeprefix("www.")
            content = f"{label} read failed: {content[:180]}; canonical URL already attempted."
        remaining = max(0, MAX_TOOL_OUTPUT_CHARS - self.output_chars)
        content = _truncate_text(content, remaining)
        self.output_chars += len(content)
        return (
            content
            or "LIMIT_REACHED: cumulative tool output exhausted; answer from existing evidence."
        )


class FinalAnswer(BaseModel):
    """Exact answer returned to the GAIA evaluator."""

    answer: str = Field(description="Only the exact final answer, with no explanation or label.")


def _extract_urls(text: str) -> list[str]:
    return [match.rstrip(".,;:!?)]") for match in URL_PATTERN.findall(text)]


def _is_youtube_url(url: str) -> bool:
    host = (urlparse(url).hostname or "").lower()
    return host in {"youtube.com", "www.youtube.com", "m.youtube.com", "youtu.be"}


def _normalize_answer(value: Any) -> str:
    text = str(value or "").strip()
    text = re.sub(r"^```(?:text)?\s*", "", text, flags=re.IGNORECASE)
    text = re.sub(r"\s*```$", "", text)
    text = re.sub(
        r"^(?:final\s+answer|answer|response)\s*:\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )
    return text.strip()


_DDG = DuckDuckGoSearchResults(output_format="list", num_results=8)
_WIKIPEDIA = WikipediaQueryRun(
    api_wrapper=WikipediaAPIWrapper(  # type: ignore[call-arg]
        top_k_results=1,
        doc_content_chars_max=10_000,
    )
)


@tool
def duckduckgo_search(query: str) -> str:
    """
    Search the public web and return relevant titles, snippets, and URLs.
    
    Args:
        query: A concise, single-purpose search query using distinctive names,
            titles, dates, identifiers, or exact quoted phrases. Operators such
            as `site:` and `filetype:` may be used when they materially narrow
            the results. Avoid full questions, filler words, unrelated topics,
            and near-duplicate searches.

    Returns:
        Relevant DuckDuckGo search results or a readable error message.
    """
    try:
        raw = _DDG.invoke(query)
        return _format_search_results(raw, query)
    except Exception as exc:
        return f"DuckDuckGo search failed for {query!r}: {type(exc).__name__}: {exc}"


@tool
def wikipedia_search(query: str) -> str:
    """
    Search English Wikipedia for one article or topic.

    Args:
        query: A concise article title, entity name, or topic, such as
            "Mercedes Sosa", "Giganotosaurus", or
            "History of artificial intelligence".

            Do not include URLs, domain names, full questions, or search
            operators such as `site:` and `inurl:`.

    Returns:
        Content from the most relevant English Wikipedia article, or a readable
        error message if retrieval fails.
    """
    try:
        return str(_WIKIPEDIA.invoke(query))
    except Exception as exc:
        return f"Wikipedia search failed for {query!r}: {type(exc).__name__}: {exc}"


def _format_search_results(raw: Any, query: str) -> str:
    """Filter and bound public-search output to three useful results."""
    if isinstance(raw, str):
        try:
            raw = ast.literal_eval(raw)
        except (SyntaxError, ValueError):
            raw = []
    if not isinstance(raw, list):
        return "ERROR: search returned no usable results."

    query_terms = set(re.findall(r"[a-z0-9]+", query.casefold()))
    blocked_hosts = {
        "facebook.com",
        "instagram.com",
        "linkedin.com",
        "pinterest.com",
        "tiktok.com",
        "x.com",
    }
    blocked_terms = {"buy", "cart", "coupon", "job", "jobs", "salary", "shop", "shopping"}
    usable: list[str] = []
    seen: set[str] = set()
    for item in raw:
        if not isinstance(item, dict):
            continue
        title = _clean_text(str(item.get("title") or ""))
        snippet = _clean_text(str(item.get("snippet") or item.get("body") or ""))
        url = str(item.get("link") or item.get("href") or item.get("url") or "").strip()
        parsed = urlparse(url)
        host = (parsed.hostname or "").removeprefix("www.").lower()
        searchable = f"{title} {snippet} {url}".casefold()
        canonical = canonicalize_url(url) if url else ""
        reflected = parsed.path in {"", "/", "/search"} and bool(parse_qsl(parsed.query))
        irrelevant = blocked_terms & set(re.findall(r"[a-z0-9]+", searchable))
        overlap = query_terms & set(re.findall(r"[a-z0-9]+", searchable))
        if (
            not title
            or not snippet
            or not url
            or parsed.scheme not in {"http", "https"}
            or any(host == domain or host.endswith(f".{domain}") for domain in blocked_hosts)
            or reflected
            or irrelevant
            or not overlap
            or canonical in seen
        ):
            continue
        seen.add(canonical)
        usable.append(f"{title}\n{snippet}\n{url}")
        if len(usable) == 3:
            break
    if not usable:
        return "ERROR: search returned no usable results."
    return _truncate_text("\n\n".join(usable), 2_000)


def _clean_text(text: str) -> str:
    """Normalize whitespace while preserving useful line boundaries."""
    lines: list[str] = []

    for line in text.splitlines():
        cleaned = " ".join(line.split())
        if cleaned:
            lines.append(cleaned)

    return "\n".join(lines)


def _normalize_match_text(text: str) -> str:
    """Normalize heading and focus text for conservative matching."""
    text = re.sub(r"\[\s*edit\s*\]", " ", text, flags=re.IGNORECASE)
    text = re.sub(r"\[[0-9]+\]", " ", text)
    return " ".join(re.findall(r"[\w]+", text.casefold(), flags=re.UNICODE))


def _heading_match_score(heading: str, focus: str) -> float:
    """Score exact and strong partial heading matches; return zero for weak ones."""
    normalized_heading = _normalize_match_text(heading)
    normalized_focus = _normalize_match_text(focus)
    if not normalized_heading or not normalized_focus:
        return 0.0
    if normalized_heading == normalized_focus:
        return 100.0

    heading_tokens = set(normalized_heading.split())
    focus_tokens = set(normalized_focus.split())
    overlap = heading_tokens & focus_tokens
    if not overlap:
        return 0.0

    shorter, longer = sorted((normalized_heading, normalized_focus), key=len)
    if shorter in longer:
        # A distinctive one-word focus (for example "November") can identify a
        # longer heading, but short generic fragments should not.
        if len(shorter) >= 5:
            return 80.0 + 15.0 * (len(shorter) / len(longer))
        return 0.0

    focus_coverage = len(overlap) / len(focus_tokens)
    heading_coverage = len(overlap) / len(heading_tokens)
    if len(overlap) >= 2 and focus_coverage >= 0.7 and heading_coverage >= 0.5:
        return 60.0 + 20.0 * focus_coverage + 10.0 * heading_coverage
    return 0.0


def _heading_level(tag: Tag) -> int | None:
    """Return the level of a heading tag or a modern Wikipedia wrapper."""
    if re.fullmatch(r"h[1-6]", tag.name or ""):
        return int(tag.name[1])
    classes = set(tag.get("class") or [])
    if tag.name == "div" and "mw-heading" in classes:
        heading = tag.find(re.compile(r"^h[1-6]$"))
        if isinstance(heading, Tag):
            return int(heading.name[1])
    return None


def _truncate_text(text: str, max_chars: int) -> str:
    """Truncate cleaned text at a line boundary when practical."""
    text = text.strip()
    if len(text) <= max_chars:
        return text
    excerpt = text[:max_chars]
    line_end = excerpt.rfind("\n")
    if line_end >= max_chars // 2:
        excerpt = excerpt[:line_end]
    return excerpt.rstrip()


def _element_text(node: Tag) -> str:
    """Render tables row-by-row and other elements with meaningful line breaks."""
    if node.name == "table":
        rows: list[str] = []
        for row in node.find_all("tr"):
            cells = [
                _clean_text(cell.get_text(" ", strip=True))
                for cell in row.find_all(["th", "td"], recursive=False)
            ]
            if cells:
                rows.append(" | ".join(cells))
        return "\n".join(rows)
    return node.get_text("\n", strip=True)


def _focused_excerpt(
    text: str,
    focus: str,
    *,
    max_chars: int = 4_000,
) -> str:
    """Return a bounded portion of text surrounding the requested topic."""
    text = text.strip()

    if len(text) <= max_chars:
        return text

    if not focus.strip():
        return text[:max_chars]

    normalized_focus = " ".join(focus.casefold().split())
    lowered_text = text.casefold()

    # Prefer an exact phrase match.
    position = lowered_text.find(normalized_focus)

    # Fall back to the line containing the most focus terms.
    if position == -1:
        terms = {term for term in re.findall(r"[a-z0-9]+", normalized_focus) if len(term) >= 3}

        best_score = 0
        best_position = -1
        current_position = 0

        for line in text.splitlines():
            lowered_line = line.lower()
            score = sum(term in lowered_line for term in terms)

            if score > best_score:
                best_score = score
                best_position = current_position

            current_position += len(line) + 1

        position = best_position

    # No meaningful focus match: return a bounded beginning of the page.
    if position == -1:
        return text[:max_chars]

    # Keep limited context before the match and more context after it.
    before_chars = max_chars // 5
    start = max(0, position - before_chars)
    end = min(len(text), start + max_chars)

    # Shift the window backwards near the end of the document.
    if end == len(text):
        start = max(0, end - max_chars)

    excerpt = text[start:end]

    # Avoid beginning or ending in the middle of a line where possible.
    if start > 0 and "\n" in excerpt:
        excerpt = excerpt.split("\n", 1)[1]

    if end < len(text) and "\n" in excerpt:
        excerpt = excerpt.rsplit("\n", 1)[0]

    return _truncate_text(excerpt, max_chars)


def _strong_phrase_present(text: str, focus: str) -> bool:
    normalized_text = _normalize_match_text(text)
    normalized_focus = _normalize_match_text(focus)
    if not normalized_focus:
        return True
    if normalized_focus in normalized_text:
        return True
    focus_terms = set(normalized_focus.split())
    if len(focus_terms) < 2:
        return False
    for line in text.splitlines():
        line_terms = set(_normalize_match_text(line).split())
        if len(focus_terms & line_terms) / len(focus_terms) >= 0.8:
            return True
    return False


def _is_low_quality_extraction(text: str) -> bool:
    """Detect interface/generated markup dominating useful prose."""
    lowered = text.casefold()
    signals = (
        "property get",
        "generated metadata",
        "page actions",
        "navigation menu",
        "toggle navigation",
        "download page",
        "print page",
        "toolbar",
        "contents hide",
        "back to top",
        "skip to main content",
    )
    signal_hits = sum(lowered.count(signal) for signal in signals)
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    nav_labels = {
        "home",
        "contents",
        "search",
        "tools",
        "edit",
        "share",
        "download",
        "print",
        "next",
        "previous",
    }
    nav_count = sum(line.casefold().strip(":") in nav_labels for line in lines)
    words = re.findall(r"[a-z0-9]+", lowered)
    unique_density = len(set(words)) / max(1, len(words))
    prose_lines = sum(len(line.split()) >= 8 for line in lines)
    return (
        signal_hits >= 3
        or (nav_count >= 5 and nav_count >= len(lines) / 3)
        or (len(words) >= 40 and unique_density < 0.12)
        or (len(lines) >= 12 and prose_lines / len(lines) < 0.08 and signal_hits >= 1)
    )


def _trafilatura_excerpt(html: str, url: str, focus: str) -> str | None:
    """Return a validated generic-main-content fallback from existing HTML."""
    extracted = trafilatura.extract(
        html,
        url=url,
        output_format="markdown",
        include_tables=True,
        include_links=True,
        include_comments=False,
        favor_recall=True,
    )
    if not extracted:
        return None
    cleaned = _clean_text(
        "\n".join(
            line
            for line in extracted.splitlines()
            if not re.search(r"\\(?:re)?newcommand|\\definecolor", line)
        )
    )
    if _is_low_quality_extraction(cleaned):
        return None
    if focus.strip() and not _strong_phrase_present(cleaned, focus):
        return _semantic_excerpt(cleaned, focus)
    return _focused_excerpt(cleaned, focus, max_chars=4_000)


def _semantic_excerpt(text: str, focus: str) -> str | None:
    """Use an ephemeral small-embedding index when lexical focus matching fails."""
    chunks: list[str] = []
    start = 0
    bounded = text[:50_000]
    while start < len(bounded):
        end = min(len(bounded), start + 1_200)
        if end < len(bounded):
            boundary = bounded.rfind("\n", start + 600, end)
            if boundary > start:
                end = boundary
        chunk = bounded[start:end].strip()
        if chunk:
            chunks.append(chunk)
        if end == len(bounded):
            break
        start = max(start + 1, end - 200)

    if not chunks:
        return None
    store = Chroma.from_texts(
        chunks,
        embedding=OpenAIEmbeddings(model="text-embedding-3-small"),
        collection_name=f"web-{uuid4().hex}",
        collection_metadata={"hnsw:space": "cosine"},
    )
    try:
        matches = store.similarity_search_with_relevance_scores(focus, k=min(2, len(chunks)))
    finally:
        store.delete_collection()
    selected = [document.page_content for document, score in matches if score >= 0.25]
    if not selected:
        return None
    excerpt = _truncate_text("\n\n".join(selected), 4_000)
    return None if _is_low_quality_extraction(excerpt) else excerpt


def _structural_section(
    content_root: Tag,
    focus: str,
    *,
    max_chars: int = 5_000,
) -> tuple[str, list[Tag]] | None:
    """Extract the best matching heading and its structurally bounded section."""
    if not focus.strip():
        return None

    best: tuple[float, Tag] | None = None
    for heading in content_root.find_all(re.compile(r"^h[1-6]$")):
        if not isinstance(heading, Tag):
            continue
        score = _heading_match_score(heading.get_text(" ", strip=True), focus)
        # A partial page-title match usually identifies the page rather than the
        # requested section. Exact h1 matches remain valid for overview requests.
        if heading.name == "h1" and score < 100:
            score = 0
        if score > 0 and (best is None or score > best[0]):
            best = (score, heading)

    if best is None:
        return None

    heading = best[1]
    level = int(heading.name[1])
    wrapper = heading.parent
    if not (
        isinstance(wrapper, Tag)
        and wrapper.name == "div"
        and "mw-heading" in set(wrapper.get("class") or [])
    ):
        wrapper = heading

    selected: list[Tag] = [wrapper]
    for sibling in wrapper.next_siblings:
        if not isinstance(sibling, Tag):
            continue
        sibling_level = _heading_level(sibling)
        if sibling_level is not None and sibling_level <= level:
            break
        selected.append(sibling)

    section_text = _clean_text("\n".join(_element_text(node) for node in selected))
    return _truncate_text(section_text, max_chars), selected


def _relevant_links(
    roots: Tag | BeautifulSoup | Iterable[Tag],
    *,
    base_url: str,
    focus: str,
    limit: int = 3,
    context_relevant: bool = False,
) -> list[str]:
    """Return a small set of links related to the requested focus."""
    focus_terms = {term for term in re.findall(r"[a-z0-9]+", focus.lower()) if len(term) >= 3}

    if isinstance(roots, (Tag, BeautifulSoup)):
        root_items: Iterable[Tag | BeautifulSoup] = [roots]
    else:
        root_items = roots

    candidates: list[tuple[int, int, str]] = []
    seen: set[str] = set()
    order = 0

    for root in root_items:
        anchors = list(root.find_all("a", href=True))
        if root.name == "a" and root.get("href"):
            anchors = [root, *anchors]
        for anchor in anchors:
            label = " ".join(anchor.get_text(" ", strip=True).split())
            href = str(anchor.get("href") or "").strip()

            if (
                not label
                or not href
                or label.casefold() in {"edit", "reply", "add links", "add languages"}
                or href.startswith(("#", "javascript:", "mailto:"))
                or "action=edit" in href
            ):
                continue

            absolute_url = urljoin(base_url, href)

            if absolute_url in seen:
                continue

            parent_text = ""
            if anchor.parent is not None:
                parent_text = " ".join(anchor.parent.get_text(" ", strip=True).split())[:500]

            searchable = f"{label} {parent_text}".casefold()
            label_text = label.casefold()
            score = sum(2 for term in focus_terms if term in label_text)
            score += sum(1 for term in focus_terms if term in searchable)
            if context_relevant:
                score += 1

            if not focus_terms or score == 0:
                continue

            seen.add(absolute_url)
            candidates.append((score, order, f"{label}: {absolute_url}"))
            order += 1

    candidates.sort(key=lambda item: (-item[0], item[1]))
    return [value for _, _, value in candidates[:limit]]


@tool
def read_webpage(url: str, focus: str) -> str:
    """
    Open a webpage and extract evidence for a specific fact, section, or linked source.

    Args:
    url: Exact webpage URL.


    focus: A concise phrasing likely to occur on the web page to locate using headings, entities, fields,
        relationships, citations, link labels, titles, or facts.


    Returns:
    A compact relevant excerpt and up to three supporting links.


    Status values:
        INVALID_FOCUS
        NOT_FOUND
        LOW_QUALITY
        ERROR
    

    """

    try:
        if not focus.strip():
            return "INVALID_FOCUS: provide a specific heading, phrase, field, or fact."
        response = requests.get(
            url,
            timeout=25,
            headers={"User-Agent": "Mozilla/5.0 GAIA-LangChain-Agent/1.0"},
        )
        response.raise_for_status()

        content_type = response.headers.get("Content-Type", "").lower()
        is_pdf = "application/pdf" in content_type or urlparse(response.url).path.lower().endswith(
            ".pdf"
        )

        if is_pdf:
            reader = PdfReader(BytesIO(response.content))
            raw_text = "\n".join(page.extract_text() or "" for page in reader.pages)
            pdf_text = _clean_text(raw_text)
            if _is_low_quality_extraction(pdf_text):
                return (
                    "LOW_QUALITY: extracted PDF content is dominated by interface "
                    "or generated text."
                )
            if focus.strip() and not _strong_phrase_present(pdf_text, focus):
                return "NOT_FOUND: requested focus is absent from the PDF."
            excerpt = _focused_excerpt(pdf_text, focus, max_chars=4_000)

            return f"OK: FOCUS: {focus or 'document beginning'}\n\nPDF TEXT:\n{excerpt}"

        soup = BeautifulSoup(response.text, "html.parser")
        for node in soup(
            [
                "script",
                "style",
                "nav",
                "footer",
                "aside",
                "noscript",
                "form",
                "svg",
            ]
        ):
            node.decompose()

        content_root = (
            soup.find("main")
            or soup.find(id="mw-content-text")
            or soup.find("article")
            or soup.body
            or soup
        )

        structural = _structural_section(content_root, focus, max_chars=5_000)
        if structural is not None:
            excerpt, link_roots = structural
            links = _relevant_links(
                link_roots,
                base_url=response.url,
                focus=focus,
                limit=3,
                context_relevant=True,
            )
        else:
            page_text = _clean_text(content_root.get_text("\n", strip=True))
            if focus.strip() and not _strong_phrase_present(page_text, focus):
                fallback = _trafilatura_excerpt(response.text, response.url, focus)
                if fallback is None:
                    return "NOT_FOUND: requested focus is absent."
                excerpt, links = fallback, []
            elif _is_low_quality_extraction(page_text):
                fallback = _trafilatura_excerpt(response.text, response.url, focus)
                if fallback is not None:
                    excerpt, links = fallback, []
                else:
                    return (
                        "LOW_QUALITY: extracted content is dominated by navigation, "
                        "metadata, or interface text."
                    )
            else:
                excerpt = _focused_excerpt(page_text, focus, max_chars=4_000)
                links = _relevant_links(
                    content_root,
                    base_url=response.url,
                    focus=focus,
                    limit=3,
                )

        sparse_interface = len(excerpt.split()) < 100 and bool(
            re.search(r"\b(?:last updated|save as pdf|download page)\b", excerpt, re.I)
        )
        if _is_low_quality_extraction(excerpt) or sparse_interface:
            fallback = _trafilatura_excerpt(response.text, response.url, focus)
            if fallback is None:
                return (
                    "LOW_QUALITY: focused extraction is dominated by navigation, "
                    "metadata, or interface text."
                )
            excerpt, links = fallback, []
        excerpt = _truncate_text(excerpt, 4_000)

        rendered = f"OK: FOCUS: {focus or 'page beginning'}\n\nPAGE TEXT:\n{excerpt}"

        if links:
            rendered += "\n\nRELEVANT LINKS:\n" + "\n".join(links)

        return rendered
    except requests.HTTPError as exc:
        return _truncate_text(f"ERROR: HTTP request failed: {type(exc).__name__}: {exc}", 280)
    except requests.RequestException as exc:
        return _truncate_text(f"ERROR: request failed: {type(exc).__name__}: {exc}", 280)
    except Exception as exc:
        return _truncate_text(f"ERROR: processing failed: {type(exc).__name__}: {exc}", 280)


SYSTEM_PROMPT = """
You solve exact-answer questions and return only the required final answer.

## Decide whether tools are needed

Answer directly when the question can be solved from supplied text, tables, lists, arithmetic, logic, classification, or transformation.

Use tools only when external information is required.

## Tool use

* Do not emit planning prose.
* Use concise, high-signal queries containing only relevant names, titles, dates, identifiers, roles, and exact phrases.
* Do not repeat equivalent searches or retry the same failed source.
* Open an existing useful URL instead of searching for it again.
* For historical questions, verify that the source reflects the requested date rather than the current state.
* When reading a page, use a short focus: an exact heading, field, table title, entity, identifier, or distinctive phrase.
* Do not use full instructions or multi-step sentences as `focus`.
* Use the narrowest useful focus.

## Multi-step questions

1. Identify the intermediate entity or source.
2. Retrieve the requested final field or value.
3. Stop when the answer is supported.

Do not add a separate verification call unless the evidence is ambiguous, incomplete, inconsistent, or unreliable.

## After each tool result

Check whether one reliable source provides:

* the correct subject;
* the requested relationship or field;
* the corresponding value;
* the required date, category, or source scope.

If all are present, answer immediately.

If the result is insufficient, change strategy meaningfully. Do not repeat the same query, URL, focus, source variant, or failed approach.

## Verification

Check the requested entity, date range, role, category, table column, spelling, ordering, units, and precision.

For counts, identify the qualifying items before counting.

For lists, include only qualifying items and preserve the requested order and format.

## Evidence rule

Do not invent factual answers.

The final answer must either:

* appear in retrieved evidence; or
* be directly derived from supplied or retrieved evidence.

If no supported answer can be established, return `Unknown`.

## Output

Return only the requested answer.

Do not include explanations, labels, citations, Markdown, or introductory text.

Preserve required spelling, capitalization, punctuation, ordering, delimiters, units, currency symbols, decimal precision, and notation.
""".strip()


class DeterministicGuardrailMiddleware(AgentMiddleware):
    """Enforce retrieval and iteration limits independently of model compliance."""

    def wrap_model_call(self, request: ModelRequest, handler) -> ModelResponse:
        budget: ResearchBudget = request.runtime.context
        budget.model_iterations += 1
        tools = list(request.tools)
        force_final = (
            budget.model_iterations >= MAX_AGENT_ITERATIONS or budget.total_calls >= MAX_TOOL_CALLS
        )
        if force_final:
            tools = []
        else:
            if budget.search_calls >= MAX_SEARCH_CALLS:
                tools = [
                    tool
                    for tool in tools
                    if getattr(tool, "name", "") not in {"duckduckgo_search", "wikipedia_search"}
                ]
            if budget.webpage_reads >= MAX_WEBPAGE_READS:
                tools = [tool for tool in tools if getattr(tool, "name", "") != "read_webpage"]
        if len(tools) != len(request.tools):
            existing = request.system_message.text if request.system_message else ""
            suffix = "\nHard limit reached. Call no unavailable tools; answer now from existing evidence."
            request = request.override(
                tools=tools,
                tool_choice=None,
                system_message=SystemMessage(content=existing + suffix),
            )
        return handler(request)

    def wrap_tool_call(self, request, handler):
        budget: ResearchBudget = request.runtime.context
        name = request.tool_call["name"]
        args = request.tool_call.get("args") or {}
        rejection = budget.reject_tool_call(name, args)
        if rejection:
            return ToolMessage(
                content=rejection, tool_call_id=request.tool_call["id"], status="error"
            )
        result = handler(request)
        if isinstance(result, ToolMessage):
            compacted = budget.compact_output(name, args, str(result.content))
            return result.model_copy(update={"content": compacted})
        return result


@lru_cache(maxsize=1)
def _model() -> ChatOpenAI:
    return ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", "gpt-5"),
        temperature=0,
        max_retries=2,
    )


@lru_cache(maxsize=1)
def _agent():
    return create_agent(
        model=_model(),
        tools=[duckduckgo_search, wikipedia_search, read_webpage],
        system_prompt=SYSTEM_PROMPT,
        middleware=[DeterministicGuardrailMiddleware()],
        context_schema=ResearchBudget,
        response_format=FinalAnswer,
    )


def _direct_openai(question: str) -> str:
    response = _model().invoke(
        [
            ("system", SYSTEM_PROMPT + "\nDo not call tools. Return only the exact answer."),
            ("user", question),
        ]
    )
    return _normalize_answer(response.content)


def _answer_youtube(question: str, video_url: str) -> str:
    from google import genai
    from google.genai import types

    api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
    client = genai.Client(api_key=api_key) if api_key else genai.Client()
    response = client.models.generate_content(
        model=os.getenv("GEMINI_VIDEO_MODEL", "gemini-3.5-flash"),
        contents=types.Content(
            parts=[
                types.Part(file_data=types.FileData(file_uri=video_url)),
                types.Part(
                    text=(
                        question + "\nReturn only the exact requested answer. Do not include an "
                        "explanation, label, citation, or Markdown."
                    )
                ),
            ]
        ),
    )
    return _normalize_answer(response.text or "")


def answer_question(question: str, file_name: str | None = None) -> str | None:
    """Answer one GAIA question, or return None for a named file attachment task."""
    if file_name and file_name.strip():
        return None

    youtube_urls = [url for url in _extract_urls(question) if _is_youtube_url(url)]
    if youtube_urls:
        try:
            return _answer_youtube(question, youtube_urls[0])
        except Exception:
            return _direct_openai(question)

    try:
        result = _agent().invoke(
            {"messages": [{"role": "user", "content": question}]},
            context=ResearchBudget(),
        )
        structured = result.get("structured_response")
        if isinstance(structured, FinalAnswer):
            return _normalize_answer(structured.answer)
        if isinstance(structured, dict) and "answer" in structured:
            return _normalize_answer(structured["answer"])

        messages = result.get("messages") or []
        if messages:
            return _normalize_answer(messages[-1].content)
        raise RuntimeError("LangChain agent returned no final answer")
    except Exception:
        return _direct_openai(question)