How to define input/output schema for your graph¶
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. Often, in these cases, we define an "internal" schema that contains all keys relevant to graph operations. But, we use specific input and output schemas to filter what's permitted when invoking and what's returned. We use type hints below to, for example, show that the output of answer_node
will be filtered to OutputState
. In addition, we define each node's input schema (e.g., as state: OverallState
for answer_node
).
In this notebook we'll walk through an example of this. At a high level, in order to do this you simply have to pass in input=..., output=...
when defining the graph. See the conceptual docs here for more details.
Let's look at an example!
Setup¶
First, let's install the required packages
%%capture --no-stderr
%pip install -U langgraph
Define and use the graph¶
from langgraph.graph import StateGraph, START, END
from typing import TypedDict
class InputState(TypedDict):
question: str
class OutputState(TypedDict):
answer: str
class OverallState(InputState, OutputState):
pass
def answer_node(state: InputState):
return {"answer": "bye"}
builder = StateGraph(OverallState, input=InputState, output=OutputState)
builder.add_node(answer_node)
builder.add_edge(START, "answer_node")
builder.add_edge("answer_node", END)
graph = builder.compile()
graph.invoke({"question": "hi"})
{'answer': 'bye'}
Notice that the output of invoke only includes the output schema.