Spaces:
Sleeping
Sleeping
| """ | |
| 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") | |