Multi-Agent & Deep Agents

# fake.py
import itertools

from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage


class ScriptedChat(GenericFakeChatModel):
    """Replays scripted AIMessages and tolerates bind_tools()."""

    def bind_tools(self, tools, **kwargs):
        return self


def scripted(*messages):
    """Replay these messages once, one per model call."""
    return ScriptedChat(messages=iter(messages))


def always(text):
    """Answer the same text on every call (for agents invoked more than once)."""
    return ScriptedChat(messages=itertools.cycle([AIMessage(content=text)]))


def call(name, _id, **args):
    """An AIMessage that asks for one tool call."""
    return AIMessage(content="", tool_calls=[{"name": name, "args": args, "id": _id, "type": "tool_call"}])

import operator
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END


class Shared(TypedDict):
    question: str
    notes: Annotated[list[str], operator.add]


def rewrite(state: Shared):
    return {"notes": [f"rewritten: {state['question'].lower()}"]}


def retrieve(state: Shared):
    return {"notes": ["retrieved: ann-indexes.adoc"]}


retrieval = (StateGraph(Shared)
             .add_node("rewrite", rewrite).add_node("retrieve", retrieve)
             .add_edge(START, "rewrite").add_edge("rewrite", "retrieve")
             .compile())


def answer(state: Shared):
    return {"notes": [f"answer from {len(state['notes'])} notes"]}


parent = (StateGraph(Shared)
          .add_node("retrieval", retrieval)          # the compiled graph itself is the node
          .add_node("answer", answer)
          .add_edge(START, "retrieval").add_edge("retrieval", "answer")
          .compile())
print(parent.invoke({"question": "What is HNSW?", "notes": []})["notes"])

['rewritten: what is hnsw?', 'retrieved: ann-indexes.adoc', 'answer from 2 notes']

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START


class Private(TypedDict):
    query: str
    hits: list[str]


def search(state: Private):
    return {"hits": [f"hit for {state['query']}"]}


inner = (StateGraph(Private).add_node("search", search)
         .add_edge(START, "search").compile())


class Outer(TypedDict):
    question: str
    context: list[str]


def call_inner(state: Outer):
    out = inner.invoke({"query": state["question"], "hits": []})   # parent state -> subgraph input
    return {"context": out["hits"]}                                # subgraph output -> parent state


outer = (StateGraph(Outer).add_node("call_inner", call_inner)
         .add_edge(START, "call_inner").compile())
print(outer.invoke({"question": "ef_search", "context": []}))

{'question': 'ef_search', 'context': ['hit for ef_search']}

from langchain.agents import create_agent
from langchain.tools import tool
from fake import scripted, call, AIMessage


@tool
def search_docs(query: str) -> str:
    """Search the documentation site and return matching sections."""
    return "ann-indexes.adoc :: HNSW :: ef_search trades recall for latency."   # pgvector in DocsAssistant


# The worker: its own model, tools and prompt, and its own clean context window
researcher = create_agent(
    scripted(call("search_docs", "r1", query="HNSW ef_search"),
             AIMessage(content="ef_search sets the candidate list size (ann-indexes.adoc).")),
    tools=[search_docs],
    system_prompt="You research the docs. Reply in at most three lines and name the page.",
)


@tool("research", description="Research a documentation question in depth and return a short, cited summary.")
def call_researcher(query: str) -> str:
    result = researcher.invoke({"messages": [{"role": "user", "content": query}]})
    return result["messages"][-1].content     # only the final message crosses back


supervisor = create_agent(
    scripted(call("research", "s1", query="What does HNSW ef_search do?"),
             AIMessage(content="ef_search is the candidate list size; higher means better recall, slower queries.")),
    tools=[call_researcher],
    system_prompt="You are DocsAssistant. Delegate documentation research to the research tool.",
)
out = supervisor.invoke({"messages": [{"role": "user", "content": "Explain ef_search"}]})
for m in out["messages"]:
    print(type(m).__name__, m.tool_calls[0]["name"] if getattr(m, "tool_calls", None) else m.content)

HumanMessage Explain ef_search
AIMessage research
ToolMessage ef_search sets the candidate list size (ann-indexes.adoc).
AIMessage ef_search is the candidate list size; higher means better recall, slower queries.

from typing import Callable

from langchain.agents import AgentState, create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, wrap_model_call
from langchain.messages import ToolMessage
from langchain.tools import ToolRuntime, tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command
from typing_extensions import NotRequired

from fake import ScriptedChat, call, AIMessage


class SupportState(AgentState):
    current_step: NotRequired[str]           # absent until the first handoff; the hook defaults it to "triage"
    component: NotRequired[str | None]


@tool
def record_component(component: str, runtime: ToolRuntime) -> Command:
    """Record which documentation component the question is about, then move to the answer step."""
    return Command(update={
        "messages": [ToolMessage(content=f"Component recorded: {component}", tool_call_id=runtime.tool_call_id)],
        "component": component,
        "current_step": "answer",
    })


@tool
def search_docs(query: str) -> str:
    """Search the documentation site and return matching sections."""
    return "ann-indexes.adoc :: HNSW :: ef_search trades recall for latency."


STEPS = {
    "triage": {"prompt": "Ask which documentation component the user means, then record it.", "tools": [record_component]},
    "answer": {"prompt": "Answer from search_docs for the {component} component.", "tools": [search_docs]},
}


@wrap_model_call
def apply_step(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:
    step = STEPS[request.state.get("current_step", "triage")]
    request = request.override(
        system_prompt=step["prompt"].format(component=request.state.get("component")),
        tools=step["tools"],
    )
    return handler(request)


seen = []   # what the model was offered on each call


class Recorder(ScriptedChat):
    def bind_tools(self, tools, **kwargs):
        seen.append([t.name for t in tools])
        return self


agent = create_agent(
    Recorder(messages=iter([
        call("record_component", "h1", component="vector-rag"),
        call("search_docs", "h2", query="ef_search"),
        AIMessage(content="ef_search trades recall for latency (ann-indexes.adoc)."),
    ])),
    tools=[record_component, search_docs],
    state_schema=SupportState,
    middleware=[apply_step],
    checkpointer=InMemorySaver(),
)
out = agent.invoke({"messages": [{"role": "user", "content": "What is ef_search?"}]},
                   {"configurable": {"thread_id": "alice-1"}})
print("tools offered per model call:", seen)
print("current_step:", out["current_step"], "| component:", out["component"])
print(out["messages"][-1].content)

tools offered per model call: [['record_component'], ['search_docs'], ['search_docs']]
current_step: answer | component: vector-rag
ef_search trades recall for latency (ann-indexes.adoc).

from typing import Literal

from langchain.agents import AgentState, create_agent
from langchain.messages import AIMessage, ToolMessage
from langchain.tools import ToolRuntime, tool
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command
from typing_extensions import NotRequired

from fake import call, scripted


class TeamState(AgentState):
    active_agent: NotRequired[str]


def handoff(target: str, description: str):
    @tool(f"transfer_to_{target}", description=description)
    def transfer(runtime: ToolRuntime) -> Command:
        last_ai = next(m for m in reversed(runtime.state["messages"]) if isinstance(m, AIMessage))
        ack = ToolMessage(content=f"Transferred to {target}", tool_call_id=runtime.tool_call_id)
        return Command(goto=target, update={"active_agent": target, "messages": [last_ai, ack]},
                       graph=Command.PARENT)
    return transfer


docs_agent = create_agent(
    scripted(call("transfer_to_bug_agent", "t1"),
             AIMessage(content="(never reached: the handoff ends this agent's run)")),
    tools=[handoff("bug_agent", "Hand the conversation to the bug-report agent when the user reports a docs bug.")],
    system_prompt="You answer questions about the documentation.",
)
bug_agent = create_agent(
    scripted(AIMessage(content="I can file that. Which page has the wrong ef_search default?")),
    tools=[handoff("docs_agent", "Hand back to the docs agent for ordinary questions.")],
    system_prompt="You collect details for a documentation bug report.",
)


def route_after(state: TeamState) -> Literal["docs_agent", "bug_agent", "__end__"]:
    last = state["messages"][-1]
    if isinstance(last, AIMessage) and not last.tool_calls:
        return END
    return state.get("active_agent") or "docs_agent"


builder = StateGraph(TeamState)
builder.add_node("docs_agent", lambda s: docs_agent.invoke(s))
builder.add_node("bug_agent", lambda s: bug_agent.invoke(s))
builder.add_conditional_edges(START, lambda s: s.get("active_agent") or "docs_agent", ["docs_agent", "bug_agent"])
for node in ("docs_agent", "bug_agent"):
    builder.add_conditional_edges(node, route_after, ["docs_agent", "bug_agent", END])
team = builder.compile()

out = team.invoke({"messages": [{"role": "user", "content": "The ef_search default in your page is wrong."}]})
for m in out["messages"]:
    print(type(m).__name__, [c["name"] for c in m.tool_calls] if getattr(m, "tool_calls", None) else m.content)
print("active_agent:", out["active_agent"])

HumanMessage The ef_search default in your page is wrong.
AIMessage ['transfer_to_bug_agent']
ToolMessage Transferred to bug_agent
AIMessage I can file that. Which page has the wrong ef_search default?
active_agent: bug_agent

from typing import Literal, TypedDict

from langgraph.graph import END, START, StateGraph
from langgraph.types import Command


class State(TypedDict):
    query: str
    answer: str


def route_query(state: State) -> Command[Literal["vector_rag", "spring_ai"]]:
    """Single-agent routing: Command(goto=...) names the next node and can update state at the same time."""
    target = "spring_ai" if "spring" in state["query"].lower() else "vector_rag"
    return Command(goto=target)


def vector_rag(state: State):
    return {"answer": "vector-rag agent handled it"}


def spring_ai(state: State):
    return {"answer": "spring-ai agent handled it"}


graph = (StateGraph(State)
         .add_node("route_query", route_query)
         .add_node("vector_rag", vector_rag).add_node("spring_ai", spring_ai)
         .add_edge(START, "route_query")
         .add_edge("vector_rag", END).add_edge("spring_ai", END)
         .compile())
print(graph.invoke({"query": "PgVectorStore in Spring AI"})["answer"])
print(graph.invoke({"query": "HNSW ef_search"})["answer"])

spring-ai agent handled it
vector-rag agent handled it

import operator
from typing import Annotated, TypedDict

from langchain.agents import create_agent
from langgraph.graph import END, START, StateGraph
from langgraph.types import Send

from fake import always


class Classification(TypedDict):
    agent: str
    query: str


class RouterState(TypedDict):
    query: str
    results: Annotated[list[str], operator.add]
    final_answer: str


class AgentInput(TypedDict):
    query: str


vector_rag = create_agent(always("pgvector: ef_search trades recall for latency."), tools=[],
                          system_prompt="You answer from the vector-rag documentation.")
spring_ai = create_agent(always("Spring AI: PgVectorStore exposes the same search request."), tools=[],
                         system_prompt="You answer from the Spring AI documentation.")


def classify(query: str) -> list[Classification]:
    """A real router asks the model (structured output); keywords keep this example deterministic."""
    picked = []
    if "pgvector" in query or "ef_search" in query:
        picked.append({"agent": "vector_rag", "query": query})
    if "spring" in query.lower():
        picked.append({"agent": "spring_ai", "query": query})
    return picked or [{"agent": "vector_rag", "query": query}]


def route_query(state: RouterState):
    return [Send(c["agent"], {"query": c["query"]}) for c in classify(state["query"])]


def run(agent, label):
    def node(state: AgentInput):
        answer = agent.invoke({"messages": [{"role": "user", "content": state["query"]}]})["messages"][-1].content
        return {"results": [f"[{label}] {answer}"]}
    return node


def synthesize(state: RouterState):
    return {"final_answer": "\n".join(state["results"])}


builder = StateGraph(RouterState)
builder.add_node("vector_rag", run(vector_rag, "vector-rag"))
builder.add_node("spring_ai", run(spring_ai, "spring-ai"))
builder.add_node("synthesize", synthesize)
builder.add_conditional_edges(START, route_query, ["vector_rag", "spring_ai"])
builder.add_edge("vector_rag", "synthesize")
builder.add_edge("spring_ai", "synthesize")
builder.add_edge("synthesize", END)
router = builder.compile()

print(router.invoke({"query": "Tune ef_search with pgvector, and show it in Spring AI", "results": []})["final_answer"])
print("---")
print(router.invoke({"query": "How do I tune ef_search?", "results": []})["final_answer"])

[vector-rag] pgvector: ef_search trades recall for latency.
[spring-ai] Spring AI: PgVectorStore exposes the same search request.
---
[vector-rag] pgvector: ef_search trades recall for latency.

from langchain.agents import create_agent
from langchain.tools import tool
from fake import scripted, call, AIMessage


@tool
def search_docs(query: str) -> str:
    """Search across the documentation verticals and return a combined, cited answer."""
    return router.invoke({"query": query, "results": []})["final_answer"]   # the router graph built above


chat = create_agent(
    scripted(call("search_docs", "q1", query="ef_search in pgvector"),
             AIMessage(content="Higher ef_search improves recall at the cost of latency.")),
    tools=[search_docs],
    system_prompt="You are DocsAssistant. Use search_docs to answer.",
)
print(chat.invoke({"messages": [{"role": "user", "content": "ef_search?"}]})["messages"][2].content)

[vector-rag] pgvector: ef_search trades recall for latency.

from langchain.agents import create_agent
from langchain.tools import tool
from fake import call, scripted, AIMessage

SKILLS = {
    "write_issue": "Draft a GitHub issue: title, steps to reproduce, expected vs actual, page path. Never open it yourself.",
    "review_page": "Review an .adoc page: check xrefs, code blocks that run, and links to official docs.",
}


@tool
def load_skill(skill_name: str) -> str:
    """Load a specialised skill prompt by name.

    Available skills:
    - write_issue: draft a documentation bug report
    - review_page: review an AsciiDoc page
    """
    return SKILLS.get(skill_name, f"Unknown skill {skill_name!r}; choose one of {sorted(SKILLS)}")


agent = create_agent(
    scripted(call("load_skill", "k1", skill_name="write_issue"),
             AIMessage(content="Title: ef_search default is wrong on ann-indexes.adoc ...")),
    tools=[load_skill],
    system_prompt="You are DocsAssistant. You have two skills, write_issue and review_page; call load_skill to read one before using it.",
)
out = agent.invoke({"messages": [{"role": "user", "content": "The ef_search default is wrong, please file it"}]})
for m in out["messages"]:
    print(type(m).__name__, m.tool_calls[0]["args"] if getattr(m, "tool_calls", None) else m.content)

HumanMessage The ef_search default is wrong, please file it
AIMessage {'skill_name': 'write_issue'}
ToolMessage Draft a GitHub issue: title, steps to reproduce, expected vs actual, page path. Never open it yourself.
AIMessage Title: ef_search default is wrong on ann-indexes.adoc ...