""" 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")