Blame for agent.py

"""
GAIA benchmark agent: router → plan → [vision|math|search]* → synthesize → confidence → format.
"""

import os
import re
import json
import time
from dotenv import load_dotenv
from langchain_anthropic import ChatAnthropic
from typing_extensions import TypedDict
from prompts import (
    ROUTER_PROMPT,
    PLANNER_PROMPT,
    SYNTHESIS_ANSWER_PROMPT,
    CONFIDENCE_PROMPT,
    FORMAT_ANSWER_PROMPT,
    ANSWER_PROMPT,
)
from langgraph.graph import StateGraph, START, END
from tools import search_web, run_calculations, extract_image_urls, analyze_images

load_dotenv(override=True)

MAX_SEARCH_QUERIES = 3
MIN_SEARCH_QUERIES = 1
MAX_SUPPLEMENTAL_QUERIES = 2
CONFIDENCE_THRESHOLD = 0.72
MAX_LLM_RETRIES = 6
DEFAULT_RETRY_SECONDS = 15
QUESTION_DELAY_SECONDS = float(os.getenv("QUESTION_DELAY_SECONDS", "2"))

VALID_TASK_TYPES = frozenset({"factual", "numeric", "visual", "hybrid"})


def _resolve_anthropic_api_key() -> str:
    key = (os.getenv("ANTHROPIC_API_KEY") or "").strip().strip('"').strip("'")
    if not key:
        raise ValueError(
            "ANTHROPIC_API_KEY is not set. Add sk-ant-... to .env or Space secrets."
        )
    if key.startswith("sess-") or key.startswith("hf_"):
        raise ValueError("ANTHROPIC_API_KEY must be sk-ant-..., not a session/HF token.")
    if not key.startswith("sk-ant-"):
        raise ValueError("ANTHROPIC_API_KEY should start with sk-ant-.")
    return key


model = ChatAnthropic(
    model=os.getenv("ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
    api_key=_resolve_anthropic_api_key(),
    temperature=0,
)


class AgentState(TypedDict, total=False):
    question: str
    task_type: str
    task_reason: str
    plan: list[str]
    search_queries: list[str]
    calculations: list[str]
    image_urls: list[str]
    vision_results: list[str]
    math_results: list[str]
    search_results: list[str]
    synthesis_notes: str
    draft_answer: str
    answer: str
    confidence: float
    confidence_issues: list[str]
    follow_up_queries: list[str]
    retry_count: int


#--- Helpers ------------------------------------------------------------#

def _normalize(text: str) -> str:
    return " ".join(text.lower().split())


def _strip_json_fences(text: str) -> str:
    text = text.strip()
    match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text, re.IGNORECASE)
    if match:
        return match.group(1).strip()
    return text


def _retry_delay_seconds(error: Exception, attempt: int) -> float:
    msg = str(error)
    match = re.search(r"retry in (\d+(?:\.\d+)?)s", msg, re.IGNORECASE)
    if match:
        return float(match.group(1)) + 1.0
    return min(DEFAULT_RETRY_SECONDS * (attempt + 1), 60.0)


def _is_rate_limit_error(error: Exception) -> bool:
    msg = str(error).lower()
    return (
        "429" in msg
        or "rate_limit" in msg
        or "rate limit" in msg
        or "too many requests" in msg
        or "resource_exhausted" in msg
        or "quota" in msg
        or "overloaded" in msg
    )


def _message_content(response) -> str:
    content = response.content
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, str):
                parts.append(block)
            elif isinstance(block, dict) and block.get("text"):
                parts.append(str(block["text"]))
        return "\n".join(parts).strip()
    return str(content).strip()


def invoke_llm(prompt: str) -> str:
    last_error: Exception | None = None
    for attempt in range(MAX_LLM_RETRIES):
        try:
            return _message_content(model.invoke(prompt))
        except Exception as e:
            last_error = e
            if _is_rate_limit_error(e) and attempt < MAX_LLM_RETRIES - 1:
                delay = _retry_delay_seconds(e, attempt)
                print(f"Rate limited; retrying in {delay:.0f}s...")
                time.sleep(delay)
                continue
            raise
    if last_error:
        raise last_error
    raise RuntimeError("invoke_llm failed")


def _is_question_copy(text: str, question: str) -> bool:
    t, q = _normalize(text), _normalize(question)
    if not t or t == q:
        return True
    return len(t) >= len(q) * 0.85 and q in t


def _keyword_fallback_query(question: str) -> str:
    stop = {
        "what", "which", "when", "where", "who", "whom", "whose", "why", "how",
        "the", "a", "an", "is", "are", "was", "were", "of", "in", "on", "at",
        "to", "for", "and", "or", "does", "do", "did", "many", "much",
    }
    words = [
        w.strip(".,?!\"'()[]")
        for w in question.split()
        if w.lower().strip(".,?!\"'()[]") not in stop and len(w) > 2
    ]
    return " ".join(words[:8]) if words else question[:80]


def _classify_task_heuristic(question: str) -> tuple[str, str]:
    """Fast routing when the router LLM call fails."""
    q = question.lower()
    urls = extract_image_urls(question)
    if urls or any(
        w in q
        for w in (
            "image", "picture", "photo", "chart", "graph", "figure", "diagram",
            "screenshot", "plot", "table in the", "shown in",
        )
    ):
        return "visual", "Question references an image or chart."
    if _extract_inline_calculations(question) or any(
        w in q
        for w in (
            "how many", "how much", "calculate", "sum", "total", "average",
            "percent", "ratio", "difference", "multiply", "divide",
        )
    ):
        return "numeric", "Question requires counting or arithmetic."
    if any(
        w in q
        for w in ("what is the capital", "who wrote", "when was", "where is", "name of")
    ):
        return "factual", "Simple factual lookup."
    return "hybrid", "Default multi-tool pipeline."


def _parse_router_response(content: str, question: str) -> dict:
    data = json.loads(_strip_json_fences(content))
    task_type = str(data.get("task_type", "hybrid")).strip().lower()
    if task_type not in VALID_TASK_TYPES:
        task_type, reason = _classify_task_heuristic(question)
        return {"task_type": task_type, "task_reason": reason}
    reason = str(data.get("reason", "")).strip() or f"Classified as {task_type}."
    return {"task_type": task_type, "task_reason": reason}


def _extract_inline_calculations(question: str) -> list[str]:
    """Heuristic: pull simple arithmetic patterns from the question."""
    found: list[str] = []
    for match in re.finditer(
        r"\b\d+(?:\.\d+)?\s*[\+\-\*\/]\s*\d+(?:\.\d+)?(?:\s*[\+\-\*\/]\s*\d+(?:\.\d+)?)*",
        question,
    ):
        expr = re.sub(r"\s+", "", match.group(0))
        if expr and expr not in found:
            found.append(expr)
    return found[:3]


def _parse_planner_response(content: str, question: str) -> dict:
    data = json.loads(_strip_json_fences(content))
    steps = data.get("steps") or []
    queries = data.get("search_queries") or []
    calcs = data.get("calculations") or []
    images = data.get("image_urls") or []

    if not isinstance(steps, list):
        steps = [str(steps)]
    if not isinstance(queries, list):
        queries = [str(queries)]
    if not isinstance(calcs, list):
        calcs = []
    if not isinstance(images, list):
        images = []

    steps = [str(s).strip() for s in steps if str(s).strip()]
    steps = [s for s in steps if not _is_question_copy(s, question)]
    if not steps:
        steps = [
            "Identify entities, dates, and constraints",
            "Gather missing facts from search",
            "Verify numeric relationships with calculations if needed",
            "Form the minimal GAIA answer",
        ]

    queries = [str(q).strip() for q in queries if str(q).strip()]
    queries = [q for q in queries if not _is_question_copy(q, question)]
    seen: set[str] = set()
    deduped: list[str] = []
    for q in queries:
        key = _normalize(q)
        if key in seen or len(q.split()) > 10:
            continue
        seen.add(key)
        deduped.append(q)
        if len(deduped) >= MAX_SEARCH_QUERIES:
            break
    if len(deduped) < MIN_SEARCH_QUERIES:
        deduped = [_keyword_fallback_query(question)]

    calcs = [str(c).strip() for c in calcs if str(c).strip()]
    for expr in _extract_inline_calculations(question):
        if expr not in calcs:
            calcs.append(expr)
    calcs = calcs[:5]

    image_urls = [str(u).strip() for u in images if str(u).strip().startswith("http")]
    for url in extract_image_urls(question):
        if url not in image_urls:
            image_urls.append(url)
    image_urls = image_urls[:3]

    return {
        "steps": steps,
        "search_queries": deduped,
        "calculations": calcs,
        "image_urls": image_urls,
    }


def _format_plan(plan: list[str]) -> str:
    return "\n".join(f"- {s}" for s in plan)


def _format_math(math_results: list[str]) -> str:
    if not math_results:
        return "(no calculations)"
    return "\n".join(f"- {m}" for m in math_results)


def _format_vision(vision_results: list[str]) -> str:
    if not vision_results:
        return "(no images analyzed)"
    return "\n\n".join(vision_results)


def _format_search_results(results: list[str], max_chars: int = 7000) -> str:
    blocks, total = [], 0
    for i, block in enumerate(results, 1):
        chunk = f"### Result block {i}\n{block}"
        if total + len(chunk) > max_chars:
            chunk = chunk[: max_chars - total]
        blocks.append(chunk)
        total += len(chunk)
        if total >= max_chars:
            break
    return "\n\n".join(blocks) if blocks else "(no search results)"


def _heuristic_synthesis(
    search_results: list[str],
    math_results: list[str],
    vision_results: list[str] | None = None,
) -> str:
    bullets: list[str] = []
    for v in vision_results or []:
        bullets.append(f"- {v[:500]}")
    for m in math_results:
        bullets.append(f"- {m}")
    for block in search_results:
        for line in block.split("\n"):
            line = line.strip()
            if line and line != "No results found.":
                bullets.append(f"- {line[:350]}")
            if len(bullets) >= 14:
                break
        if len(bullets) >= 14:
            break
    return "\n".join(bullets) if bullets else "- No usable evidence."


def _parse_synthesis(content: str) -> dict:
    data = json.loads(_strip_json_fences(content))
    notes = str(data.get("synthesis_notes", "")).strip()
    draft = str(data.get("draft_answer", data.get("answer", ""))).strip()
    draft = re.sub(r'^["\']|["\']$', "", draft)
    if not notes:
        raise ValueError("missing synthesis_notes")
    if not draft:
        raise ValueError("missing draft_answer")
    return {"synthesis_notes": notes, "draft_answer": draft}


def _parse_confidence(content: str) -> dict:
    data = json.loads(_strip_json_fences(content))
    conf = float(data.get("confidence", 0.5))
    conf = max(0.0, min(1.0, conf))
    issues = data.get("issues") or []
    if not isinstance(issues, list):
        issues = [str(issues)]
    follow = data.get("follow_up_queries") or []
    if not isinstance(follow, list):
        follow = []
    draft = str(data.get("draft_answer", "")).strip()
    draft = re.sub(r'^["\']|["\']$', "", draft)
    return {
        "confidence": conf,
        "confidence_issues": [str(i) for i in issues if str(i).strip()],
        "follow_up_queries": [str(q).strip() for q in follow if str(q).strip()],
        "draft_answer": draft,
    }


def _clean_final_answer(text: str) -> str:
    text = text.strip()
    text = re.sub(r'^["\']|["\']$', "", text)
    text = re.sub(r"^(the answer is|answer:)\s*", "", text, flags=re.I).strip()
    return text


#--- Nodes --------------------------------------------------------------#

def router_node(state: AgentState) -> AgentState:
    """Classify question type to route the tool pipeline."""
    question = state["question"]
    try:
        parsed = _parse_router_response(
            invoke_llm(ROUTER_PROMPT.format(question=question)),
            question,
        )
    except Exception:
        task_type, reason = _classify_task_heuristic(question)
        parsed = {"task_type": task_type, "task_reason": reason}
    print(f"Router: {parsed['task_type']} — {parsed['task_reason']}")
    return parsed


def plan_node(state: AgentState) -> AgentState:
    question = state["question"]
    task_type = state.get("task_type", "hybrid")
    try:
        data = _parse_planner_response(
            invoke_llm(
                PLANNER_PROMPT.format(question=question, task_type=task_type)
            ),
            question,
        )
    except Exception:
        data = {
            "steps": [
                "Identify entities, dates, and constraints",
                "Search for authoritative facts",
                "Apply calculations if the question is numeric",
            ],
            "search_queries": [_keyword_fallback_query(question)],
            "calculations": _extract_inline_calculations(question),
            "image_urls": extract_image_urls(question),
        }
    return {
        "plan": data["steps"],
        "search_queries": data["search_queries"],
        "calculations": data["calculations"],
        "image_urls": data.get("image_urls", []),
        "retry_count": state.get("retry_count", 0),
    }


def vision_node(state: AgentState) -> AgentState:
    """Analyze images in the question with Claude vision."""
    urls = state.get("image_urls") or []
    if not urls:
        return {"vision_results": []}
    print(f"Analyzing {len(urls)} image(s)...")
    results = analyze_images(urls, state["question"], model)
    return {"vision_results": results}


def math_node(state: AgentState) -> AgentState:
    calcs = state.get("calculations") or []
    return {"math_results": run_calculations(calcs)}


def search_node(state: AgentState) -> AgentState:
    queries = state.get("search_queries") or []
    results = [search_web(q) for q in queries[:MAX_SEARCH_QUERIES]]
    return {"search_results": results}


def synthesize_node(state: AgentState) -> AgentState:
    question = state["question"]
    plan = state.get("plan") or []
    math_results = state.get("math_results") or []
    vision_results = state.get("vision_results") or []
    search_results = state.get("search_results") or []

    try:
        prompt = SYNTHESIS_ANSWER_PROMPT.format(
            question=question,
            plan=_format_plan(plan),
            math_results=_format_math(math_results),
            vision_results=_format_vision(vision_results),
            search_results=_format_search_results(search_results),
        )
        parsed = _parse_synthesis(invoke_llm(prompt))
        return parsed
    except Exception as e:
        print(f"Synthesize failed ({e}); using fallback.")
        notes = _heuristic_synthesis(search_results, math_results, vision_results)
        try:
            draft = _clean_final_answer(
                invoke_llm(
                    ANSWER_PROMPT.format(
                        question=question,
                        plan=_format_plan(plan),
                        synthesis_notes=notes,
                    )
                )
            )
        except Exception:
            draft = "Unknown"
        return {"synthesis_notes": notes, "draft_answer": draft}


def confidence_node(state: AgentState) -> AgentState:
    question = state["question"]
    draft = state.get("draft_answer") or ""
    try:
        prompt = CONFIDENCE_PROMPT.format(
            question=question,
            plan=_format_plan(state.get("plan") or []),
            synthesis_notes=state.get("synthesis_notes") or "",
            math_results=_format_math(state.get("math_results") or []),
            vision_results=_format_vision(state.get("vision_results") or []),
            draft_answer=draft,
        )
        parsed = _parse_confidence(invoke_llm(prompt))
        if parsed["draft_answer"]:
            draft = parsed["draft_answer"]
        return {
            "confidence": parsed["confidence"],
            "confidence_issues": parsed["confidence_issues"],
            "follow_up_queries": parsed["follow_up_queries"][:MAX_SUPPLEMENTAL_QUERIES],
            "draft_answer": draft,
        }
    except Exception as e:
        print(f"Confidence check failed ({e}); assuming moderate confidence.")
        return {
            "confidence": 0.75,
            "confidence_issues": [],
            "follow_up_queries": [],
            "draft_answer": draft,
        }


def supplemental_search_node(state: AgentState) -> AgentState:
    follow = state.get("follow_up_queries") or []
    follow = [q for q in follow if not _is_question_copy(q, state["question"])]
    if not follow:
        return {"retry_count": state.get("retry_count", 0) + 1}
    new_results = [search_web(q) for q in follow[:MAX_SUPPLEMENTAL_QUERIES]]
    existing = state.get("search_results") or []
    return {
        "search_results": existing + new_results,
        "retry_count": state.get("retry_count", 0) + 1,
    }


def format_node(state: AgentState) -> AgentState:
    question = state["question"]
    draft = state.get("draft_answer") or ""
    issues = state.get("confidence_issues") or []
    issues_text = "; ".join(issues) if issues else "(none)"

    try:
        answer = _clean_final_answer(
            invoke_llm(
                FORMAT_ANSWER_PROMPT.format(
                    question=question,
                    draft_answer=draft,
                    issues=issues_text,
                )
            )
        )
    except Exception:
        answer = _clean_final_answer(draft)
    if not answer:
        answer = "Unknown"
    return {"answer": answer}


def route_after_confidence(state: AgentState) -> str:
    if state.get("confidence", 0) >= CONFIDENCE_THRESHOLD:
        return "format"
    if state.get("retry_count", 0) >= 1:
        return "format"
    if state.get("follow_up_queries"):
        return "supplemental_search"
    return "format"


def route_after_plan(state: AgentState) -> str:
    """Pick the first tool node based on task type and planned resources."""
    task = state.get("task_type", "hybrid")
    urls = state.get("image_urls") or []
    calcs = state.get("calculations") or []

    if task == "factual" and not urls and not calcs:
        return "search"
    if task == "visual" or urls:
        return "vision"
    if task == "numeric" or calcs:
        return "math"
    if urls:
        return "vision"
    if calcs:
        return "math"
    return "search"


def route_after_vision(state: AgentState) -> str:
    """After vision, run math only if calculations were planned."""
    if state.get("calculations"):
        return "math"
    return "search"


#--- Graph --------------------------------------------------------------#

agent_builder = StateGraph(AgentState)
agent_builder.add_node("router", router_node)
agent_builder.add_node("plan", plan_node)
agent_builder.add_node("vision", vision_node)
agent_builder.add_node("math", math_node)
agent_builder.add_node("search", search_node)
agent_builder.add_node("synthesize", synthesize_node)
agent_builder.add_node("confidence", confidence_node)
agent_builder.add_node("supplemental_search", supplemental_search_node)
agent_builder.add_node("format", format_node)

agent_builder.add_edge(START, "router")
agent_builder.add_edge("router", "plan")
agent_builder.add_conditional_edges(
    "plan",
    route_after_plan,
    {
        "vision": "vision",
        "math": "math",
        "search": "search",
    },
)
agent_builder.add_conditional_edges(
    "vision",
    route_after_vision,
    {
        "math": "math",
        "search": "search",
    },
)
agent_builder.add_edge("math", "search")
agent_builder.add_edge("search", "synthesize")
agent_builder.add_edge("synthesize", "confidence")
agent_builder.add_conditional_edges(
    "confidence",
    route_after_confidence,
    {
        "format": "format",
        "supplemental_search": "supplemental_search",
    },
)
agent_builder.add_edge("supplemental_search", "synthesize")
agent_builder.add_edge("format", END)

agent = agent_builder.compile()


def run_agent(question: str) -> AgentState:
    """Run the full GAIA pipeline for one question."""
    return agent.invoke(
        {
            "question": question,
            "retry_count": 0,
            "search_results": [],
            "math_results": [],
            "vision_results": [],
        }
    )


if __name__ == "__main__":
    questions = ["What is the capital of France?"]
    for i, question in enumerate(questions):
        if i > 0:
            time.sleep(QUESTION_DELAY_SECONDS)
        result = run_agent(question)
        print(f"Question: {result['question']}")
        print(f"Task: {result.get('task_type')} — {result.get('task_reason')}")
        print(f"Plan: {result.get('plan')}")
        print(f"Calculations: {result.get('calculations')}")
        print(f"Images: {result.get('image_urls')}")
        print(f"Vision: {result.get('vision_results')}")
        print(f"Math: {result.get('math_results')}")
        print(f"Queries: {result.get('search_queries')}")
        print(f"Confidence: {result.get('confidence')}")
        print(f"Synthesis:\n{result.get('synthesis_notes')}")
        print(f"Answer: {result.get('answer')}\n")