Slim graph + cap tool rounds for gateway latency; pin gpt-oss-120b

Drop planner LLM (always probe for file), merge evidence into solver, make formatter pure-Python: ~8-12 calls/question -> ~3-6. Add max_tool_rounds cap so research-heavy questions terminate fast (no more 300s timeouts). Pin openai/gpt-oss-120b on the gateway (fast + reliable tool/structured); force function_calling for structured output since routed models ignore json_schema.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

gaia_agent/config.py

@@ -35,7 +35,9 @@ class Settings(BaseSettings):
     # --- OpenAI-compatible gateway (freellmapi) ---
     openai_compatible_base_url: str = ""  # e.g. https://my-gateway.example.com/v1
     openai_compatible_api_key: str = ""
-    openai_compatible_model: str = "auto"
+    # Pin a model that reliably does tool calls + function-calling structured output.
+    # "auto"/Balanced routing bounces to models that don't, breaking the graph.
+    openai_compatible_model: str = "openai/gpt-oss-120b"
     groq_vision_model: str = "meta-llama/llama-4-scout-17b-16e-instruct"
     groq_whisper_model: str = "whisper-large-v3"
 
@@ -46,6 +48,9 @@ class Settings(BaseSettings):
     # --- API + control knobs ---
     gaia_api_url: str = "https://agents-course-unit4-scoring.hf.space"
     max_judge_retries: int = 1
+    # Hard cap on research tool-calling rounds; after this, the model must answer
+    # without tools. Bounds worst-case latency on research-heavy questions.
+    max_tool_rounds: int = 3
     recursion_limit: int = 40
     question_timeout: int = 300
 

gaia_agent/graph.py

@@ -10,14 +10,11 @@ from langgraph.graph import END, START, StateGraph
 from langgraph.prebuilt import ToolNode
 
 from gaia_agent.nodes import (
-    evidence,
     formatter,
     ingest_file,
     judge,
-    planner,
     research,
     route_after_judge,
-    route_after_planner,
     route_after_research,
     solver,
 )
@@ -29,25 +26,19 @@ def build_graph():
     """Build and compile the GAIA agent graph."""
     g = StateGraph(GraphState)
 
-    g.add_node("planner", planner)
     g.add_node("ingest_file", ingest_file)
     g.add_node("research", research)
     g.add_node("tools", ToolNode(RESEARCH_TOOLS))
-    g.add_node("evidence", evidence)
     g.add_node("solver", solver)
     g.add_node("judge", judge)
     g.add_node("formatter", formatter)
 
-    g.add_edge(START, "planner")
-    g.add_conditional_edges(
-        "planner", route_after_planner, {"ingest_file": "ingest_file", "research": "research"}
-    )
+    g.add_edge(START, "ingest_file")
     g.add_edge("ingest_file", "research")
     g.add_conditional_edges(
-        "research", route_after_research, {"tools": "tools", "evidence": "evidence"}
+        "research", route_after_research, {"tools": "tools", "solver": "solver"}
     )
     g.add_edge("tools", "research")
-    g.add_edge("evidence", "solver")
     g.add_edge("solver", "judge")
     g.add_conditional_edges(
         "judge", route_after_judge, {"research": "research", "formatter": "formatter"}

gaia_agent/llm.py

@@ -72,6 +72,19 @@ def get_text_llm(temperature: float = 0.0):
     )
 
 
+def get_structured_llm(schema, temperature: float = 0.0):
+    """Return an LLM that emits an instance of ``schema`` (Pydantic).
+
+    On the openai_compatible gateway the routed model often ignores json_schema
+    response_format and returns prose, so we force tool/function-calling-based
+    structured output (which the gateway models do support).
+    """
+    llm = get_text_llm(temperature)
+    if get_settings().llm_provider.lower() == "openai_compatible":
+        return llm.with_structured_output(schema, method="function_calling")
+    return llm.with_structured_output(schema)
+
+
 @lru_cache(maxsize=1)
 def get_vision_llm():
     """Return the Groq multimodal model used for image understanding."""

gaia_agent/nodes.py

@@ -6,26 +6,22 @@ Flow: planner -> [ingest_file] -> research <-> tools -> evidence -> solver
 
 from __future__ import annotations
 
+import re
+
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
 from gaia_agent.config import get_settings
-from gaia_agent.llm import get_text_llm
+from gaia_agent.llm import get_structured_llm, get_text_llm
 from gaia_agent.prompts import (
-    EVIDENCE_PROMPT,
-    FORMATTER_PROMPT,
     GAIA_RULES,
     JUDGE_PROMPT,
-    PLANNER_PROMPT,
     RESEARCH_PROMPT,
     SOLVER_PROMPT,
 )
 from gaia_agent.schemas import (
     Candidate,
-    Evidence,
     FileExtract,
-    FinalAnswer,
     JudgeVerdict,
-    Plan,
 )
 from gaia_agent.state import GraphState
 from gaia_agent.tools import (
@@ -39,39 +35,20 @@ from gaia_agent.tools import (
 )
 
 # --------------------------------------------------------------------------- #
-# Planner
-# --------------------------------------------------------------------------- #
-
-
-def planner(state: GraphState) -> dict:
-    """Produce a structured Plan and initialise loop bookkeeping."""
-    try:
-        llm = get_text_llm().with_structured_output(Plan)
-        plan = llm.invoke(
-            PLANNER_PROMPT.format(question=state["question"], task_id=state.get("task_id", ""))
-        )
-    except Exception:  # noqa: BLE001
-        plan = Plan(needs_file=False, reasoning="planner-fallback")
-    return {"plan": plan, "attempts": 0, "context_notes": ""}
-
-
-def route_after_planner(state: GraphState) -> str:
-    """Go fetch the file only if the plan says one is needed."""
-    plan = state.get("plan")
-    return "ingest_file" if (plan and plan.needs_file) else "research"
-
-
-# --------------------------------------------------------------------------- #
-# File ingestion
+# File ingestion (entry node: always probe for an attached file)
 # --------------------------------------------------------------------------- #
 
 
 def ingest_file(state: GraphState) -> dict:
-    """Download and extract the task's attached file into context_notes."""
+    """Download and extract the task's attached file (if any) into context_notes.
+
+    Always attempts the fetch (cheap HTTP) instead of asking an LLM whether a file
+    is needed -- saves a call per question and never misses an attachment.
+    """
     task_id = state.get("task_id", "")
     path = fetch_task_file(task_id)
     if not path:
-        return {"file_extract": None}
+        return {"file_extract": None, "attempts": 0, "context_notes": ""}
 
     kind = classify_file(path)
     question = state["question"]
@@ -90,8 +67,8 @@ def ingest_file(state: GraphState) -> dict:
         extracted = f"File extraction failed: {exc}"
 
     fe = FileExtract(kind=kind, summary=f"Attached {kind} file.", extracted=str(extracted)[:8000])
-    notes = (state.get("context_notes", "") + f"\n[FILE:{kind}]\n{fe.extracted}").strip()
-    return {"file_extract": fe, "context_notes": notes}
+    notes = f"[FILE:{kind}]\n{fe.extracted}".strip()
+    return {"file_extract": fe, "context_notes": notes, "attempts": 0}
 
 
 # --------------------------------------------------------------------------- #
@@ -122,8 +99,16 @@ def _invoke_with_tools(llm, msgs: list):
 
 
 def research(state: GraphState) -> dict:
-    """Run one step of the research tool-calling loop."""
-    llm = get_text_llm().bind_tools(RESEARCH_TOOLS)
+    """Run one step of the research tool-calling loop.
+
+    Tools are bound only while under the max_tool_rounds cap; past it, the model is
+    called without tools so it must answer, guaranteeing the loop terminates fast.
+    """
+    rounds = state.get("tool_rounds", 0)
+    base = get_text_llm()
+    use_tools = rounds < get_settings().max_tool_rounds
+    llm = base.bind_tools(RESEARCH_TOOLS) if use_tools else base
+
     if not state.get("messages"):
         system = RESEARCH_PROMPT.format(
             rules=GAIA_RULES,
@@ -138,47 +123,34 @@ def research(state: GraphState) -> dict:
             ),
         ]
         ai = _invoke_with_tools(llm, seed)
-        return {"messages": seed + [ai]}
+        return {"messages": seed + [ai], "tool_rounds": rounds + 1}
     ai = _invoke_with_tools(llm, state["messages"])
-    return {"messages": [ai]}
+    return {"messages": [ai], "tool_rounds": rounds + 1}
 
 
 def route_after_research(state: GraphState) -> str:
-    """Route to tools if the model requested any, else summarise evidence."""
+    """Route to tools if the model requested any, else go straight to the solver."""
     last = state["messages"][-1]
     if getattr(last, "tool_calls", None):
         return "tools"
-    return "evidence"
+    return "solver"
 
 
 # --------------------------------------------------------------------------- #
-# Evidence / Solver / Judge / Formatter
+# Solver / Judge / Formatter
 # --------------------------------------------------------------------------- #
 
 
-def evidence(state: GraphState) -> dict:
-    """Distil the research conversation into structured Evidence."""
-    try:
-        llm = get_text_llm().with_structured_output(Evidence)
-        ev = llm.invoke([SystemMessage(EVIDENCE_PROMPT), *state.get("messages", [])])
-    except Exception:  # noqa: BLE001
-        ev = Evidence(findings=[], sources=[], confidence=0.0)
-    return {"evidence": ev}
-
-
 def solver(state: GraphState) -> dict:
-    """Synthesize evidence + context into a single Candidate answer."""
-    ev = state.get("evidence")
+    """Synthesize the research conversation + context into one Candidate answer."""
     try:
-        llm = get_text_llm().with_structured_output(Candidate)
-        cand = llm.invoke(
-            SOLVER_PROMPT.format(
-                rules=GAIA_RULES,
-                question=state["question"],
-                evidence=ev.model_dump() if ev else "{}",
-                context_notes=state.get("context_notes", "") or "(none)",
-            )
+        llm = get_structured_llm(Candidate)
+        system = SOLVER_PROMPT.format(
+            rules=GAIA_RULES,
+            question=state["question"],
+            context_notes=state.get("context_notes", "") or "(none)",
         )
+        cand = llm.invoke([SystemMessage(system), *state.get("messages", [])])
     except Exception:  # noqa: BLE001
         cand = Candidate(answer="", justification="solver-fallback")
     return {"candidate": cand}
@@ -187,14 +159,13 @@ def solver(state: GraphState) -> dict:
 def judge(state: GraphState) -> dict:
     """LLM-as-judge: PASS, or REVISE with feedback fed back into research."""
     cand = state.get("candidate")
-    ev = state.get("evidence")
     try:
-        llm = get_text_llm().with_structured_output(JudgeVerdict)
+        llm = get_structured_llm(JudgeVerdict)
         verdict = llm.invoke(
             JUDGE_PROMPT.format(
                 rules=GAIA_RULES,
                 question=state["question"],
-                evidence=ev.model_dump() if ev else "{}",
+                context_notes=state.get("context_notes", "") or "(none)",
                 candidate=cand.answer if cand else "",
             )
         )
@@ -227,17 +198,13 @@ def route_after_judge(state: GraphState) -> str:
 
 
 def formatter(state: GraphState) -> dict:
-    """Render the candidate into the bare exact-match answer string."""
+    """Pure-Python scrub of the candidate into the bare exact-match string (no LLM).
+
+    The solver already formats per GAIA rules; this just strips stray prefixes,
+    quotes, and whitespace.
+    """
     cand = state.get("candidate")
-    candidate_text = cand.answer if cand else ""
-    try:
-        llm = get_text_llm().with_structured_output(FinalAnswer)
-        final = llm.invoke(
-            FORMATTER_PROMPT.format(
-                rules=GAIA_RULES, question=state["question"], candidate=candidate_text
-            )
-        )
-        answer = final.answer
-    except Exception:  # noqa: BLE001
-        answer = candidate_text
-    return {"final_answer": (answer or candidate_text or "").strip()}
+    text = (cand.answer if cand else "") or ""
+    text = re.sub(r"(?is)^\s*(final answer|answer)\s*:?\s*", "", text).strip()
+    text = text.strip().strip('"').strip("'").strip()
+    return {"final_answer": text}

gaia_agent/prompts.py

@@ -50,17 +50,15 @@ their sources, and your confidence (0-1) that they pin down the answer.
 """
 
 SOLVER_PROMPT = """\
-You are the SOLVER. Using the evidence and context below, produce ONE candidate answer.
-Apply the GAIA formatting rules to your candidate, but you may keep a short justification.
+You are the SOLVER. Using the research conversation above and the context below,
+produce ONE candidate answer. The `answer` field MUST already obey the GAIA
+exact-match formatting rules (bare answer, no prose, no "FINAL ANSWER").
 
 {rules}
 
 Question:
 {question}
 
-Evidence:
-{evidence}
-
 Context notes:
 {context_notes}
 """
@@ -75,8 +73,8 @@ rules. Be strict: a near-miss in format is a REVISE.
 Question:
 {question}
 
-Evidence:
-{evidence}
+Context notes:
+{context_notes}
 
 Candidate answer:
 {candidate}

gaia_agent/state.py

@@ -42,6 +42,8 @@ class GraphState(TypedDict, total=False):
 
     # Eval-loop bookkeeping
     attempts: int
+    # Research tool-calling rounds used (capped by max_tool_rounds)
+    tool_rounds: int
 
     # Output
     final_answer: str

tests/test_graph.py

@@ -6,10 +6,9 @@ from unittest.mock import MagicMock, patch
 from gaia_agent.nodes import (
     judge,
     route_after_judge,
-    route_after_planner,
     route_after_research,
 )
-from gaia_agent.schemas import JudgeVerdict, Plan
+from gaia_agent.schemas import JudgeVerdict
 
 
 def test_graph_compiles():
@@ -17,19 +16,14 @@ def test_graph_compiles():
 
     assert graph.name == "GAIA Agent"
     nodes = set(graph.get_graph().nodes.keys())
-    assert {"planner", "research", "tools", "judge", "formatter"} <= nodes
+    assert {"ingest_file", "research", "tools", "solver", "judge", "formatter"} <= nodes
 
 
-def test_route_after_planner():
-    assert route_after_planner({"plan": Plan(needs_file=True)}) == "ingest_file"
-    assert route_after_planner({"plan": Plan(needs_file=False)}) == "research"
-
-
-def test_route_after_research_tools_vs_evidence():
+def test_route_after_research_tools_vs_solver():
     with_calls = SimpleNamespace(tool_calls=[{"name": "tavily_search"}])
     without = SimpleNamespace(tool_calls=[])
     assert route_after_research({"messages": [with_calls]}) == "tools"
-    assert route_after_research({"messages": [without]}) == "evidence"
+    assert route_after_research({"messages": [without]}) == "solver"
 
 
 def test_route_after_judge_pass():
@@ -38,23 +32,24 @@ def test_route_after_judge_pass():
 
 
 def test_route_after_judge_revise_within_budget():
-    # max_judge_retries default is 2; one attempt used -> still loop back.
-    state = {"verdict": JudgeVerdict(verdict="REVISE"), "attempts": 1}
-    assert route_after_judge(state) == "research"
+    state = {"verdict": JudgeVerdict(verdict="REVISE"), "attempts": 0}
+    with patch("gaia_agent.nodes.get_settings", return_value=SimpleNamespace(max_judge_retries=1)):
+        assert route_after_judge(state) == "research"
 
 
 def test_route_after_judge_revise_budget_exhausted():
-    state = {"verdict": JudgeVerdict(verdict="REVISE"), "attempts": 2}
-    assert route_after_judge(state) == "formatter"
+    state = {"verdict": JudgeVerdict(verdict="REVISE"), "attempts": 1}
+    with patch("gaia_agent.nodes.get_settings", return_value=SimpleNamespace(max_judge_retries=1)):
+        assert route_after_judge(state) == "formatter"
 
 
 def test_judge_increments_attempts_on_revise():
-    fake_llm = MagicMock()
-    fake_llm.with_structured_output.return_value.invoke.return_value = JudgeVerdict(
+    structured = MagicMock()
+    structured.invoke.return_value = JudgeVerdict(
         verdict="REVISE", feedback="wrong format", missing=["units"]
     )
-    with patch("gaia_agent.nodes.get_text_llm", return_value=fake_llm):
-        out = judge({"question": "q", "candidate": None, "evidence": None, "attempts": 0})
+    with patch("gaia_agent.nodes.get_structured_llm", return_value=structured):
+        out = judge({"question": "q", "candidate": None, "attempts": 0})
     assert out["attempts"] == 1
     assert out["verdict"].verdict == "REVISE"
     assert "messages" in out  # feedback injected back into the loop