langgraph: update get_state to handle nested subgraph state

libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py

@@ -255,9 +255,6 @@ def _search_where(
         if config:
             wheres.append("thread_id = %s ")
             param_values.append(config["configurable"]["thread_id"])
-            checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-            wheres.append("checkpoint_ns = %s")
-            param_values.append(checkpoint_ns)
 
         # construct predicate for metadata filter
         if filter:

libs/checkpoint-postgres/tests/test_async.py

@@ -87,29 +87,15 @@ async def test_asearch(self):
             search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = [
                 c
                 async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
             ]
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = [
-                c
-                async for c in saver.alist(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            ]
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-postgres/tests/test_sync.py

@@ -88,27 +88,14 @@ def test_search(self):
             search_results_4 = list(saver.list(None, filter=query_4))
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = list(
                 saver.list({"configurable": {"thread_id": "thread-2"}})
             )
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = list(
-                saver.list(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            )
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py

@@ -70,9 +70,6 @@ def search_where(
     if config is not None:
         wheres.append("thread_id = ?")
         param_values.append(config["configurable"]["thread_id"])
-        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
-        wheres.append("checkpoint_ns = ?")
-        param_values.append(checkpoint_ns)
 
     # construct predicate for metadata filter
     if filter:

libs/checkpoint-sqlite/tests/test_aiosqlite.py

@@ -84,29 +84,15 @@ async def test_asearch(self):
             search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = [
                 c
                 async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
             ]
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = [
-                c
-                async for c in saver.alist(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            ]
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params

libs/checkpoint-sqlite/tests/test_sqlite.py

@@ -87,28 +87,15 @@ def test_search(self):
             search_results_4 = list(saver.list(None, filter=query_4))
             assert len(search_results_4) == 0
 
-            # search by config (defaults to root graph checkpoints)
+            # search by config (defaults to checkpoints across all namespaces)
             search_results_5 = list(
                 saver.list({"configurable": {"thread_id": "thread-2"}})
             )
-            assert len(search_results_5) == 1
-            assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-            # search by config and checkpoint_ns
-            search_results_6 = list(
-                saver.list(
-                    {
-                        "configurable": {
-                            "thread_id": "thread-2",
-                            "checkpoint_ns": "inner",
-                        }
-                    }
-                )
-            )
-            assert len(search_results_6) == 1
-            assert (
-                search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
-            )
+            assert len(search_results_5) == 2
+            assert {
+                search_results_5[0].config["configurable"]["checkpoint_ns"],
+                search_results_5[1].config["configurable"]["checkpoint_ns"],
+            } == {"", "inner"}
 
             # TODO: test before and limit params
 

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -177,62 +177,66 @@ def list(
             Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
         """
         thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
-        checkpoint_ns = (
-            config["configurable"].get("checkpoint_ns", "") if config else ""
-        )
         for thread_id in thread_ids:
-            for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted(
-                self.storage[thread_id][checkpoint_ns].items(),
-                key=lambda x: x[0],
-                reverse=True,
-            ):
-                # filter by checkpoint ID
-                if (
-                    before
-                    and (before_checkpoint_id := get_checkpoint_id(before))
-                    and checkpoint_id >= before_checkpoint_id
-                ):
-                    continue
-
-                # filter by metadata
-                metadata = self.serde.loads_typed(metadata_b)
-                if filter and not all(
-                    query_value == metadata[query_key]
-                    for query_key, query_value in filter.items()
+            for checkpoint_ns in self.storage[thread_id].keys():
+                for checkpoint_id, (
+                    checkpoint,
+                    metadata_b,
+                    parent_checkpoint_id,
+                ) in sorted(
+                    self.storage[thread_id][checkpoint_ns].items(),
+                    key=lambda x: x[0],
+                    reverse=True,
                 ):
-                    continue
-
-                # limit search results
-                if limit is not None and limit <= 0:
-                    break
-                elif limit is not None:
-                    limit -= 1
-
-                writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
-
-                yield CheckpointTuple(
-                    config={
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": checkpoint_id,
+                    # filter by checkpoint ID
+                    if (
+                        before
+                        and (before_checkpoint_id := get_checkpoint_id(before))
+                        and checkpoint_id >= before_checkpoint_id
+                    ):
+                        continue
+
+                    # filter by metadata
+                    metadata = self.serde.loads_typed(metadata_b)
+                    if filter and not all(
+                        query_value == metadata.get(query_key)
+                        for query_key, query_value in filter.items()
+                    ):
+                        continue
+
+                    # limit search results
+                    if limit is not None and limit <= 0:
+                        break
+                    elif limit is not None:
+                        limit -= 1
+
+                    writes = self.writes[
+                        (thread_id, checkpoint_ns, checkpoint_id)
+                    ].values()
+
+                    yield CheckpointTuple(
+                        config={
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": checkpoint_id,
+                            }
+                        },
+                        checkpoint=self.serde.loads_typed(checkpoint),
+                        metadata=metadata,
+                        parent_config={
+                            "configurable": {
+                                "thread_id": thread_id,
+                                "checkpoint_ns": checkpoint_ns,
+                                "checkpoint_id": parent_checkpoint_id,
+                            }
                         }
-                    },
-                    checkpoint=self.serde.loads_typed(checkpoint),
-                    metadata=metadata,
-                    parent_config={
-                        "configurable": {
-                            "thread_id": thread_id,
-                            "checkpoint_ns": checkpoint_ns,
-                            "checkpoint_id": parent_checkpoint_id,
-                        }
-                    }
-                    if parent_checkpoint_id
-                    else None,
-                    pending_writes=[
-                        (id, c, self.serde.loads_typed(v)) for id, c, v in writes
-                    ],
-                )
+                        if parent_checkpoint_id
+                        else None,
+                        pending_writes=[
+                            (id, c, self.serde.loads_typed(v)) for id, c, v in writes
+                        ],
+                    )
 
     def put(
         self,
@@ -338,7 +342,14 @@ async def alist(
         """
         loop = asyncio.get_running_loop()
         iter = await loop.run_in_executor(
-            None, partial(self.list, before=before, limit=limit, filter=filter), config
+            None,
+            partial(
+                self.list,
+                before=before,
+                limit=limit,
+                filter=filter,
+            ),
+            config,
         )
         while True:
             # handling StopIteration exception inside coroutine won't work

libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py

@@ -104,7 +104,7 @@ def _default(self, obj):
             return self._encode_constructor_args(obj.__class__, args=[obj.value])
         elif isinstance(obj, SendProtocol):
             return self._encode_constructor_args(
-                obj.__class__, kwargs={"node": obj.node, "arg": obj.arg}
+                obj.__class__, kwargs={"node": obj.node, "arg": obj.arg, "id": obj.id}
             )
         elif isinstance(obj, (bytes, bytearray)):
             return self._encode_constructor_args(

libs/checkpoint/langgraph/checkpoint/serde/types.py

@@ -57,6 +57,7 @@ class SendProtocol(Protocol):
     # Mirrors langgraph.constants.Send
     node: str
     arg: Any
+    id: str
 
     def __hash__(self) -> int:
         ...

libs/checkpoint/tests/test_memory.py

@@ -82,26 +82,20 @@ async def test_search(self):
         assert search_results_2[0].metadata == self.metadata_2
 
         search_results_3 = list(self.memory_saver.list(None, filter=query_3))
-        assert len(search_results_3) == 2
+        assert len(search_results_3) == 3
 
         search_results_4 = list(self.memory_saver.list(None, filter=query_4))
         assert len(search_results_4) == 0
 
-        # search by config (defaults to root graph checkpoints)
+        # search by config (defaults to checkpoints across all namespaces)
         search_results_5 = list(
             self.memory_saver.list({"configurable": {"thread_id": "thread-2"}})
         )
-        assert len(search_results_5) == 1
-        assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
-
-        # search by config and checkpoint_ns
-        search_results_6 = list(
-            self.memory_saver.list(
-                {"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}}
-            )
-        )
-        assert len(search_results_6) == 1
-        assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
+        assert len(search_results_5) == 2
+        assert {
+            search_results_5[0].config["configurable"]["checkpoint_ns"],
+            search_results_5[1].config["configurable"]["checkpoint_ns"],
+        } == {"", "inner"}
 
         # TODO: test before and limit params
 
@@ -110,6 +104,7 @@ async def test_asearch(self):
         # save checkpoints
         self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
         self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
+        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
 
         # call method / assertions
         query_1: CheckpointMetadata = {"source": "input"}  # search by 1 key
@@ -135,7 +130,7 @@ async def test_asearch(self):
         search_results_3 = [
             c async for c in self.memory_saver.alist(None, filter=query_3)
         ]
-        assert len(search_results_3) == 2
+        assert len(search_results_3) == 3
 
         search_results_4 = [
             c async for c in self.memory_saver.alist(None, filter=query_4)