Chapter 10

multiple schemas

NotebookPython 3 (ipykernel)15 cells

Multiple Schemas

Review

We just covered state schema and reducers.

Typically, all graph nodes communicate with a single schema.

Also, this single schema contains the graph's input and output keys / channels.

Goals

But, there are cases where we may want a bit more control over this:

Internal nodes may pass information that is not required in the graph's input / output.
We may also want to use different input / output schemas for the graph. The output might, for example, only contain a single relevant output key.

We'll discuss a few ways to customize graphs with multiple schemas.

In [1]python · cell 3

python

%%capture --no-stderr
%pip install --quiet -U langgraph

Private State

First, let's cover the case of passing private state between nodes.

This is useful for anything needed as part of the intermediate working logic of the graph, but not relevant for the overall graph input or output.

We'll define an OverallState and a PrivateState.

node_2 uses PrivateState as input, but writes out to OverallState.

In [2]python · cell 5

python

from typing_extensions import TypedDict
from IPython.display import Image, display
from langgraph.graph import StateGraph, START, END

class OverallState(TypedDict):
    foo: int

class PrivateState(TypedDict):
    baz: int

def node_1(state: OverallState) -> PrivateState:
    print("---Node 1---")
    return {"baz": state['foo'] + 1}

def node_2(state: PrivateState) -> OverallState:
    print("---Node 2---")
    return {"foo": state['baz'] + 1}

# Build graph
builder = StateGraph(OverallState)
builder.add_node("node_1", node_1)
builder.add_node("node_2", node_2)

# Logic
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
builder.add_edge("node_2", END)

# Add
graph = builder.compile()

# View
display(Image(graph.get_graph().draw_mermaid_png()))

Output

<IPython.core.display.Image object>

In [3]python · cell 6

python

graph.invoke({"foo" : 1})

Output

---Node 1---
---Node 2---

{'foo': 3}

baz is only included in PrivateState.

node_2 uses PrivateState as input, but writes out to OverallState.

So, we can see that baz is excluded from the graph output because it is not in OverallState.

Input / Output Schema

By default, StateGraph takes in a single schema and all nodes are expected to communicate with that schema.

However, it is also possible to define explicit input and output schemas for a graph.

In these cases, we often define an "internal" schema that contains all keys relevant to graph operations.

But we use specific input and output schemas to constrain the input and output.

First, let's just run the graph with a single schema.

In [4]python · cell 9

python

class OverallState(TypedDict):
    question: str
    answer: str
    notes: str

def thinking_node(state: OverallState):
    return {"answer": "bye", "notes": "... his name is Lance"}

def answer_node(state: OverallState):
    return {"answer": "bye Lance"}

graph = StateGraph(OverallState)
graph.add_node("answer_node", answer_node)
graph.add_node("thinking_node", thinking_node)
graph.add_edge(START, "thinking_node")
graph.add_edge("thinking_node", "answer_node")
graph.add_edge("answer_node", END)

graph = graph.compile()

# View
display(Image(graph.get_graph().draw_mermaid_png()))

Output

<IPython.core.display.Image object>

Notice that the output of invoke contains all keys in OverallState.

In [5]python · cell 11

python

graph.invoke({"question":"hi"})

Output

{'question': 'hi', 'answer': 'bye Lance', 'notes': '... his name is Lance'}

Now, let's use a specific input and output schema with our graph.

Here, input / output schemas perform filtering on what keys are permitted on the input and output of the graph.

In addition, we can use a type hint state: InputState to specify the input schema of each of our nodes.

This is important when the graph is using multiple schemas.

We use type hints below to, for example, show that the output of answer_node will be filtered to OutputState.

In [ ]python · cell 13

python

class InputState(TypedDict):
    question: str

class OutputState(TypedDict):
    answer: str

class OverallState(TypedDict):
    question: str
    answer: str
    notes: str

def thinking_node(state: InputState):
    return {"answer": "bye", "notes": "... his is name is Lance"}

def answer_node(state: OverallState) -> OutputState:
    return {"answer": "bye Lance"}

graph = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)
graph.add_node("answer_node", answer_node)
graph.add_node("thinking_node", thinking_node)
graph.add_edge(START, "thinking_node")
graph.add_edge("thinking_node", "answer_node")
graph.add_edge("answer_node", END)

graph = graph.compile()

# View
display(Image(graph.get_graph().draw_mermaid_png()))

graph.invoke({"question":"hi"})

Output

<IPython.core.display.Image object>

{'question': 'hi', 'answer': 'bye Lance', 'notes': '... his is name is Lance'}

We can see the output schema constrains the output to only the answer key.

Goals

But, there are cases where we may want a bit more control over this:

Internal nodes may pass information that is not required in the graph's input / output.

We may also want to use different input / output schemas for the graph. The output might, for example, only contain a single relevant output key.

We'll discuss a few ways to customize graphs with multiple schemas.

Private State

First, let's cover the case of passing private state between nodes.

This is useful for anything needed as part of the intermediate working logic of the graph, but not relevant for the overall graph input or output.

We'll define an OverallState and a PrivateState.

node_2 uses PrivateState as input, but writes out to OverallState.

from typing_extensions import TypedDict from IPython.display import Image, display from langgraph.graph import StateGraph, START, END class OverallState(TypedDict): foo: int class PrivateState(TypedDict): baz: int def node_1(state: OverallState) -> PrivateState: print("---Node 1---") return {"baz": state['foo'] + 1} def node_2(state: PrivateState) -> OverallState: print("---Node 2---") return {"foo": state['baz'] + 1} # Build graph builder = StateGraph(OverallState) builder.add_node("node_1", node_1) builder.add_node("node_2", node_2) # Logic builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") builder.add_edge("node_2", END) # Add graph = builder.compile() # View display(Image(graph.get_graph().draw_mermaid_png()))

Input / Output Schema

By default, StateGraph takes in a single schema and all nodes are expected to communicate with that schema.

In these cases, we often define an "internal" schema that contains all keys relevant to graph operations.

But we use specific input and output schemas to constrain the input and output.

First, let's just run the graph with a single schema.

class OverallState(TypedDict): question: str answer: str notes: str def thinking_node(state: OverallState): return {"answer": "bye", "notes": "... his name is Lance"} def answer_node(state: OverallState): return {"answer": "bye Lance"} graph = StateGraph(OverallState) graph.add_node("answer_node", answer_node) graph.add_node("thinking_node", thinking_node) graph.add_edge(START, "thinking_node") graph.add_edge("thinking_node", "answer_node") graph.add_edge("answer_node", END) graph = graph.compile() # View display(Image(graph.get_graph().draw_mermaid_png()))

class InputState(TypedDict): question: str class OutputState(TypedDict): answer: str class OverallState(TypedDict): question: str answer: str notes: str def thinking_node(state: InputState): return {"answer": "bye", "notes": "... his is name is Lance"} def answer_node(state: OverallState) -> OutputState: return {"answer": "bye Lance"} graph = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState) graph.add_node("answer_node", answer_node) graph.add_node("thinking_node", thinking_node) graph.add_edge(START, "thinking_node") graph.add_edge("thinking_node", "answer_node") graph.add_edge("answer_node", END) graph = graph.compile() # View display(Image(graph.get_graph().draw_mermaid_png())) graph.invoke({"question":"hi"})