steroids

README.md

@@ -13,6 +13,6 @@ hf_oauth_expiration_minutes: 480
 
 # GAIA Agent Evaluation
 
-LangGraph agent with web search for the Hugging Face agents course final assignment.
+LangGraph GAIA agent: router (factual|numeric|visual|hybrid) → plan → conditional tools → synthesize → confidence → format.
 
 **Setup:** Add `ANTHROPIC_API_KEY` in Space secrets (default model: `claude-haiku-4-5-20251001`). Log in with Hugging Face, then run evaluation submit.

agent.py

@@ -1,5 +1,5 @@
 """
-This file contains the agent that will be used to evaluate GAIA benchmark questions.
+GAIA benchmark agent: router → plan → [vision|math|search]* → synthesize → confidence → format.
 """
 
 import os
@@ -9,56 +9,69 @@ import time
 from dotenv import load_dotenv
 from langchain_anthropic import ChatAnthropic
 from typing_extensions import TypedDict
-from prompts import PLANNER_PROMPT, SYNTHESIS_ANSWER_PROMPT, ANSWER_PROMPT
+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
+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:
-    """Use ANTHROPIC_API_KEY from .env (override=True), not IDE session tokens."""
     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 (see .env.example) "
-            "or Space secrets. Create keys at https://console.anthropic.com/settings/keys"
+            "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 an Anthropic key (sk-ant-...), not a Cursor/HF token."
-        )
+        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-. "
-            "Create one at https://console.anthropic.com/settings/keys"
-        )
+        raise ValueError("ANTHROPIC_API_KEY should start with sk-ant-.")
     return key
 
 
-MAX_SEARCH_QUERIES = 3
-MIN_SEARCH_QUERIES = 1
-MAX_LLM_RETRIES = 6
-DEFAULT_RETRY_SECONDS = 15
-QUESTION_DELAY_SECONDS = float(os.getenv("QUESTION_DELAY_SECONDS", "2"))
-
-#--- Model --------------------------------------------------------------#
-
 model = ChatAnthropic(
     model=os.getenv("ANTHROPIC_MODEL", "claude-haiku-4-5-20251001"),
     api_key=_resolve_anthropic_api_key(),
     temperature=0,
 )
 
-#--- State --------------------------------------------------------------#
 
-class AgentState(TypedDict):
+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 ------------------------------------------------------------#
 
@@ -111,34 +124,28 @@ def _message_content(response) -> str:
 
 
 def invoke_llm(prompt: str) -> str:
-    """Call Claude with retries on rate-limit errors."""
     last_error: Exception | None = None
     for attempt in range(MAX_LLM_RETRIES):
         try:
-            response = model.invoke(prompt)
-            return _message_content(response)
+            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 (attempt {attempt + 1}/{MAX_LLM_RETRIES})...")
+                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 without an exception")
+    raise RuntimeError("invoke_llm failed")
 
 
 def _is_question_copy(text: str, question: str) -> bool:
     t, q = _normalize(text), _normalize(question)
-    if not t:
-        return True
-    if t == q:
+    if not t or t == q:
         return True
-    if len(t) >= len(q) * 0.85 and q in t:
-        return True
-    return False
+    return len(t) >= len(q) * 0.85 and q in t
 
 
 def _keyword_fallback_query(question: str) -> str:
@@ -155,24 +162,81 @@ def _keyword_fallback_query(question: str) -> str:
     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:
-    raw = _strip_json_fences(content)
-    data = json.loads(raw)
+    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 and constraints in the question",
-            "Determine what facts must be verified externally",
-            "Cross-check candidates against authoritative sources",
+            "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()]
@@ -187,22 +251,49 @@ def _parse_planner_response(content: str, question: str) -> dict:
         deduped.append(q)
         if len(deduped) >= MAX_SEARCH_QUERIES:
             break
-
     if len(deduped) < MIN_SEARCH_QUERIES:
         deduped = [_keyword_fallback_query(question)]
 
-    return {"steps": steps, "search_queries": deduped}
+    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"- {step}" for step in plan)
+    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_search_results(results: list[str], max_chars: int = 6000) -> str:
-    blocks = []
-    total = 0
+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"### Query {i}\n{block}"
+        chunk = f"### Result block {i}\n{block}"
         if total + len(chunk) > max_chars:
             chunk = chunk[: max_chars - total]
         blocks.append(chunk)
@@ -212,133 +303,349 @@ def _format_search_results(results: list[str], max_chars: int = 6000) -> str:
     return "\n\n".join(blocks) if blocks else "(no search results)"
 
 
-def _heuristic_synthesis(search_results: list[str]) -> str:
+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 not line or line == "No results found.":
-                continue
-            bullets.append(f"- {line[:400]}")
-            if len(bullets) >= 12:
+            if line and line != "No results found.":
+                bullets.append(f"- {line[:350]}")
+            if len(bullets) >= 14:
                 break
-        if len(bullets) >= 12:
+        if len(bullets) >= 14:
             break
-    return "\n".join(bullets) if bullets else "- No usable search evidence."
+    return "\n".join(bullets) if bullets else "- No usable evidence."
 
 
-def _parse_synthesis_answer(content: str) -> dict:
+def _parse_synthesis(content: str) -> dict:
     data = json.loads(_strip_json_fences(content))
     notes = str(data.get("synthesis_notes", "")).strip()
-    answer = str(data.get("answer", "")).strip()
-    answer = re.sub(r'^["\']|["\']$', "", answer)
+    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 answer:
-        raise ValueError("missing answer")
-    return {"synthesis_notes": notes, "answer": answer}
+    if not draft:
+        raise ValueError("missing draft_answer")
+    return {"synthesis_notes": notes, "draft_answer": draft}
 
 
-def _answer_from_notes(question: str, plan: list[str], synthesis_notes: str) -> str:
-    prompt = ANSWER_PROMPT.format(
-        question=question,
-        plan=_format_plan(plan),
-        synthesis_notes=synthesis_notes,
-    )
-    content = invoke_llm(prompt)
-    answer = re.sub(r'^["\']|["\']$', "", content.strip())
-    return answer
+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 plan_node(state: AgentState) -> AgentState:
-    """Plan steps and capped search queries."""
+def router_node(state: AgentState) -> AgentState:
+    """Classify question type to route the tool pipeline."""
     question = state["question"]
-    prompt = PLANNER_PROMPT.format(question=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:
-        content = invoke_llm(prompt)
-        data = _parse_planner_response(content, question)
+        data = _parse_planner_response(
+            invoke_llm(
+                PLANNER_PROMPT.format(question=question, task_type=task_type)
+            ),
+            question,
+        )
     except Exception:
         data = {
             "steps": [
-                "Identify entities and constraints in the question",
-                "Determine what external facts are required",
-                "Verify the answer against reliable sources",
+                "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:
-    """Run 1–3 web searches."""
-    search_queries = state["search_queries"][:MAX_SEARCH_QUERIES]
-    results = [search_web(query) for query in search_queries]
+    queries = state.get("search_queries") or []
+    results = [search_web(q) for q in queries[:MAX_SEARCH_QUERIES]]
     return {"search_results": results}
 
 
-def finalize_node(state: AgentState) -> AgentState:
-    """Synthesize evidence and produce answer in one LLM call (saves quota)."""
+def synthesize_node(state: AgentState) -> AgentState:
     question = state["question"]
-    plan = state["plan"]
-    search_results = state["search_results"]
-    formatted_results = _format_search_results(search_results)
+    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),
-            search_results=formatted_results,
+            math_results=_format_math(math_results),
+            vision_results=_format_vision(vision_results),
+            search_results=_format_search_results(search_results),
         )
-        parsed = _parse_synthesis_answer(invoke_llm(prompt))
+        parsed = _parse_synthesis(invoke_llm(prompt))
         return parsed
     except Exception as e:
-        print(f"Finalize LLM failed ({e}); using heuristic synthesis.")
-        notes = _heuristic_synthesis(search_results)
+        print(f"Synthesize failed ({e}); using fallback.")
+        notes = _heuristic_synthesis(search_results, math_results, vision_results)
         try:
-            answer = _answer_from_notes(question, plan, notes)
+            draft = _clean_final_answer(
+                invoke_llm(
+                    ANSWER_PROMPT.format(
+                        question=question,
+                        plan=_format_plan(plan),
+                        synthesis_notes=notes,
+                    )
+                )
+            )
         except Exception:
-            first_line = notes.split("\n", 1)[0].lstrip("- ").strip()
-            answer = first_line[:200] if first_line else "Unknown"
-        return {"synthesis_notes": notes, "answer": answer}
+            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("finalize", finalize_node)
-agent_builder.add_edge(START, "plan")
-agent_builder.add_edge("plan", "search")
-agent_builder.add_edge("search", "finalize")
-agent_builder.add_edge("finalize", END)
+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 pipeline for a single question (2 LLM calls: plan + finalize)."""
-    return agent.invoke({"question": question})
+    """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__":
-    # Use 1 question while testing to save API usage.
-    questions = [
-        "What is the capital of France?",
-    ]
+    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"Plan: {result['plan']}")
-        print(f"Queries: {result['search_queries']}")
-        print(f"Synthesis:\n{result['synthesis_notes']}")
-        print(f"Answer: {result['answer']}\n")
+        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")

prompts.py

@@ -1,50 +1,94 @@
+ROUTER_PROMPT = """
+You classify GAIA benchmark questions by the primary skill needed to answer.
+
+Output ONLY valid JSON:
+{{
+  "task_type": "factual|numeric|visual|hybrid",
+  "reason": "one short sentence"
+}}
+
+Task types:
+- factual: lookup a name, place, date, definition, or single known fact (web search is enough).
+- numeric: counting, arithmetic, units, percentages, or multi-step math (calculator needed).
+- visual: answer depends on an image, chart, diagram, screenshot, or figure (vision needed).
+- hybrid: needs multiple skills (e.g. read a chart AND search AND calculate).
+
+Pick exactly one task_type. Use hybrid only when two or more of search, math, and vision are clearly required.
+
+Question:
+{question}
+"""
+
 PLANNER_PROMPT = """
-You are a planning agent for a GAIA-style question-answering benchmark.
+You are a planning agent for the GAIA benchmark (level 1).
+
+Task type for this question: {task_type}
 
-Break the question into:
-1. A step-by-step reasoning plan (3–6 atomic steps)
-2. Web search queries to gather missing facts
+Analyze the question and output JSON only (no markdown fences).
 
-Hard rules:
-- Output ONLY valid JSON. No markdown, no code fences, no commentary.
-- Use EXACTLY this schema (array lengths: steps 3–6, search_queries 1–3):
+Schema (all four keys required):
 {{
-  "steps": ["..."],
-  "search_queries": ["..."]
+  "steps": ["3-6 actionable reasoning steps"],
+  "search_queries": ["1-3 short keyword queries"],
+  "calculations": ["0-5 math expressions, or []"],
+  "image_urls": ["0-3 image URLs from the question, or []"]
 }}
-- NEVER copy or paraphrase the full question as a step or as a search query.
-- NEVER use the question text verbatim in search_queries.
-- Steps must be actionable reasoning actions (e.g. "Identify the entity", "Verify the date from an authoritative source").
-- search_queries: 1–3 items only. Short keyword phrases (2–8 words). Each query must target a distinct fact gap.
-- Prefer precision over breadth: one excellent query beats three redundant ones.
-- If the question is answerable without search, still return 1 minimal disambiguation query.
-
-Penalties (invalid output — avoid):
-- steps or search_queries containing the full question string
-- search_queries longer than 10 words
-- more than 3 search_queries or fewer than 1 search_queries
-- placeholder text like "step1", "query1", "..."
+
+Rules for steps:
+- Atomic, specific (identify entities, dates, units, constraints).
+- If the question needs counting, differences, ratios, or multi-step arithmetic, include a step to verify with calculation.
+
+Rules for search_queries:
+- 1-3 items, 2-8 words each, distinct fact gaps.
+- NEVER copy the full question verbatim.
+- Target authoritative facts (names, dates, numbers, definitions).
+
+Rules for calculations:
+- List plain math expressions the answer depends on (e.g. "1250-380", "24*365", "sqrt(144)").
+- Use only numbers and operators + - * / // ** % and functions sqrt, abs, round.
+- Empty array [] if no math is needed.
+
+Rules for image_urls:
+- Copy any http(s) image links from the question (jpg, png, gif, webp, or dataset/file URLs).
+- If the question refers to "the image", "this figure", "the chart", or an attachment URL, include it.
+- Max 3 URLs. Empty array [] if no image is involved.
 
 Question:
 {question}
 """
 
-SYNTHESIS_ANSWER_PROMPT = """
-You are a GAIA research and answer agent. Do two tasks in one response.
+IMAGE_ANALYSIS_PROMPT = """
+You are a vision analyst for GAIA benchmark questions.
+
+Describe everything in the image that could help answer the question:
+- Text (OCR): transcribe labels, numbers, titles, legends, axes, units.
+- Objects, people, places, charts, tables, maps.
+- Counts, colors, dates, and relationships visible in the image.
+- If a chart/graph: report series names, values, trends, min/max.
+
+Be factual and specific. Do not answer the question yet — only report visual evidence.
+
+Question:
+{question}
 
-Task 1 — Synthesis: compress search results into bullet notes (max 12, prefix "- ").
-Task 2 — Answer: produce the final GAIA answer from those notes.
+Image URL:
+{url}
+"""
+
+SYNTHESIS_ANSWER_PROMPT = """
+You are a GAIA research agent. Use search results, math tool outputs, image vision notes, and the plan to answer.
 
-Answer rules:
-- Return ONLY valid JSON. No markdown fences.
-- Schema:
+Return ONLY valid JSON:
 {{
-  "synthesis_notes": "- fact1\\n- fact2",
-  "answer": "final answer only"
+  "synthesis_notes": "- bullet facts (max 14, prefix '- ')",
+  "draft_answer": "best answer before final formatting"
 }}
-- synthesis_notes: bullets only, no preamble.
-- answer: concise; no quotes, meta-commentary, or search artifacts (no broken fragments like "Rome., is good?").
-- Do NOT repeat the question text in the answer.
+
+Rules:
+- synthesis_notes: relevant facts only; note conflicts if sources disagree.
+- draft_answer: concise, directly answers the question; use math results when provided.
+- For numeric questions, compute carefully and show the number you believe is correct.
+- Do not copy long search snippets into draft_answer.
 
 Question:
 {question}
@@ -52,23 +96,88 @@ Question:
 Plan:
 {plan}
 
+Math tool results:
+{math_results}
+
+Image vision analysis:
+{vision_results}
+
 Search results:
 {search_results}
 """
 
+CONFIDENCE_PROMPT = """
+You are a GAIA answer verifier. Judge whether the draft answer is well-supported and correctly formatted.
+
+Return ONLY valid JSON:
+{{
+  "confidence": 0.0,
+  "issues": ["list of problems, or empty"],
+  "follow_up_queries": ["0-2 NEW search queries if evidence is weak, else empty"],
+  "draft_answer": "revised short answer if you can improve it, else repeat the current draft"
+}}
+
+Scoring confidence (0.0-1.0):
+- 0.9+: answer directly matches question, supported by notes/math, format likely correct for GAIA.
+- 0.7-0.89: plausible but missing verification or ambiguous evidence.
+- Below 0.7: wrong, vague, unsupported, or format wrong — add follow_up_queries to fill gaps.
+
+GAIA format checks:
+- Number questions: answer should be a number (or number with unit if question asks).
+- Yes/no: only yes or no.
+- Name/place: proper noun, minimal words.
+- List: match requested format.
+- No sentences from search titles, no "based on", no hedging.
+
+Question:
+{question}
+
+Plan:
+{plan}
+
+Synthesis notes:
+{synthesis_notes}
+
+Math results:
+{math_results}
+
+Image vision analysis:
+{vision_results}
+
+Current draft answer:
+{draft_answer}
+"""
+
+FORMAT_ANSWER_PROMPT = """
+You format the final GAIA submission answer. Output ONLY the answer string, nothing else.
+
+Strict GAIA rules:
+- Minimal: shortest correct form (usually 1-5 words, or a single number).
+- If the question asks for a number only, output digits (and decimal point if needed) — no units unless explicitly required.
+- If yes/no, output exactly: yes or no (lowercase).
+- No preamble, quotes, explanation, or "The answer is".
+- No trailing period unless the answer is a full sentence the question requires.
+- Strip Wikipedia-style titles, colons, or citation artifacts.
+- Prefer the draft if correct; fix format only.
+
+Question:
+{question}
+
+Draft answer:
+{draft_answer}
+
+Confidence notes (issues to fix if any):
+{issues}
+"""
+
 ANSWER_PROMPT = """
 You are a GAIA benchmark answer agent. Produce the final answer only.
 
 Rules:
 - Use the question, plan, and synthesis notes — not raw search snippets.
-- Return ONLY the final answer text. No preamble, labels, quotes, or explanation unless the question requires it.
-- Be concise: prefer a short phrase, name, number, or yes/no when appropriate.
-- Do NOT echo search artifacts, typos, or broken fragments (e.g. "Rome., is good?", trailing punctuation glitches).
-- Do NOT repeat the question, plan, or notes.
-- Do NOT hedge with "based on the search results" or similar meta-commentary.
-- For factual questions: give the single best-supported answer.
-- For lists: use the format the question implies (comma-separated or one item per line).
-- If evidence is insufficient, give the most likely short answer; do not invent citations.
+- Return ONLY the final answer text.
+- Be concise: short phrase, name, number, or yes/no.
+- Do NOT echo search artifacts or broken fragments.
 
 Question:
 {question}

tools.py

@@ -1,13 +1,199 @@
+"""Tools for the GAIA agent: web search, arithmetic, and image recognition."""
+
+import ast
+import base64
+import math
+import operator
+import re
+from typing import Any
+
+import requests
 from ddgs import DDGS
+from langchain_core.messages import HumanMessage
+from prompts import IMAGE_ANALYSIS_PROMPT
+
+MAX_IMAGES_PER_QUESTION = 3
+_IMAGE_EXTENSIONS = r"(?:jpg|jpeg|png|gif|webp|bmp|svg)"
+_URL_TAIL = r"(?:\?[^\s\)\]\"\'<>]*)?"
+
+# Safe operators for GAIA-style numeric questions
+_SAFE_OPERATORS = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.FloorDiv: operator.floordiv,
+    ast.Pow: operator.pow,
+    ast.Mod: operator.mod,
+    ast.USub: operator.neg,
+    ast.UAdd: operator.pos,
+}
+
+
+def _eval_ast(node: ast.AST) -> float:
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, (int, float)):
+            return float(node.value)
+        raise ValueError(f"Unsupported constant: {node.value}")
+    if isinstance(node, ast.Num):  # pragma: no cover - py<3.8
+        return float(node.n)
+    if isinstance(node, ast.BinOp):
+        op = _SAFE_OPERATORS.get(type(node.op))
+        if not op:
+            raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
+        return op(_eval_ast(node.left), _eval_ast(node.right))
+    if isinstance(node, ast.UnaryOp):
+        op = _SAFE_OPERATORS.get(type(node.op))
+        if not op:
+            raise ValueError(f"Unsupported unary op: {type(node.op).__name__}")
+        return op(_eval_ast(node.operand))
+    if isinstance(node, ast.Call):
+        if isinstance(node.func, ast.Name):
+            name = node.func.id
+            args = [_eval_ast(a) for a in node.args]
+            if name == "sqrt" and len(args) == 1:
+                return math.sqrt(args[0])
+            if name == "round" and len(args) in (1, 2):
+                return float(round(args[0], int(args[1]) if len(args) == 2 else 0))
+            if name == "abs" and len(args) == 1:
+                return abs(args[0])
+        raise ValueError("Unsupported function call")
+    raise ValueError(f"Unsupported expression node: {type(node).__name__}")
+
+
+def _normalize_expression(expr: str) -> str:
+    expr = expr.strip()
+    expr = expr.replace("×", "*").replace("÷", "/").replace("^", "**")
+    expr = re.sub(r"\s+", "", expr)
+    return expr
+
+
+def calculate(expression: str) -> str:
+    """
+    Evaluate a numeric expression safely (numbers, + - * / // ** %, sqrt, abs, round).
+    Returns a string result for the agent to use as evidence.
+    """
+    expr = _normalize_expression(expression)
+    if not expr:
+        return "Error: empty expression"
+    try:
+        tree = ast.parse(expr, mode="eval")
+        result = _eval_ast(tree.body)
+        if result == int(result) and "." not in str(expr).split("/")[-1]:
+            return str(int(result))
+        return str(round(result, 10)).rstrip("0").rstrip(".")
+    except Exception as e:
+        return f"Error: {e}"
+
+
+def run_calculations(expressions: list[str]) -> list[str]:
+    """Run multiple calculations; skip empty or duplicate expressions."""
+    seen: set[str] = set()
+    outputs: list[str] = []
+    for raw in expressions:
+        expr = str(raw).strip()
+        if not expr or expr in seen:
+            continue
+        seen.add(expr)
+        result = calculate(expr)
+        outputs.append(f"{expr} = {result}")
+    return outputs
+
+
+def _message_content(response: Any) -> 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 extract_image_urls(text: str) -> list[str]:
+    """Find image URLs in question text (markdown, HTML, plain links)."""
+    candidates: list[str] = []
+    patterns = [
+        rf"https?://[^\s\)\]\"\'<>]+\.{_IMAGE_EXTENSIONS}{_URL_TAIL}",
+        r"!\[[^\]]*\]\((https?://[^\)]+)\)",
+        r"https?://huggingface\.co/datasets/[^\s\)\]\"\'<>]+",
+        r"https?://[^\s\)\]\"\'<>]*(?:/image|/images|/file|/files)[^\s\)\]\"\'<>]*",
+    ]
+    for pattern in patterns:
+        for match in re.finditer(pattern, text, re.IGNORECASE):
+            url = match.group(1) if match.lastindex else match.group(0)
+            url = url.rstrip(".,;:")
+            if url.startswith("http"):
+                candidates.append(url)
+
+    seen: set[str] = set()
+    deduped: list[str] = []
+    for url in candidates:
+        key = url.lower()
+        if key not in seen:
+            seen.add(key)
+            deduped.append(url)
+        if len(deduped) >= MAX_IMAGES_PER_QUESTION:
+            break
+    return deduped
+
+
+def _fetch_image_base64(url: str) -> tuple[str, str]:
+    """Download image and return (media_type, base64_data)."""
+    headers = {"User-Agent": "Mozilla/5.0 (compatible; GAIA-Agent/1.0)"}
+    resp = requests.get(url, timeout=45, headers=headers)
+    resp.raise_for_status()
+    content_type = (resp.headers.get("content-type") or "image/jpeg").split(";")[0].strip()
+    if not content_type.startswith("image/"):
+        content_type = "image/jpeg"
+    return content_type, base64.standard_b64encode(resp.content).decode("ascii")
+
+
+def describe_image(url: str, question: str, vision_model: Any) -> str:
+    """
+    Analyze an image URL with Claude vision and return a text description.
+    """
+    try:
+        media_type, data = _fetch_image_base64(url)
+        prompt = IMAGE_ANALYSIS_PROMPT.format(question=question, url=url)
+        message = HumanMessage(
+            content=[
+                {"type": "text", "text": prompt},
+                {
+                    "type": "image",
+                    "source": {
+                        "type": "base64",
+                        "media_type": media_type,
+                        "data": data,
+                    },
+                },
+            ]
+        )
+        return _message_content(vision_model.invoke([message]))
+    except Exception as e:
+        return f"Vision error: {e}"
+
+
+def analyze_images(urls: list[str], question: str, vision_model: Any) -> list[str]:
+    """Run vision on up to MAX_IMAGES_PER_QUESTION URLs."""
+    results: list[str] = []
+    for url in urls[:MAX_IMAGES_PER_QUESTION]:
+        description = describe_image(url, question, vision_model)
+        results.append(f"[Image: {url}]\n{description}")
+    return results
+
 
 def search_web(query: str) -> str:
     """Search the web for information related to the query."""
-    results = DDGS().text(query, max_results=10)
+    results = DDGS().text(query, max_results=8)
     if not results:
         return "No results found."
     return "\n\n".join(
         f"{r.get('title', '')}: {r.get('body', r.get('href', ''))}"
         for r in results
     )
-
-