Add missing input and output types for state graph

langgraph/graph/state.py

@@ -78,13 +78,13 @@ class StateGraph(Graph):
         node_inboxes = {
             # we take any value written to channel because all writers
             # write the entire state as of that step, which is equal for all
-            f"{key}:inbox": AnyValue(Any)
-            for key in self.nodes
+            f"{key}:inbox": AnyValue(self.schema)
+            for key in list(self.nodes) + [START]
         }
         node_outboxes = {
             # we clear outbox channels after each step
             key: EphemeralValue(Any)
-            for key in self.nodes
+            for key in list(self.nodes) + [START]
         }
 
         for key in self.nodes:
@@ -114,7 +114,12 @@ class StateGraph(Graph):
         return CompiledGraph(
             graph=self,
             nodes=nodes,
-            channels={**self.channels, **node_inboxes, **node_outboxes},
+            channels={
+                **self.channels,
+                **node_inboxes,
+                **node_outboxes,
+                END: LastValue(self.schema),
+            },
             input=f"{START}:inbox",
             output=END,
             hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,

tests/test_pregel.py

@@ -1074,7 +1074,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None:
     ]
 
 
-def test_prebuilt_tool_chat() -> None:
+def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
     from langchain.chat_models.fake import FakeMessagesListChatModel
     from langchain_community.tools import tool
     from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
@@ -1137,6 +1137,11 @@ def test_prebuilt_tool_chat() -> None:
         tools,
     )
 
+    assert app.get_input_schema().schema_json() == snapshot
+    assert app.get_output_schema().schema_json() == snapshot
+    assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
+    assert app.get_graph().draw_ascii() == snapshot
+
     assert app.invoke(
         {"messages": [HumanMessage(content="what is weather in sf")]}
     ) == {
@@ -1318,7 +1323,7 @@ def test_prebuilt_tool_chat() -> None:
     ]
 
 
-def test_prebuilt_chat() -> None:
+def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None:
     from langchain.chat_models.fake import FakeMessagesListChatModel
     from langchain_community.tools import tool
     from langchain_core.messages import AIMessage, FunctionMessage, HumanMessage
@@ -1361,6 +1366,11 @@ def test_prebuilt_chat() -> None:
         tools,
     )
 
+    assert app.get_input_schema().schema_json() == snapshot
+    assert app.get_output_schema().schema_json() == snapshot
+    assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
+    assert app.get_graph().draw_ascii() == snapshot
+
     assert app.invoke(
         {"messages": [HumanMessage(content="what is weather in sf")]}
     ) == {
@@ -1463,7 +1473,7 @@ def test_prebuilt_chat() -> None:
     ]
 
 
-def test_message_graph() -> None:
+def test_message_graph(snapshot: SnapshotAssertion) -> None:
     from langchain.chat_models.fake import FakeMessagesListChatModel
     from langchain_community.tools import tool
     from langchain_core.agents import AgentAction
@@ -1574,6 +1584,11 @@ def test_message_graph() -> None:
     # meaning you can use it as you would any other runnable
     app = workflow.compile()
 
+    assert app.get_input_schema().schema_json() == snapshot
+    assert app.get_output_schema().schema_json() == snapshot
+    assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
+    assert app.get_graph().draw_ascii() == snapshot
+
     assert app.invoke(HumanMessage(content="what is weather in sf")) == [
         HumanMessage(content="what is weather in sf"),
         AIMessage(