116 lines
5.3 KiB
Python
116 lines
5.3 KiB
Python
"""
|
|
Graph assembly — wires together the supervisor and worker nodes.
|
|
|
|
This is the LangGraph equivalent of Loki's runtime agent execution engine
|
|
(src/supervisor/mod.rs + src/config/request_context.rs).
|
|
|
|
In Loki, the runtime:
|
|
1. Loads the agent config (config.yaml)
|
|
2. Compiles tools (tools.sh → binary)
|
|
3. Starts a chat loop: user → LLM → tool calls → LLM → ...
|
|
4. For orchestrators with can_spawn_agents: true, the supervisor module
|
|
manages child agent lifecycle (spawn, check, collect, cancel).
|
|
|
|
In LangGraph, all of this is declarative:
|
|
1. Define nodes (supervisor, explore, oracle, coder)
|
|
2. Define edges (workers always return to supervisor)
|
|
3. Compile the graph (with optional checkpointer for persistence)
|
|
4. Invoke with initial state
|
|
|
|
The graph topology:
|
|
|
|
┌─────────────────────────────────────────────┐
|
|
│ SUPERVISOR │
|
|
│ (classifies intent, routes to workers) │
|
|
└─────┬──────────┬──────────┬─────────────────┘
|
|
│ │ │
|
|
▼ ▼ ▼
|
|
┌────────┐ ┌────────┐ ┌────────┐
|
|
│EXPLORE │ │ ORACLE │ │ CODER │
|
|
│(search)│ │(advise)│ │(build) │
|
|
└───┬────┘ └───┬────┘ └───┬────┘
|
|
│ │ │
|
|
└──────────┼──────────┘
|
|
│
|
|
(back to supervisor)
|
|
|
|
Every worker returns to the supervisor. The supervisor decides what to do next:
|
|
route to another worker, or end the graph.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langgraph.graph import END, START, StateGraph
|
|
|
|
from sisyphus_langchain.agents.coder import create_coder_node
|
|
from sisyphus_langchain.agents.explore import create_explore_node
|
|
from sisyphus_langchain.agents.oracle import create_oracle_node
|
|
from sisyphus_langchain.agents.supervisor import create_supervisor_node
|
|
from sisyphus_langchain.state import SisyphusState
|
|
|
|
|
|
def build_graph(
|
|
*,
|
|
supervisor_model: str = "gpt-4o",
|
|
explore_model: str = "gpt-4o-mini",
|
|
oracle_model: str = "gpt-4o",
|
|
coder_model: str = "gpt-4o",
|
|
use_checkpointer: bool = True,
|
|
):
|
|
"""
|
|
Build and compile the Sisyphus LangGraph.
|
|
|
|
This is the main entry point for creating the agent system. It wires
|
|
together all nodes and edges, optionally adds a checkpointer for
|
|
persistence, and returns a compiled graph ready to invoke.
|
|
|
|
Args:
|
|
supervisor_model: Model for the routing supervisor.
|
|
explore_model: Model for the explore agent (can be cheaper).
|
|
oracle_model: Model for the oracle agent (should be strong).
|
|
coder_model: Model for the coder agent.
|
|
use_checkpointer: Whether to add MemorySaver for session persistence.
|
|
|
|
Returns:
|
|
A compiled LangGraph ready to .invoke() or .stream().
|
|
|
|
Model cost optimization (mirrors Loki's per-agent model config):
|
|
- supervisor: expensive (accurate routing is critical)
|
|
- explore: cheap (just searching, not reasoning deeply)
|
|
- oracle: expensive (deep reasoning, architecture advice)
|
|
- coder: expensive (writing correct code matters)
|
|
"""
|
|
# Create the graph builder with our typed state
|
|
builder = StateGraph(SisyphusState)
|
|
|
|
# ── Register nodes ─────────────────────────────────────────────────
|
|
# Each node is a function that takes state and returns state updates.
|
|
# This mirrors Loki's agent registration (agents are discovered by
|
|
# their config.yaml in the agents/ directory).
|
|
builder.add_node("supervisor", create_supervisor_node(supervisor_model))
|
|
builder.add_node("explore", create_explore_node(explore_model))
|
|
builder.add_node("oracle", create_oracle_node(oracle_model))
|
|
builder.add_node("coder", create_coder_node(coder_model))
|
|
|
|
# ── Define edges ───────────────────────────────────────────────────
|
|
# Entry point: every invocation starts at the supervisor
|
|
builder.add_edge(START, "supervisor")
|
|
|
|
# Workers always return to supervisor (the hub-and-spoke pattern).
|
|
# In Loki, this is implicit: agent__collect returns output to the parent,
|
|
# and the parent (sisyphus) decides what to do next.
|
|
builder.add_edge("explore", "supervisor")
|
|
builder.add_edge("oracle", "supervisor")
|
|
builder.add_edge("coder", "supervisor")
|
|
|
|
# The supervisor node itself uses Command(goto=...) to route,
|
|
# so we don't need add_conditional_edges — the Command API
|
|
# handles dynamic routing internally.
|
|
|
|
# ── Compile ────────────────────────────────────────────────────────
|
|
checkpointer = MemorySaver() if use_checkpointer else None
|
|
graph = builder.compile(checkpointer=checkpointer)
|
|
|
|
return graph
|