Mastering LangGraph - State Management - 02

Mastering LangGraph – State Management – 02

by
Mastering LangGraph - State Management - 02
How to Define the Input and Output Schema of a Graph
By default, the StateGraph operates using a single schema, and all nodes should communicate using that schema. However, different input and output schemas can also be defined for the graph.
When different schemas are specified, the internal schema will still be used for communication between nodes. The input schema ensures that the provided input matches the expected structure, while the output schema filters internal data to return only relevant information based on the defined output schema.
In this example, we will see how to define different input and output schemas.
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

# Define the input schema
class InputState(TypedDict):
    question: str

# Define the output schema
class OutputState(TypedDict):
    answer: str

# Define the overall schema, combining input and output
class OverallState(InputState, OutputState):
    pass

# Define the node that processes input and generates an answer
def answer_node(state: InputState):
    # Example answer and an additional key
    return {"answer": "bye", "question": state["question"]}

# Build the graph with specified input and output schemas
builder = StateGraph(OverallState, input=InputState, output=OutputState)
builder.add_node(answer_node)  # Add answer node
builder.add_edge(START, "answer_node")  # Define starting edge
builder.add_edge("answer_node", END)  # Define ending edge
graph = builder.compile()  # Compile the graph

# Call the graph with input and print the result
print(graph.invoke({"question": "hi"}))
The key point is this
StateGraph(OverallState, input=InputState, output=OutputState)

Leave a Comment