"""
Supervisor agent that coordinates multiple specialized agents.
This agent uses a hierarchical approach - it delegates tasks to specialized
sub-agents (Engineering, Search, etc.) rather than having all tools directly.
"""
import logging
import os
import uuid
from typing import Annotated, Any, Literal, cast
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from langchain_core.runnables import RunnableConfig
from langgraph.errors import GraphRecursionError
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field
from typing_extensions import TypedDict
from config import config
from src.agents.arxiv_agent import ArXivAgent
from src.agents.base_agent import FORCED_STOP_MARKER, BaseAgent
from src.agents.cli_agent import CLIAgent
from src.agents.engineering_agent import EngineeringAgent
from src.agents.excalidraw_agent import ExcalidrawAgent
from src.agents.hpc_agent import HPCAgent
from src.agents.prusa_agent import PrusaAgent
from src.agents.rag_agent import RAGAgent
from src.agents.search_agent import SearchAgent
from src.checkpoint import get_checkpointer
from src.models.state import MessagesState
from src.utils.prompts import (
EXCALIDRAW_CAPABILITY_BLOCK,
SUPERVISOR_AGENT_SYSTEM_PROMPT,
SUPERVISOR_CAPABILITIES_PROMPT,
_is_eval_mode,
strip_suggested_prompts,
)
from src.utils.response_recovery import recover_local_response
logger = logging.getLogger(__name__)
# Maximum consecutive re-routings to the same agent before the supervisor
# forces FINISH. Smaller models often fail to emit FINISH and keep
# re-routing to the same agent (sometimes with spurious tool calls).
# The first delegation is always allowed; the counter tracks re-routings.
# Limit of 1: if an agent just completed with tool results, re-routing to
# the same agent is almost always a mistake (the LLM failed to recognise
# the step was done).
_MAX_CONSECUTIVE_SAME_AGENT_REROUTINGS = 1
[docs]
class RouteDecision(BaseModel):
"""Structured output for routing decision."""
agent: Literal[
"engineering_agent",
"hpc_agent",
"search_agent",
"rag_agent",
"arxiv_agent",
"prusa_agent",
"excalidraw_agent",
"cli_agent",
"supervisor_response",
"FINISH",
] = Field(description="The agent to route the task to based on the user's request")
reasoning: str = Field(
description="Brief explanation of why this agent was selected"
)
task_instruction: str = Field(
default="",
description=(
"Specific sub-task instruction for the selected agent. For multi-step "
"workflows, scope the instruction to ONLY the next step(s) that this "
"agent should perform, then stop. E.g. 'Generate the SLURM training "
"command. Do NOT download models or simulate designs — that will be "
"handled after HPC training completes.' Leave empty for single-step tasks."
),
)
more_steps_after: bool = Field(
default=False,
description=(
"True ONLY if the user's request requires a DIFFERENT agent after this "
"one completes (e.g. rag_agent finds parameters, then engineering_agent "
"optimises). False for single-agent tasks (the vast majority of requests)."
),
)
final_response: str = Field(
default="",
description=(
"User-facing reply, written in second person as a direct message to the "
"user. REQUIRED when agent='FINISH' AND no sub-agent has produced a "
"reply in this turn — e.g. the user is repeating a request that was "
"already completed earlier in the conversation, or no action is needed. "
"In that case, briefly tell the user what was already done (including "
"any file paths or key results) and offer next steps. Leave empty in "
"all other cases (when routing to a sub-agent, or when FINISH follows "
"a sub-agent that already replied)."
),
)
[docs]
class SupervisorState(TypedDict):
"""State for the supervisor agent."""
messages: Annotated[list, add_messages]
next: str
[docs]
class SupervisorAgent:
"""
Supervisor agent that coordinates specialized agents.
Uses a hierarchical delegation model where the supervisor routes
tasks to appropriate specialized agents based on the request.
"""
[docs]
def __init__(
self,
model_name: str | None = None,
temperature: float | None = None,
seed: int | None = None,
rag_read_only: bool = False,
):
"""Initialize the supervisor agent.
Args:
model_name: Name of the LLM model to use (defaults to config.llm_model)
temperature: Model temperature (defaults to config.llm_temperature)
seed: Random seed for model (defaults to config.llm_seed)
rag_read_only: If True, sub-agents only get read-only RAG tools (search, list).
"""
self.model_name = model_name or config.llm_model
self.temperature = (
temperature if temperature is not None else config.llm_temperature
)
self.seed = seed if seed is not None else config.llm_seed
# Build model kwargs
model_kwargs: dict[str, Any] = {"temperature": self.temperature}
if self.seed is not None:
model_kwargs["seed"] = self.seed
is_ollama = self.model_name.startswith("ollama:")
if is_ollama:
model_kwargs.update(
{
"num_ctx": config.ollama_num_ctx,
"reasoning": False,
"num_predict": 1024,
}
)
self.llm = init_chat_model(self.model_name, **model_kwargs)
# Create structured LLM for routing decisions
self.router_model_name = config.llm_router_model or self.model_name
router_is_ollama = self.router_model_name.startswith("ollama:")
routing_model = self.llm
if self.router_model_name != self.model_name or router_is_ollama:
routing_kwargs: dict[str, Any] = {
"temperature": config.llm_router_temperature
}
if self.seed is not None:
routing_kwargs["seed"] = self.seed
if router_is_ollama:
routing_kwargs.update(
{
"num_ctx": config.ollama_num_ctx,
"reasoning": False,
"num_predict": 256,
}
)
routing_model = init_chat_model(
self.router_model_name,
**routing_kwargs,
)
self.routing_llm = routing_model.with_structured_output(
RouteDecision, method="function_calling"
)
# Initialize specialized sub-agents
self.engineering_agent = EngineeringAgent(
model_name=self.model_name,
temperature=self.temperature,
seed=self.seed,
)
self.hpc_agent = HPCAgent(
model_name=self.model_name, temperature=self.temperature, seed=self.seed
)
skip_search = os.getenv("SKIP_SEARCH", "false").lower() == "true"
self.search_agent: SearchAgent | None = None
if not skip_search and config.search_available:
self.search_agent = SearchAgent(
model_name=self.model_name,
temperature=self.temperature,
seed=self.seed,
)
else:
logger.info(
"SearchAgent disabled (SKIP_SEARCH=%s, provider=%s, available=%s)",
skip_search,
config.search_provider,
config.search_available,
)
# Initialize RAG agent with MMORE
self.rag_agent = RAGAgent(
model_name=self.model_name,
temperature=self.temperature,
seed=self.seed,
rag_read_only=rag_read_only,
)
# Initialize ArXiv agent with MMORE
self.arxiv_agent = ArXivAgent(
model_name=self.model_name,
temperature=self.temperature,
seed=self.seed,
)
# PrusaAgent will check SKIP_MCP env var automatically
self.prusa_agent = PrusaAgent(
model_name=self.model_name, temperature=self.temperature, seed=self.seed
)
# None when the canvas server is disabled (SKIP_CANVAS, or SKIP_MCP as a
# master switch). _excalidraw_node answers for it in that case rather
# than routing into an agent with no tools -- a drawing agent that
# cannot draw would otherwise reply "I've drawn it!" beside a blank
# whiteboard.
self.excalidraw_agent = (
ExcalidrawAgent(
model_name=self.model_name,
temperature=self.temperature,
seed=self.seed,
)
if config.canvas_enabled
else None
)
# Enabled but unreachable is worse than disabled. The agent absorbs a
# connection failure and comes back with no tools, and everything
# downstream keys on "is there an agent" -- so it would be advertised in
# the capabilities answer, routed to, and then reply in prose about the
# diagram it had not drawn. Dropping it here keeps one invariant:
# advertised means it works.
if self.excalidraw_agent is not None and not self.excalidraw_agent.tools:
logger.warning(
"Canvas MCP server returned no tools (%s); the whiteboard agent "
"will not be offered",
config.canvas_mcp_url,
)
self.excalidraw_agent = None
# CLI confirmation is handled by supervisor-level interrupt, not internally
self.cli_agent = CLIAgent(
model_name=self.model_name,
require_confirmation=False,
temperature=self.temperature,
seed=self.seed,
)
# Loop detection state (reset per invoke)
self._last_routed_agent: str | None = None
self._consecutive_same_agent_count: int = 0
self._last_delegation_had_tools: bool = True
self._expects_followup: bool = False
# Build the supervisor graph
self.graph = self._build_graph()
def _strip_disabled_capabilities(self, prompt: str) -> str:
"""Remove capability blocks for agents that are not wired up.
Applied to BOTH consumers of AGENT_CAPABILITIES -- the routing prompt
and SUPERVISOR_CAPABILITIES_PROMPT. Stripping only the router would fix
which agents get *considered* while leaving the assistant telling users
it can draw on a whiteboard that is switched off, and would then answer
the follow-up with the bypass message.
Removal rather than a "do NOT route here" paragraph: AGENT_CAPABILITIES
is on the hot path of every benchmark run, so appending would add ~340
characters of context to evals that cannot use the agent anyway.
Removing keeps the disabled prompts byte-identical to before this agent
existed, so routing numbers stay comparable.
Keyed on the agent instance rather than re-reading ``config`` so the
prompt and the graph cannot disagree: ``canvas_enabled`` reads the
environment on every access, while the agent is built once in
``__init__``.
"""
if self._canvas_available() is None:
prompt = prompt.replace(EXCALIDRAW_CAPABILITY_BLOCK, "")
return prompt
def _canvas_available(self) -> ExcalidrawAgent | None:
"""The whiteboard agent, but only if it can actually draw.
Not just ``is None``. ``ExcalidrawAgent._create_tools`` returns ``[]``
when the canvas server is unreachable while ``canvas_enabled`` is still
true, so the agent exists and can draw nothing -- and a tool-less agent
does not merely fail, it answers with an *empty message*, because the
model's whole response is a tool call that BaseAgent's guard strips.
Both the prompt and the node ask this one question, so "advertised means
it works" survives a server that died rather than being switched off.
Benchmarks are unaffected: they set SKIP_MCP=true, which leaves the
agent None exactly as before.
Returns the agent rather than a bool so the caller gets a narrowed,
non-optional handle -- a bool would leave the node re-checking for None
to satisfy the type checker, which is how the two conditions drift
apart in the first place.
"""
if self.excalidraw_agent is None or not self.excalidraw_agent.tools:
return None
return self.excalidraw_agent
def _build_routing_prompt(self) -> str:
"""Get the routing system prompt for the supervisor.
Returns:
Routing prompt from centralized prompts file
"""
prompt = SUPERVISOR_AGENT_SYSTEM_PROMPT
if os.getenv("SKIP_ARXIV", "false").lower() == "true":
prompt += (
"\n\nIMPORTANT: ArXiv search is currently unavailable. "
"Do NOT route to arxiv_agent under any circumstances. "
"For any paper or document lookup, route to rag_agent first. "
"After rag_agent returns the information, route to engineering_agent "
"for any follow-up design or optimization tasks. "
"Keep task_instruction for rag_agent concise: just state which "
"parameter values to find and from which papers. Do NOT ask for "
"full citations, DOIs, authors, or extensive quotes — just the "
"numeric values needed for the next step."
)
prompt = self._strip_disabled_capabilities(prompt)
if os.getenv("SKIP_SEARCH", "false").lower() == "true":
prompt += (
"\n\nIMPORTANT: Web search is currently unavailable. "
"Do NOT route to search_agent under any circumstances. "
"Use rag_agent for document lookup and engineering_agent "
"for design tasks."
)
return prompt
def _supervisor_node(self, state: SupervisorState):
"""Supervisor decides which agent should act next using LLM-based routing.
After each agent completes, the supervisor re-evaluates the full message
history to decide whether to route to another agent or finish.
Prior [SUPERVISOR INSTRUCTION] messages are kept so the routing LLM
can see which sub-tasks were already delegated and completed.
"""
messages = [
{"role": "system", "content": self._build_routing_prompt()},
*state["messages"],
]
# Use structured output to get routing decision from LLM
route_decision = self._invoke_routing_llm(messages)
next_agent = route_decision.agent
# Store follow-up flag so _route_after_agent knows whether to
# loop back to the supervisor or end directly after delegation.
self._expects_followup = route_decision.more_steps_after
# Log the routing decision with reasoning
logger.info(
f"[SUPERVISOR ROUTING] Agent: '{next_agent}' "
f"| Reasoning: {route_decision.reasoning} "
f"| more_steps_after: {route_decision.more_steps_after}"
)
# ── Same-agent loop detection ────────────────────────────────
# If the LLM tries to re-route to the agent that just completed,
# it likely failed to recognise the step was done. Rather than
# forcing FINISH (which would skip remaining steps), re-invoke the
# routing LLM with a redirect hint so it picks a different agent.
if next_agent not in ("FINISH", "supervisor_response"):
if next_agent == self._last_routed_agent:
self._consecutive_same_agent_count += 1
else:
self._consecutive_same_agent_count = 0
if (
self._consecutive_same_agent_count
>= _MAX_CONSECUTIVE_SAME_AGENT_REROUTINGS
):
logger.warning(
f"[SUPERVISOR] Loop detected: '{next_agent}' re-routed "
f"{self._consecutive_same_agent_count} consecutive times. "
"Re-invoking routing with redirect hint."
)
# Re-invoke routing with a hint to pick a different agent
redirect_hint = HumanMessage(
content=(
f"[SYSTEM] You already delegated to '{next_agent}' and "
f"it completed its work (see results above). Do NOT "
f"route to '{next_agent}' again. Choose a DIFFERENT "
f"agent for the next step, or FINISH if all steps are done."
)
)
retry_messages = [*messages, redirect_hint]
retry_decision = self._invoke_routing_llm(retry_messages)
next_agent = retry_decision.agent
self._expects_followup = retry_decision.more_steps_after
logger.info(
f"[SUPERVISOR REDIRECT] Agent: '{next_agent}' "
f"| Reasoning: {retry_decision.reasoning} "
f"| more_steps_after: {retry_decision.more_steps_after}"
)
# If the LLM STILL picks the same agent, force FINISH
if next_agent == self._last_routed_agent:
logger.warning(
f"[SUPERVISOR] Redirect failed — still '{next_agent}'. "
"Forcing FINISH."
)
self._consecutive_same_agent_count = 0
self._last_routed_agent = None
return {
"next": "FINISH",
"messages": [
AIMessage(
content=(
"I've completed the task based on the "
"information gathered so far."
)
)
],
}
# Redirect succeeded — use the new decision going forward
route_decision = retry_decision
self._consecutive_same_agent_count = 0
self._last_routed_agent = next_agent
else:
# Routing to FINISH or supervisor_response — reset tracking
self._last_routed_agent = None
self._consecutive_same_agent_count = 0
# ── Inject scoped task instruction ───────────────────────────
# HumanMessage so the delegated agent's LLM treats it as a new
# user directive (not just a prior assistant turn). This
# prevents the agent from ignoring the scope and trying to
# handle the entire workflow with its own tools.
new_messages: list = []
if route_decision.task_instruction and next_agent not in (
"FINISH",
"supervisor_response",
):
new_messages.append(
HumanMessage(
content=(
f"[SUPERVISOR INSTRUCTION — follow this exactly] "
f"{route_decision.task_instruction}"
)
)
)
logger.info(
f"[SUPERVISOR] Task instruction for {next_agent}: {route_decision.task_instruction}"
)
if next_agent == "FINISH":
final_response = (route_decision.final_response or "").strip()
if final_response:
new_messages.append(
recover_local_response(
AIMessage(content=final_response),
[
SystemMessage(
content=self._strip_disabled_capabilities(
SUPERVISOR_CAPABILITIES_PROMPT
)
),
*state["messages"],
],
model_name=self.model_name,
answer_model=self.llm,
allow_tools=False,
)
)
else:
# Safety net: if the routing LLM omitted final_response and
# no sub-agent has replied this turn (last message is the
# user's), emit a minimal reply so the UI never renders blank.
last_message = state["messages"][-1] if state["messages"] else None
if last_message is None or isinstance(last_message, HumanMessage):
new_messages.append(
AIMessage(
content="Done. Let me know if you'd like anything else."
)
)
return {
"next": next_agent,
"messages": new_messages,
}
def _supervisor_response_node(self, state: SupervisorState):
"""Supervisor responds directly to informational questions."""
cap_prompt = self._strip_disabled_capabilities(SUPERVISOR_CAPABILITIES_PROMPT)
if _is_eval_mode():
cap_prompt = strip_suggested_prompts(cap_prompt)
messages = [
SystemMessage(content=cap_prompt),
*state["messages"],
]
response = self.llm.invoke(messages)
message = (
response
if isinstance(response, AIMessage)
else AIMessage(content=response.content)
)
message = recover_local_response(
message,
messages,
model_name=self.model_name,
answer_model=self.llm,
allow_tools=False,
)
return {"messages": [message], "next": "FINISH"}
def _invoke_routing_llm(self, messages: list) -> RouteDecision:
"""Return a bounded routing decision and never crash the user chat."""
for attempt in range(2):
try:
decision = self.routing_llm.invoke(messages)
except Exception as exc:
logger.warning(
"Routing attempt %d failed with %s",
attempt + 1,
type(exc).__name__,
)
continue
if isinstance(decision, RouteDecision):
return decision
logger.warning("Routing attempt %d returned no decision", attempt + 1)
logger.error(
"Routing unavailable; falling back to a direct supervisor response"
)
return RouteDecision(
agent="supervisor_response",
reasoning="The structured router was unavailable; respond safely.",
)
def _engineering_node(self, state: SupervisorState):
"""Delegate to engineering agent."""
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self.engineering_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"engineering_{uuid.uuid4().hex[:8]}"}},
)
return self._extract_agent_result(state, result)
def _invoke_with_salvage(
self, agent: BaseAgent, agent_state: MessagesState, label: str
) -> MessagesState:
"""Invoke a delegated agent, recovering its trace if it runs away.
The thread_id is held in a variable rather than inlined so the partial
conversation can be read back from the checkpointer when the agent
exhausts its recursion limit while self-correcting -- otherwise the
whole delegation is lost (see the _build_predict_output salvage in
evaluate_agent.py for the outer backstop).
Worth knowing how far that limit is: LangGraph's default is 10007
steps, not the 25 it is often assumed to be, so an agent looping on
tool calls burns thousands of LLM calls before raising. Nothing else
bounds these graphs.
"""
inner_config = {
"configurable": {"thread_id": f"{label}_{uuid.uuid4().hex[:8]}"}
}
try:
return agent.invoke(agent_state, inner_config)
except GraphRecursionError:
logger.warning(
"%s agent hit its recursion limit; salvaging partial trace.", label
)
try:
snapshot = agent.agent.get_state(inner_config)
salvaged = list(snapshot.values.get("messages", []))
except Exception: # salvage must never raise
logger.warning(
"Failed to salvage %s partial trace", label, exc_info=True
)
salvaged = list(agent_state["messages"])
# Recovering nothing new is the blank-reply failure one layer up:
# _extract_agent_result slices off the input, so an unchanged list
# surfaces zero messages and the user is told nothing at all.
if len(salvaged) <= len(agent_state["messages"]):
salvaged.append(
AIMessage(
content=(
"I got stuck in a loop working on that and had to "
"stop, and I could not recover what I had done. "
"Please try again, ideally with a narrower request."
)
)
)
return cast(MessagesState, {"messages": salvaged})
def _hpc_node(self, state: SupervisorState):
"""Delegate to HPC agent."""
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self._invoke_with_salvage(self.hpc_agent, agent_state, "hpc")
return self._extract_agent_result(state, result)
def _search_node(self, state: SupervisorState):
"""Delegate to search agent."""
if (
os.getenv("SKIP_SEARCH", "false").lower() == "true"
or self.search_agent is None
):
self._last_delegation_had_tools = True
return {
"messages": [
AIMessage(
content=(
"Web search is not available in this context. "
"Please use the document search tool (search_documents) instead."
)
)
],
"next": "",
}
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self.search_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"search_{uuid.uuid4().hex[:8]}"}},
)
return self._extract_agent_result(state, result)
def _rag_node(self, state: SupervisorState):
"""Delegate to RAG agent for document Q&A."""
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self.rag_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"rag_{uuid.uuid4().hex[:8]}"}},
)
return self._extract_agent_result(state, result)
def _arxiv_node(self, state: SupervisorState):
"""Delegate to ArXiv agent for paper search and analysis."""
if os.getenv("SKIP_ARXIV", "false").lower() == "true":
# System bypass, not an idle agent — supervisor should re-route
# (typically to rag_agent), so mark as "productive".
self._last_delegation_had_tools = True
return {
"messages": [
AIMessage(
content=(
"ArXiv search is not available in this context. "
"Please use the document search tool (search_documents) instead."
)
)
],
"next": "",
}
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self.arxiv_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"arxiv_{uuid.uuid4().hex[:8]}"}},
)
return self._extract_agent_result(state, result)
def _prusa_node(self, state: SupervisorState):
"""Delegate to Prusa agent.
A tool-less agent answers with nothing at all, so it is bypassed here.
Measured, not assumed: with the printer server unreachable the model's
entire response is a tool call, which BaseAgent's guard strips, leaving
an empty message. The user asks about printers and gets a blank reply.
This is the ``_search_node`` shape rather than the ``_excalidraw_node``
one -- the capability block stays in the routing prompt instead of being
stripped. Benchmarks run with SKIP_MCP=true, so stripping would rewrite
the routing prompt for every eval and move routing numbers that are
meant to stay comparable across runs.
"""
if not self.prusa_agent.tools:
# System bypass, not an idle agent -- mark productive so the
# supervisor re-routes rather than ending the turn.
self._last_delegation_had_tools = True
return {
"messages": [
AIMessage(
content=(
"3D printer control is not available in this "
"context — the Prusa MCP server is either switched "
"off or was unreachable at startup. Everything else "
"still works."
)
)
],
"next": "",
}
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self._invoke_with_salvage(self.prusa_agent, agent_state, "prusa")
return self._extract_agent_result(state, result)
def _excalidraw_node(
self, state: SupervisorState, config: RunnableConfig | None = None
):
"""Delegate to the whiteboard agent.
``config`` is injected by LangGraph for any node that declares a
parameter of that name, which is how the conversation's thread_id gets
here. Setting it on the agent before delegating is what makes the
canvas per-conversation: the tool loader re-reads it on every call, so
a chat switch between turns is picked up without rebuilding tools.
The name is forced by that injection, so it SHADOWS the module-level
``config`` singleton for the body of this method -- ``config.canvas_enabled``
and friends are unreachable here. Read them before delegating, or via
``self``.
"""
canvas_agent = self._canvas_available()
if canvas_agent is None:
# System bypass, not an idle agent -- mark productive so the
# supervisor re-routes (e.g. to answer in text) rather than ending.
self._last_delegation_had_tools = True
return {
"messages": [
AIMessage(
content=(
"The whiteboard is not available in this context. "
"Describe the drawing in words instead."
)
)
],
"next": "",
}
thread_id = ((config or {}).get("configurable") or {}).get("thread_id")
canvas_agent.canvas_id = thread_id
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = canvas_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"excalidraw_{uuid.uuid4().hex[:8]}"}},
)
return self._extract_agent_result(state, result)
def _cli_node(self, state: SupervisorState):
"""Delegate to CLI agent."""
logger.info("[SUPERVISOR] _cli_node invoked - delegating to CLI agent")
agent_state = cast(MessagesState, {"messages": state["messages"]})
result = self.cli_agent.invoke(
agent_state,
{"configurable": {"thread_id": f"cli_{uuid.uuid4().hex[:8]}"}},
)
logger.info("[SUPERVISOR] CLI agent returned result")
return self._extract_agent_result(state, result)
def _route_after_agent(self, state: SupervisorState) -> str: # noqa: ARG002
"""Decide whether to loop back to supervisor or end directly.
Uses two signals:
- ``_last_delegation_had_tools``: whether the agent made any tool calls.
- ``_expects_followup``: set from ``RouteDecision.more_steps_after``
by the supervisor node before delegation.
Returns to supervisor only when both conditions hold (agent was
productive AND the supervisor indicated more agents are needed).
Otherwise ends directly to avoid wasteful LLM re-evaluation.
"""
if not self._last_delegation_had_tools:
logger.info("[SUPERVISOR] Agent produced no tool calls — ending directly.")
return END
if self._expects_followup:
logger.info(
"[SUPERVISOR] Agent completed with tools — more steps expected, "
"returning to supervisor."
)
return "supervisor"
logger.info("[SUPERVISOR] Agent completed with tools — task done, ending.")
return END
def _build_graph(self):
"""Build the supervisor workflow graph with agent routing."""
# Build the graph
workflow = StateGraph(SupervisorState)
# Add nodes
workflow.add_node("supervisor", self._supervisor_node)
workflow.add_node("supervisor_response", self._supervisor_response_node)
workflow.add_node("engineering_agent", self._engineering_node)
workflow.add_node("hpc_agent", self._hpc_node)
workflow.add_node("search_agent", self._search_node)
workflow.add_node("rag_agent", self._rag_node)
workflow.add_node("arxiv_agent", self._arxiv_node)
workflow.add_node("prusa_agent", self._prusa_node)
workflow.add_node("excalidraw_agent", self._excalidraw_node)
workflow.add_node("cli_agent", self._cli_node)
# Add edges - start with supervisor
workflow.add_edge(START, "supervisor")
# Build conditional routing dictionary
routing_dict = {
"supervisor_response": "supervisor_response",
"engineering_agent": "engineering_agent",
"hpc_agent": "hpc_agent",
"search_agent": "search_agent",
"rag_agent": "rag_agent",
"arxiv_agent": "arxiv_agent",
"prusa_agent": "prusa_agent",
"excalidraw_agent": "excalidraw_agent",
"cli_agent": "cli_agent",
"FINISH": END,
}
# Conditional routing based on supervisor decision
workflow.add_conditional_edges(
"supervisor",
lambda state: state["next"],
routing_dict,
)
# Agents loop back to supervisor only if they used tools (productive).
# If an agent produced no tool calls it had nothing to do — go
# directly to END to avoid wasteful supervisor re-evaluation
# (common with smaller models that fail to emit FINISH).
agent_nodes = [
"engineering_agent",
"hpc_agent",
"search_agent",
"rag_agent",
"arxiv_agent",
"prusa_agent",
"excalidraw_agent",
"cli_agent",
]
after_agent_map = {"supervisor": "supervisor", END: END}
for agent_name in agent_nodes:
workflow.add_conditional_edges(
agent_name, self._route_after_agent, after_agent_map
)
# supervisor_response always goes directly to END
workflow.add_edge("supervisor_response", END)
# Compile with persistent checkpointer
checkpointer = get_checkpointer()
# Interrupt before cli_agent for user confirmation of commands
return workflow.compile(
checkpointer=checkpointer,
interrupt_before=["cli_agent"],
)
def _extract_agent_result(
self, state: SupervisorState, result: MessagesState
) -> dict:
"""Extract new messages from an agent result and track tool usage.
Shared by all agent delegation nodes to avoid duplicated logic.
Sets ``_last_delegation_had_tools`` for idle-loop detection in the
supervisor node.
"""
input_len = len(state["messages"])
new_messages = result["messages"][input_len:]
# A forced stop closes the call it abandoned with a ToolMessage (see
# BaseAgent._force_end_node). Counting it here would report the exact
# opposite of what happened -- no tool ran -- and re-route to an agent
# that had just given up.
self._last_delegation_had_tools = any(
isinstance(m, ToolMessage)
and not m.additional_kwargs.get(FORCED_STOP_MARKER)
for m in new_messages
)
return {"messages": new_messages, "next": ""}
[docs]
def invoke(self, state, config):
"""Invoke the supervisor agent.
Args:
state: Current conversation state with messages, or None to resume from interrupt
config: Configuration including thread_id
Returns:
Updated state with agent responses
"""
# Reset loop detection state for each top-level invocation
self._last_routed_agent = None
self._consecutive_same_agent_count = 0
self._last_delegation_had_tools = True
self._expects_followup = False
# If state is None, we're resuming from an interrupt - pass None to graph
if state is None:
result = self.graph.invoke(None, config)
return {"messages": result["messages"]}
# Convert MessagesState to SupervisorState if needed
if "next" not in state:
supervisor_state = {"messages": state["messages"], "next": ""}
else:
supervisor_state = state
result = self.graph.invoke(supervisor_state, config)
# Return in MessagesState format for compatibility
return {"messages": result["messages"]}
