Mastering LangGraph: Subgraphs

Mastering LangGraph: Subgraphs
How to Add and Use Subgraphs
Subgraphs allow the construction of complex systems with multiple components, each of which is a graph. A common use case for using subgraphs is building multi-agent systems.
The main issue when adding subgraphs is how the parent graph and the subgraph communicate, i.e., how they pass state to each other during graph execution. There are two scenarios:
  • The parent graph and subgraph share the architecture key. In this case, a node of the compiled subgraph can be added.
  • The parent graph and subgraph have different architectures. In this case, a node function that calls the subgraph must be added: this is useful when the parent graph and subgraph have different state patterns and need to convert states before or after calling the subgraph.
Below we demonstrate how to add subgraphs for each scenario.
Mastering LangGraph: Subgraphs
Using a Compiled Subgraph as a Node
A common scenario is that the parent graph and subgraph communicate via a shared state key (channel). For example, in multi-agent systems, agents typically communicate through a shared messages key. The steps to use this architecture are as follows:
  1. Define the subgraph workflow (subgraph_builder in the example below) and compile it.
  2. Pass the compiled subgraph to the .add_node method when defining the parent graph workflow.
from langgraph.graph import START, StateGraph
from typing import TypedDict

# Define subgraph
class SubgraphState(TypedDict):
    foo: str  # Note: this key is shared with parent graph state
    bar: str

def subgraph_node_1(state: SubgraphState):
    # Only setting bar here
    return {"bar": "bar"}

def subgraph_node_2(state: SubgraphState):
    # The value taken from state["foo"] is actually passed from parent graph
    return {"foo": state["foo"] + state["bar"]}

subgraph_builder = StateGraph(SubgraphState)
subgraph_builder.add_node(subgraph_node_1)
subgraph_builder.add_node(subgraph_node_2)
subgraph_builder.add_edge(START, "subgraph_node_1")
subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
subgraph = subgraph_builder.compile()

# Define parent graph
class ParentState(TypedDict):
    foo: str

def node_1(state: ParentState):
    return {"foo": "hi! " + state["foo"]}

builder = StateGraph(ParentState)
builder.add_node("node_1", node_1)
# Added a compiled subgraph as a node
builder.add_node("node_2", subgraph)
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
graph = builder.compile()

for chunk in graph.stream({"foo": "foo"}):
    print(chunk)
Output:
{'node_1': {'foo': 'hi! foo'}}
{'node_2': {'foo': 'hi! foobar'}}
If you want to see the output of the subgraph, you can specify subgraphs=True during streaming.
for chunk in graph.stream({"foo": "foo"}, subgraphs=True):
    print(chunk)
((), {'node_1': {'foo': 'hi! foo'}})(('node_2:e58e5673-a661-ebb0-70d4-e298a7fc28b7',), {'subgraph_node_1': {'bar': 'bar'}})(('node_2:e58e5673-a661-ebb0-70d4-e298a7fc28b7',), {'subgraph_node_2': {'foo': 'hi! foobar'}})((), {'node_2': {'foo': 'hi! foobar'}})
Executing Subgraphs via Node Functions
For more complex systems, it may be necessary to define subgraphs that have completely different patterns from the parent graph (without shared keys).
For example, in a multi-agent RAG system, a search agent may only need to track queries and the retrieved documents. If your application is of this kind, you need to define a node function that calls the subgraph.
This function needs to convert the input (parent) state to the subgraph state before calling the subgraph, and convert the result back to the parent state before returning the state updates from the node.
Below we demonstrate how to modify the original example to call the subgraph from within the node.
⚠️ You cannot call multiple subgraphs within the same node.
from typing import TypedDict
from langgraph.constants import START
from langgraph.graph import StateGraph

class SubgraphState(TypedDict):
    bar: str
    baz: str

def subgraph_node_1(state: SubgraphState):
    return {"baz": "baz"}

def subgraph_node_2(state: SubgraphState):
    return {"bar": state["bar"] + state["baz"]}

subgraph_builder = StateGraph(SubgraphState)
subgraph_builder.add_node(subgraph_node_1)
subgraph_builder.add_node(subgraph_node_2)
subgraph_builder.add_edge(START, "subgraph_node_1")
subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
subgraph = subgraph_builder.compile()

# Define parent graph
class ParentState(TypedDict):
    foo: str

def node_1(state: ParentState):
    return {"foo": "hi! " + state["foo"]}

def node_2(state: ParentState):
    # Convert state to subgraph state
    response = subgraph.invoke({"bar": state["foo"]})
    # Convert response back to parent state
    return {"foo": response["bar"]}

builder = StateGraph(ParentState)
builder.add_node("node_1", node_1)
# Note that instead of using the compiled subgraph we are using `node_2` function that is calling the subgraph
builder.add_node("node_2", node_2)
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
graph = builder.compile()

for chunk in graph.stream({"foo": "foo"}, subgraphs=True):
    print(chunk)
((), {'node_1': {'foo': 'hi! foo'}})(('node_2:6525365b-44ed-3ff4-9824-a9dd11ff9386',), {'subgraph_node_1': {'baz': 'baz'}})(('node_2:6525365b-44ed-3ff4-9824-a9dd11ff9386',), {'subgraph_node_2': {'bar': 'hi! foobaz'}})((), {'node_2': {'foo': 'hi! foobaz'}})

Leave a Comment