Fix get_state

langchain-ai · Nov 26, 2024 · 5fdba64 · 5fdba64
1 parent 2221dbe
commit 5fdba64
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
@@ -71,6 +71,7 @@ def _imp(func: types.FunctionType) -> Pregel:
             channels={START: EphemeralValue(Any), END: LastValue(Any, END)},
             input_channels=START,
             output_channels=END,
+            stream_channels=END,
             stream_mode="updates",
             checkpointer=checkpointer,
             store=store,

diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
@@ -2039,6 +2039,8 @@ def graph(state: dict) -> dict:
         {"graph": {"a": "0foobarbaz", "c": "something else"}},
     ]
 
+    assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"}
+
 
 @pytest.mark.repeat(20)
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)