graph: improve generics on StateGraph etc + move typing utils to private file (#4982)

This commit is contained in:
Sydney Runkle
2025-06-06 19:51:05 -04:00
committed by GitHub
17 changed files with 226 additions and 171 deletions
+2 -2
View File
@@ -89,7 +89,7 @@ def node_3(state: PrivateState) -> OutputState:
# Read from PrivateState, write to OutputState
return {"graph_output": state["bar"] + " Lance"}
builder = StateGraph(OverallState,input=InputState,output=OutputState)
builder = StateGraph(OverallState,input_schema=InputState,output_schema=OutputState)
builder.add_node("node_1", node_1)
builder.add_node("node_2", node_2)
builder.add_node("node_3", node_3)
@@ -107,7 +107,7 @@ There are two subtle and important points to note here:
1. We pass `state: InputState` as the input schema to `node_1`. But, we write out to `foo`, a channel in `OverallState`. How can we write out to a state channel that is not included in the input schema? This is because a node _can write to any state channel in the graph state._ The graph state is the union of the state channels defined at initialization, which includes `OverallState` and the filters `InputState` and `OutputState`.
2. We initialize the graph with `StateGraph(OverallState,input=InputState,output=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because _nodes can also declare additional state channels_ as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it.
2. We initialize the graph with `StateGraph(OverallState,input_schema=InputState,output_schema=OutputState)`. So, how can we write to `PrivateState` in `node_2`? How does the graph gain access to this schema if it was not passed in the `StateGraph` initialization? We can do this because _nodes can also declare additional state channels_ as long as the state schema definition exists. In this case, the `PrivateState` schema is defined, so we can add `bar` as a new state channel in the graph and write to it.
### Reducers
+1 -1
View File
@@ -94,7 +94,7 @@ def answer_node(state: InputState):
return {"answer": "bye", "question": state["question"]}
# Build the graph with explicit schemas
builder = StateGraph(OverallState, input=InputState, output=OutputState)
builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)
builder.add_node(answer_node)
builder.add_edge(START, "answer_node")
builder.add_edge("answer_node", END)
+3 -3
View File
@@ -439,7 +439,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"id": "6ec0eb77-874e-443e-8c73-93125b515106",
"metadata": {},
"outputs": [
@@ -478,7 +478,7 @@
"\n",
"\n",
"# Build the graph with input and output schemas specified\n",
"builder = StateGraph(OverallState, input=InputState, output=OutputState)\n",
"builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState)\n",
"builder.add_node(answer_node) # Add the answer node\n",
"builder.add_edge(START, \"answer_node\") # Define the starting edge\n",
"builder.add_edge(\"answer_node\", END) # Define the ending edge\n",
@@ -3430,7 +3430,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.9"
"version": "3.9.6"
}
},
"nbformat": 4,