From 094a9bd40aa8beaa20be8c0bec7bdfb0dc26b184 Mon Sep 17 00:00:00 2001 From: kgridou <32600911+kgridou@users.noreply.github.com> Date: Tue, 9 Jun 2026 22:46:16 +0200 Subject: [PATCH] Supervisor graph: route requests to knowledge vs action specialists Replace the single linear graph (guardrails -> retrieve -> generate) with a supervisor that classifies each request and dispatches to one of two specialists: guardrails -> supervisor --(knowledge)--> retrieve -> answer -> END --(action)----> act_agent --(sensitive?)--> approval -> act -> END --(otherwise)----------------------> END - supervisor_node uses the fast model to classify into "knowledge" | "action" (one-word reply, forgiving parse that defaults to the safe knowledge path). - answer_node (knowledge) binds no tools, so it is structurally incapable of triggering an action -- it can only answer or refuse. - act_agent_node (action) binds tools; sensitive ones still flow through the existing human-in-the-loop approval gate, non-sensitive execute inline. - AgentState gains a "route" field. Streaming: scope token forwarding to the answer/act_agent nodes so the supervisor's one-word routing label never leaks into the user's stream. Tests: add tests/fakes.route_model helper; existing graph tests now script the classifier (patch get_fast_model); new tests/test_supervisor.py pins both routes, the safe default, and that the action path never retrieves. Docs: README Phase 2 + roadmap, OBSERVABILITY trace example updated. Co-Authored-By: Claude Opus 4.8 --- README.md | 10 +++--- agentforge/agents/graph.py | 29 +++++++++++---- agentforge/agents/nodes.py | 72 +++++++++++++++++++++++++++++++------ agentforge/agents/state.py | 4 +++ agentforge/api/main.py | 8 ++++- docs/OBSERVABILITY.md | 2 +- tests/fakes.py | 9 +++++ tests/test_graph_flow.py | 6 ++-- tests/test_hitl_approval.py | 4 +-- tests/test_pii_flow.py | 3 +- tests/test_supervisor.py | 58 ++++++++++++++++++++++++++++++ 11 files changed, 177 insertions(+), 28 deletions(-) create mode 100644 tests/test_supervisor.py diff --git a/README.md b/README.md index b4b55b0..5aaa5c5 100644 --- a/README.md +++ b/README.md @@ -140,7 +140,7 @@ Ingestion → chunking → embeddings → pgvector → grounded answers with cit *Skills: RAG, Prompt Engineering, LLM.* ### Phase 2 — Agentic orchestration (LangGraph) -Single stateful graph: planner → retrieval/tool nodes → **human-in-the-loop approval** before sensitive actions. Durable execution that survives restarts. Guardrails middleware (PII redaction, tool scoping). +Stateful graph with a **supervisor** that routes each request to a knowledge or action specialist → **human-in-the-loop approval** before sensitive actions. Durable execution that survives restarts. Guardrails middleware (PII redaction, tool scoping). *Skills: LangGraph, LangChain, agents. **The core of the profile — invest the most time here.*** ### Phase 3 — Observability & Evals @@ -161,11 +161,11 @@ Chat UI + admin views: live traces, eval scores, approval queue. The full-stack *Skills: full-stack differentiator.* ### Beyond v1.0 — community milestones -- Langfuse adapter (fully self-hosted observability) -- Multi-agent / supervisor graphs -- Qdrant + alternative vector-store adapters +- ✅ Langfuse adapter (fully self-hosted observability) +- ✅ Multi-agent / supervisor graphs +- ✅ Qdrant + alternative vector-store adapters +- ✅ Additional reference examples (insurance, support) - Multi-tenancy -- Additional reference examples (insurance, support) --- diff --git a/agentforge/agents/graph.py b/agentforge/agents/graph.py index 824899f..d37f03c 100644 --- a/agentforge/agents/graph.py +++ b/agentforge/agents/graph.py @@ -17,32 +17,49 @@ from langgraph.graph import END, START, StateGraph from agentforge.agents.nodes import ( + act_agent_node, act_node, + answer_node, approval_node, - generate_node, guardrails_node, retrieve_node, + supervisor_node, ) from agentforge.agents.state import AgentState -def _route_after_generate(state: AgentState) -> str: +def _route_from_supervisor(state: AgentState) -> str: + """Dispatch to the specialist the supervisor selected (default: knowledge).""" + return "action" if state.get("route") == "action" else "knowledge" + + +def _route_after_action(state: AgentState) -> str: return "approval" if state.get("proposed_action") else "end" def build_graph() -> StateGraph: graph = StateGraph(AgentState) graph.add_node("guardrails", guardrails_node) + graph.add_node("supervisor", supervisor_node) graph.add_node("retrieve", retrieve_node) - graph.add_node("generate", generate_node) + graph.add_node("answer", answer_node) + graph.add_node("act_agent", act_agent_node) graph.add_node("approval", approval_node) graph.add_node("act", act_node) graph.add_edge(START, "guardrails") - graph.add_edge("guardrails", "retrieve") - graph.add_edge("retrieve", "generate") + graph.add_edge("guardrails", "supervisor") + graph.add_conditional_edges( + "supervisor", + _route_from_supervisor, + {"knowledge": "retrieve", "action": "act_agent"}, + ) + # Knowledge path: ground, answer, done. + graph.add_edge("retrieve", "answer") + graph.add_edge("answer", END) + # Action path: propose a tool, gate sensitive ones through human approval. graph.add_conditional_edges( - "generate", _route_after_generate, {"approval": "approval", "end": END} + "act_agent", _route_after_action, {"approval": "approval", "end": END} ) graph.add_edge("approval", "act") graph.add_edge("act", END) diff --git a/agentforge/agents/nodes.py b/agentforge/agents/nodes.py index 35f41b9..6f44c51 100644 --- a/agentforge/agents/nodes.py +++ b/agentforge/agents/nodes.py @@ -1,13 +1,18 @@ """Graph nodes — the units of work the agent executes. -The flow encodes the platform's safety posture: +A supervisor routes each request to one of two specialists, which keeps the +safety posture legible: the knowledge path can never trigger an action (no +tools are bound), and every action goes through the same human gate. - guardrails → retrieve → generate ──(sensitive tool?)──→ approval → act - └──(otherwise)─────────────────────→ end + guardrails → supervisor ─(knowledge)→ retrieve → answer ──────────────→ end + └(action)──→ act_agent ─(sensitive tool?)→ approval → act → end + └(otherwise)──────────────────────→ end -- ``guardrails`` strips PII before anything reaches the model or a trace. +- ``guardrails`` strips PII before anything reaches a model or a trace. +- ``supervisor`` classifies the request as "knowledge" or "action". - ``retrieve`` grounds the answer; an empty result drives a refusal. -- ``generate`` answers with citations, or proposes a sensitive tool call. +- ``answer`` replies with citations — no tools bound, so it can't act. +- ``act_agent`` proposes a tool call for the requested operation. - ``approval`` pauses the graph (durable) until a human approves/rejects. - ``act`` runs the tool only after sign-off. """ @@ -21,7 +26,7 @@ from agentforge.agents.tools import TOOLS, execute_tool from agentforge.config import get_settings from agentforge.guardrails import redact_pii, requires_approval -from agentforge.llm import get_chat_model +from agentforge.llm import get_chat_model, get_fast_model from agentforge.rag import retrieve # Default persona (banking-compliance reference example). Override per-domain via @@ -35,6 +40,15 @@ - For sensitive actions (escalating a case, filing a SAR), call the appropriate \ tool. A human will review before it runs — do not claim the action is done.""" +# Router persona for the supervisor. Kept deliberately narrow: one word out. +ROUTER_PROMPT = """You route a user request to exactly one specialist. Reply with \ +a single lowercase word and nothing else: +- "knowledge" — the user is asking a question to be answered from policy or \ +reference documents. +- "action" — the user is asking you to perform a sensitive operation (for \ +example escalating a case or filing a report). +When unsure, answer "knowledge".""" + def _latest_question(state: AgentState) -> str: """Prefer an explicit ``question``; otherwise the last human message.""" @@ -52,6 +66,22 @@ def guardrails_node(state: AgentState) -> dict: return {"question": question, "redacted_question": redacted, "pii_found": found} +def supervisor_node(state: AgentState) -> dict: + """Classify the request and route it to a specialist. + + Uses the cheap/fast model (a one-word classification, not reasoning). The + parse is forgiving and biased to "knowledge" — the safe default, since the + knowledge path has no tools and so can never trigger an action. + """ + model = get_fast_model() + response = model.invoke( + [SystemMessage(ROUTER_PROMPT), HumanMessage(state["redacted_question"])] + ) + text = (response.content or "").strip().lower() + route = "action" if "action" in text else "knowledge" + return {"route": route} + + def retrieve_node(state: AgentState) -> dict: result = retrieve(state["redacted_question"]) return { @@ -61,15 +91,20 @@ def retrieve_node(state: AgentState) -> dict: } -def generate_node(state: AgentState) -> dict: - model = get_chat_model().bind_tools(TOOLS) +def answer_node(state: AgentState) -> dict: + """Knowledge specialist: a grounded, cited answer with no tools bound. + + Because no tools are bound, this path is structurally incapable of + proposing or running an action — the model can only answer or refuse. + """ + model = get_chat_model() context = state.get("context") or "(no relevant context found)" grounding_note = ( "Relevant context is below.\n\n" + context if state.get("grounded") - else "No relevant context was found. Refuse and mark the question out of scope, " - "unless the user is requesting a sensitive action." + else "No relevant context was found. Say you don't know and that the " + "question is out of scope." ) messages = [ SystemMessage(get_settings().system_prompt or SYSTEM_PROMPT), @@ -77,6 +112,23 @@ def generate_node(state: AgentState) -> dict: HumanMessage(state["redacted_question"]), ] response = model.invoke(messages) + return {"messages": [response], "answer": response.content, "proposed_action": None} + + +def act_agent_node(state: AgentState) -> dict: + """Action specialist: bind tools and propose the requested operation.""" + model = get_chat_model().bind_tools(TOOLS) + messages = [ + SystemMessage(get_settings().system_prompt or SYSTEM_PROMPT), + SystemMessage( + "The user is requesting an action. If it requires a sensitive tool " + "(escalating a case, filing a SAR), call the appropriate tool — a human " + "will review before it runs, so do not claim it is done. If no tool " + "applies, answer directly." + ), + HumanMessage(state["redacted_question"]), + ] + response = model.invoke(messages) tool_calls = getattr(response, "tool_calls", []) or [] if tool_calls: diff --git a/agentforge/agents/state.py b/agentforge/agents/state.py index d6c79fe..8f414d1 100644 --- a/agentforge/agents/state.py +++ b/agentforge/agents/state.py @@ -20,6 +20,10 @@ class AgentState(TypedDict, total=False): redacted_question: str pii_found: list[str] + # Supervisor routing decision: "knowledge" (answer from docs) | "action" + # (perform a sensitive operation). Set by ``supervisor_node``. + route: str + # Retrieval output. context: str citations: list[dict[str, Any]] diff --git a/agentforge/api/main.py b/agentforge/api/main.py index 2672243..2e1e509 100644 --- a/agentforge/api/main.py +++ b/agentforge/api/main.py @@ -205,13 +205,19 @@ async def chat_stream(req: ChatRequest) -> EventSourceResponse: graph = get_compiled_graph() config = _run_config(thread_id) + # Only stream tokens from the answer-producing specialists — never the + # supervisor, whose model emits a one-word routing label, not user output. + _streaming_nodes = {"answer", "act_agent"} + async def event_generator(): yield {"event": "thread", "data": thread_id} - # Stream LLM tokens as they are produced by the generate node. + # Stream LLM tokens as the answer/act_agent specialists produce them. async for event in graph.astream_events( {"question": req.message}, config=config, version="v2" ): if event["event"] == "on_chat_model_stream": + if event["metadata"].get("langgraph_node") not in _streaming_nodes: + continue chunk = event["data"]["chunk"] if getattr(chunk, "content", ""): yield {"event": "token", "data": chunk.content} diff --git a/docs/OBSERVABILITY.md b/docs/OBSERVABILITY.md index 19522cb..38fc2e1 100644 --- a/docs/OBSERVABILITY.md +++ b/docs/OBSERVABILITY.md @@ -61,7 +61,7 @@ docker compose -f docker-compose.yml -f docker-compose.langfuse.yml up --build ``` 3. Restart the api: `docker compose ... up -d api`. 4. Ask a question, then refresh the Langfuse project — a trace per graph node - (guardrails → retrieve → generate → …) appears. + (guardrails → supervisor → retrieve → answer → …) appears. `LANGFUSE_HOST` is already set to `http://langfuse:3000` by the overlay. Set `LANGFUSE_NEXTAUTH_SECRET` and `LANGFUSE_SALT` in `.env` for anything beyond a diff --git a/tests/fakes.py b/tests/fakes.py index ffb49d6..1ea09d6 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -34,6 +34,15 @@ def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResu return ChatResult(generations=[ChatGeneration(message=message)]) +def route_model(route: str) -> FakeChatModel: + """A fake fast-model standing in for the supervisor's classifier. + + Returns ``route`` ("knowledge" or "action") as the one-word reply the + supervisor parses, so tests can pin which specialist a request reaches. + """ + return FakeChatModel(responses=[AIMessage(content=route)]) + + def grounded_retrieval(text: str = "Enhanced due diligence applies at $10,000.") -> RetrievalResult: """A retrieval result with one relevant chunk (drives a grounded answer).""" return RetrievalResult( diff --git a/tests/test_graph_flow.py b/tests/test_graph_flow.py index c0004ed..62330ca 100644 --- a/tests/test_graph_flow.py +++ b/tests/test_graph_flow.py @@ -5,7 +5,7 @@ import agentforge.agents.nodes as nodes from langchain_core.messages import AIMessage -from tests.fakes import FakeChatModel, empty_retrieval, grounded_retrieval +from tests.fakes import FakeChatModel, empty_retrieval, grounded_retrieval, route_model def _config(thread_id: str) -> dict: @@ -13,6 +13,7 @@ def _config(thread_id: str) -> dict: def test_grounded_answer_includes_citations(fresh_graph, monkeypatch): + monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("knowledge")) monkeypatch.setattr(nodes, "retrieve", lambda q: grounded_retrieval()) monkeypatch.setattr( nodes, @@ -31,6 +32,7 @@ def test_grounded_answer_includes_citations(fresh_graph, monkeypatch): def test_out_of_scope_question_is_refused(fresh_graph, monkeypatch): + monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("knowledge")) monkeypatch.setattr(nodes, "retrieve", lambda q: empty_retrieval()) monkeypatch.setattr( nodes, @@ -50,8 +52,8 @@ def test_out_of_scope_question_is_refused(fresh_graph, monkeypatch): def test_non_sensitive_tool_executes_without_approval(fresh_graph, monkeypatch): # Treat everything as non-sensitive for this test: the auto-exec path runs. + monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("action")) monkeypatch.setattr(nodes, "requires_approval", lambda name: False) - monkeypatch.setattr(nodes, "retrieve", lambda q: grounded_retrieval()) tool_call = { "name": "escalate_case", "args": {"customer_ref": "C-1", "reason": "unusual activity"}, diff --git a/tests/test_hitl_approval.py b/tests/test_hitl_approval.py index a4dd428..9658ca7 100644 --- a/tests/test_hitl_approval.py +++ b/tests/test_hitl_approval.py @@ -6,7 +6,7 @@ from langchain_core.messages import AIMessage from langgraph.types import Command -from tests.fakes import FakeChatModel, grounded_retrieval +from tests.fakes import FakeChatModel, route_model _SAR_CALL = { "name": "file_sar", @@ -17,7 +17,7 @@ def _setup(monkeypatch): - monkeypatch.setattr(nodes, "retrieve", lambda q: grounded_retrieval()) + monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("action")) monkeypatch.setattr( nodes, "get_chat_model", diff --git a/tests/test_pii_flow.py b/tests/test_pii_flow.py index 2388ac2..b6ddb38 100644 --- a/tests/test_pii_flow.py +++ b/tests/test_pii_flow.py @@ -5,7 +5,7 @@ import agentforge.agents.nodes as nodes from langchain_core.messages import AIMessage -from tests.fakes import FakeChatModel, empty_retrieval +from tests.fakes import FakeChatModel, empty_retrieval, route_model def test_pii_redacted_before_retrieval(fresh_graph, monkeypatch): @@ -15,6 +15,7 @@ def spy_retrieve(query: str): seen["query"] = query return empty_retrieval() + monkeypatch.setattr(nodes, "get_fast_model", lambda: route_model("knowledge")) monkeypatch.setattr(nodes, "retrieve", spy_retrieve) monkeypatch.setattr( nodes, "get_chat_model", lambda: FakeChatModel(responses=[AIMessage("ok")]) diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py new file mode 100644 index 0000000..bcc7ad4 --- /dev/null +++ b/tests/test_supervisor.py @@ -0,0 +1,58 @@ +"""Supervisor routing: knowledge vs action, with a safe default. + +These pin the structural guarantees of the split graph — the action path +never retrieves, the knowledge path never binds tools — and the forgiving +parse that biases an unrecognized classification back to "knowledge". +""" + +from __future__ import annotations + +import agentforge.agents.nodes as nodes +from langchain_core.messages import AIMessage + +from tests.fakes import FakeChatModel, empty_retrieval + + +def _config(thread_id: str) -> dict: + return {"configurable": {"thread_id": thread_id}} + + +def _no_retrieve(query: str): + raise AssertionError("knowledge path was taken on an action request") + + +def test_action_request_routes_to_act_agent_without_retrieval(fresh_graph, monkeypatch): + monkeypatch.setattr(nodes, "get_fast_model", lambda: FakeChatModel(responses=[AIMessage("action")])) + # Retrieval must never run on the action path. + monkeypatch.setattr(nodes, "retrieve", _no_retrieve) + sar_call = { + "name": "file_sar", + "args": {"customer_ref": "C-2", "summary": "possible structuring"}, + "id": "call_x", + "type": "tool_call", + } + monkeypatch.setattr( + nodes, + "get_chat_model", + lambda: FakeChatModel(responses=[AIMessage(content="", tool_calls=[sar_call])]), + ) + + result = fresh_graph.invoke({"question": "File a SAR for C-2"}, config=_config("sup-act")) + + assert result["route"] == "action" + assert result["proposed_action"]["name"] == "file_sar" + assert "__interrupt__" in result # paused for human approval + + +def test_unrecognized_classification_defaults_to_knowledge(fresh_graph, monkeypatch): + # A garbage one-word reply must fall back to the safe (tool-less) path. + monkeypatch.setattr(nodes, "get_fast_model", lambda: FakeChatModel(responses=[AIMessage("???")])) + monkeypatch.setattr(nodes, "retrieve", lambda q: empty_retrieval()) + monkeypatch.setattr( + nodes, "get_chat_model", lambda: FakeChatModel(responses=[AIMessage("Out of scope.")]) + ) + + result = fresh_graph.invoke({"question": "anything"}, config=_config("sup-default")) + + assert result["route"] == "knowledge" + assert result["proposed_action"] is None