prebuilt: release 0.1.4 (#3977)

- This makes our checkpoint benchmarks more closely resemble the
behavior of our prod checkpointers
- Also found and fixed a bug w multiple subgraphs in same node
accidentally sharing checkpoints

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

@@ -7,7 +7,7 @@ from collections import defaultdict
 from collections.abc import AsyncIterator, Iterator, Sequence
 from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
 from types import TracebackType
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from langchain_core.runnables import RunnableConfig
 
@@ -70,6 +70,12 @@ class InMemorySaver(
         tuple[str, str, str],
         dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
     ]
+    blobs: dict[
+        tuple[
+            str, str, str, Union[str, int, float]
+        ],  # thread id, checkpoint ns, channel, version
+        tuple[str, bytes],
+    ]
 
     def __init__(
         self,
@@ -80,6 +86,7 @@ class InMemorySaver(
         super().__init__(serde=serde)
         self.storage = factory(lambda: defaultdict(dict))
         self.writes = factory(dict)
+        self.blobs = factory()
         self.stack = ExitStack()
         if factory is not defaultdict:
             self.stack.enter_context(self.storage)  # type: ignore[arg-type]
@@ -107,6 +114,18 @@ class InMemorySaver(
     ) -> Optional[bool]:
         return self.stack.__exit__(__exc_type, __exc_value, __traceback)
 
+    def _load_blobs(
+        self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
+    ) -> dict[str, Any]:
+        channel_values: dict[str, Any] = {}
+        for k, v in versions.items():
+            kk = (thread_id, checkpoint_ns, k, v)
+            if kk in self.blobs:
+                vv = self.blobs[kk]
+                if vv[0] != "empty":
+                    channel_values[k] = self.serde.loads_typed(vv)
+        return channel_values
+
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the in-memory storage.
 
@@ -121,8 +140,8 @@ class InMemorySaver(
         Returns:
             Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
         """
-        thread_id = config["configurable"]["thread_id"]
+        thread_id: str = config["configurable"]["thread_id"]
-        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
         if checkpoint_id := get_checkpoint_id(config):
             if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
                 checkpoint, metadata, parent_checkpoint_id = saved
@@ -140,10 +159,14 @@ class InMemorySaver(
                     )
                 else:
                     sends = []
+                checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
                 return CheckpointTuple(
                     config=config,
                     checkpoint={
-                        **self.serde.loads_typed(checkpoint),
+                        **checkpoint_,
+                        "channel_values": self._load_blobs(
+                            thread_id, checkpoint_ns, checkpoint_["channel_versions"]
+                        ),
                         "pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
                     },
                     metadata=self.serde.loads_typed(metadata),
@@ -180,6 +203,9 @@ class InMemorySaver(
                     )
                 else:
                     sends = []
+
+                checkpoint_ = self.serde.loads_typed(checkpoint)
+
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -189,7 +215,10 @@ class InMemorySaver(
                         }
                     },
                     checkpoint={
-                        **self.serde.loads_typed(checkpoint),
+                        **checkpoint_,
+                        "channel_values": self._load_blobs(
+                            thread_id, checkpoint_ns, checkpoint_["channel_versions"]
+                        ),
                         "pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
                     },
                     metadata=self.serde.loads_typed(metadata),
@@ -297,6 +326,8 @@ class InMemorySaver(
                     else:
                         sends = []
 
+                    checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
+
                     yield CheckpointTuple(
                         config={
                             "configurable": {
@@ -306,7 +337,12 @@ class InMemorySaver(
                             }
                         },
                         checkpoint={
-                            **self.serde.loads_typed(checkpoint),
+                            **checkpoint_,
+                            "channel_values": self._load_blobs(
+                                thread_id,
+                                checkpoint_ns,
+                                checkpoint_["channel_versions"],
+                            ),
                             "pending_sends": [
                                 self.serde.loads_typed(s[2]) for s in sends
                             ],
@@ -353,6 +389,11 @@ class InMemorySaver(
         c.pop("pending_sends")  # type: ignore[misc]
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
+        values: dict[str, Any] = c.pop("channel_values")  # type: ignore[misc]
+        for k, v in new_versions.items():
+            self.blobs[(thread_id, checkpoint_ns, k, v)] = (
+                self.serde.dumps_typed(values[k]) if k in values else ("empty", b"")
+            )
         self.storage[thread_id][checkpoint_ns].update(
             {
                 checkpoint["id"]: (

libs/checkpoint/tests/test_memory.py

@@ -68,7 +68,9 @@ class TestMemorySaver:
             },
             "metadata": {"run_id": "my_run_id"},
         }
-        self.memory_saver.put(config, self.chkpnt_2, self.metadata_2, {})
+        self.memory_saver.put(
+            config, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"]
+        )
         checkpoint = self.memory_saver.get_tuple(config)
         assert checkpoint is not None
         assert checkpoint.metadata == {
@@ -80,9 +82,24 @@ class TestMemorySaver:
     async def test_search(self) -> None:
         # set up test
         # save checkpoints
-        self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
+        self.memory_saver.put(
-        self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
+            self.config_1,
-        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
+            self.chkpnt_1,
+            self.metadata_1,
+            self.chkpnt_1["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_2,
+            self.chkpnt_2,
+            self.metadata_2,
+            self.chkpnt_2["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_3,
+            self.chkpnt_3,
+            self.metadata_3,
+            self.chkpnt_3["channel_versions"],
+        )
 
         # call method / assertions
         query_1 = {"source": "input"}  # search by 1 key
@@ -129,9 +146,24 @@ class TestMemorySaver:
     async def test_asearch(self) -> None:
         # set up test
         # save checkpoints
-        self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
+        self.memory_saver.put(
-        self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
+            self.config_1,
-        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
+            self.chkpnt_1,
+            self.metadata_1,
+            self.chkpnt_1["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_2,
+            self.chkpnt_2,
+            self.metadata_2,
+            self.chkpnt_2["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_3,
+            self.chkpnt_3,
+            self.metadata_3,
+            self.chkpnt_3["channel_versions"],
+        )
 
         # call method / assertions
         query_1 = {"source": "input"}  # search by 1 key

libs/cli/langgraph_cli/cli.py

@@ -575,7 +575,7 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
     default=False,
 )
 @click.option(
-    "--studio-url",
+    "--studio_url",
     type=str,
     default=None,
     help="URL of the LangGraph Studio instance to connect to. Defaults to https://smith.langchain.com",

libs/langgraph/langgraph/pregel/loop.py

@@ -269,13 +269,13 @@ class PregelLoop(LoopProtocol):
             self.checkpoint_config = patch_configurable(
                 self.config,
                 {
-                    CONFIG_KEY_CHECKPOINT_ID: config[CONF][CONFIG_KEY_CHECKPOINT_MAP][
+                    CONFIG_KEY_CHECKPOINT_ID: self.config[CONF][
-                        self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]
+                        CONFIG_KEY_CHECKPOINT_MAP
-                    ]
+                    ][self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]]
                 },
             )
         else:
-            self.checkpoint_config = config
+            self.checkpoint_config = self.config
         self.checkpoint_ns = (
             tuple(cast(str, self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP))
             if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)

libs/langgraph/poetry.lock

@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
 
 [[package]]
 name = "aiosqlite"
@@ -1348,7 +1348,7 @@ typing-extensions = ">=4.7"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "2.0.18"
+version = "2.0.21"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1366,7 +1366,7 @@ url = "../checkpoint"
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.16"
+version = "2.0.19"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1375,7 +1375,7 @@ files = []
 develop = true
 
 [package.dependencies]
-langgraph-checkpoint = "^2.0.15"
+langgraph-checkpoint = "^2.0.21"
 orjson = ">=3.10.1"
 psycopg = "^3.2.0"
 psycopg-pool = "^3.2.0"
@@ -1404,7 +1404,7 @@ url = "../checkpoint-sqlite"
 
 [[package]]
 name = "langgraph-prebuilt"
-version = "0.1.2"
+version = "0.1.4"
 description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -1422,7 +1422,7 @@ url = "../prebuilt"
 
 [[package]]
 name = "langgraph-sdk"
-version = "0.1.55"
+version = "0.1.58"
 description = "SDK for interacting with LangGraph API"
 optional = false
 python-versions = "^3.9.0,<4.0"

libs/langgraph/tests/memory_assert.py

@@ -1,4 +1,3 @@
-import asyncio
 import os
 import tempfile
 from collections import defaultdict
@@ -13,7 +12,6 @@ from langgraph.checkpoint.base import (
     CheckpointMetadata,
     CheckpointTuple,
     SerializerProtocol,
-    copy_checkpoint,
 )
 from langgraph.checkpoint.memory import InMemorySaver, PersistentDict
 
@@ -63,69 +61,14 @@ class MemorySaverAssertImmutable(InMemorySaver):
                     self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]]
                 )
                 == saved
-            )
+            ), config["configurable"]["checkpoint_ns"]
         self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
-            self.serde.dumps_typed(copy_checkpoint(checkpoint))
+            self.serde.dumps_typed(checkpoint)
         )
         # call super to write checkpoint
         return super().put(config, checkpoint, metadata, new_versions)
 
 
-class MemorySaverAssertCheckpointMetadata(InMemorySaver):
-    """This custom checkpointer is for verifying that a run's configurable
-    fields are merged with the previous checkpoint config for each step in
-    the run. This is the desired behavior. Because the checkpointer's (a)put()
-    method is called for each step, the implementation of this checkpointer
-    should produce a side effect that can be asserted.
-    """
-
-    def put(
-        self,
-        config: RunnableConfig,
-        checkpoint: Checkpoint,
-        metadata: CheckpointMetadata,
-        new_versions: ChannelVersions,
-    ) -> None:
-        """The implementation of put() merges config["configurable"] (a run's
-        configurable fields) with the metadata field. The state of the
-        checkpoint metadata can be asserted to confirm that the run's
-        configurable fields were merged with the previous checkpoint config.
-        """
-        configurable = config["configurable"].copy()
-
-        # remove checkpoint_id to make testing simpler
-        checkpoint_id = configurable.pop("checkpoint_id", None)
-        thread_id = config["configurable"]["thread_id"]
-        checkpoint_ns = config["configurable"]["checkpoint_ns"]
-        self.storage[thread_id][checkpoint_ns].update(
-            {
-                checkpoint["id"]: (
-                    self.serde.dumps_typed(checkpoint),
-                    # merge configurable fields and metadata
-                    self.serde.dumps_typed({**configurable, **metadata}),
-                    checkpoint_id,
-                )
-            }
-        )
-        return {
-            "configurable": {
-                "thread_id": config["configurable"]["thread_id"],
-                "checkpoint_id": checkpoint["id"],
-            }
-        }
-
-    async def aput(
-        self,
-        config: RunnableConfig,
-        checkpoint: Checkpoint,
-        metadata: CheckpointMetadata,
-        new_versions: ChannelVersions,
-    ) -> RunnableConfig:
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.put, config, checkpoint, metadata, new_versions
-        )
-
-
 class MemorySaverNoPending(InMemorySaver):
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         result = super().get_tuple(config)

libs/langgraph/tests/test_pregel.py

@@ -84,7 +84,6 @@ from tests.conftest import (
     REGULAR_CHECKPOINTERS_SYNC,
     SHOULD_CHECK_SNAPSHOTS,
 )
-from tests.memory_assert import MemorySaverAssertCheckpointMetadata
 from tests.messages import (
     _AnyIdAIMessage,
     _AnyIdAIMessageChunk,
@@ -4213,11 +4212,11 @@ def test_checkpoint_metadata() -> None:
     workflow.add_edge("tools", "agent")
 
     # graph w/o interrupt
-    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_1 = InMemorySaver()
     app = workflow.compile(checkpointer=checkpointer_1)
 
     # graph w/ interrupt
-    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_2 = InMemorySaver()
     app_w_interrupt = workflow.compile(
         checkpointer=checkpointer_2, interrupt_before=["tools"]
     )
@@ -4635,59 +4634,6 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None:
     assert app.get_graph(xray=True).draw_mermaid() == snapshot
 
 
-def test_subgraph_retries():
-    class State(TypedDict):
-        count: int
-
-    class ChildState(State):
-        some_list: Annotated[list, operator.add]
-
-    called_times = 0
-
-    class RandomError(ValueError):
-        """This will be retried on."""
-
-    def parent_node(state: State):
-        return {"count": state["count"] + 1}
-
-    def child_node_a(state: ChildState):
-        nonlocal called_times
-        # We want it to retry only on node_b
-        # NOT re-compute the whole graph.
-        assert not called_times
-        called_times += 1
-        return {"some_list": ["val"]}
-
-    def child_node_b(state: ChildState):
-        raise RandomError("First attempt fails")
-
-    child = StateGraph(ChildState)
-    child.add_node(child_node_a)
-    child.add_node(child_node_b)
-    child.add_edge("__start__", "child_node_a")
-    child.add_edge("child_node_a", "child_node_b")
-
-    parent = StateGraph(State)
-    parent.add_node("parent_node", parent_node)
-    parent.add_node(
-        "child_graph",
-        child.compile(),
-        retry=RetryPolicy(
-            max_attempts=3,
-            retry_on=(RandomError,),
-            backoff_factor=0.0001,
-            initial_interval=0.0001,
-        ),
-    )
-    parent.add_edge("parent_node", "child_graph")
-    parent.set_entry_point("parent_node")
-
-    checkpointer = InMemorySaver()
-    app = parent.compile(checkpointer=checkpointer)
-    with pytest.raises(RandomError):
-        app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}})
-
-
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 @pytest.mark.parametrize("store_name", ALL_STORES_SYNC)
 def test_store_injected(
@@ -6294,6 +6240,7 @@ def test_double_interrupt_subgraph(
     def invoke_sub_agent(state: AgentState):
         return subgraph.invoke(state)
 
+    thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
     parent_agent = (
         StateGraph(AgentState)
         .add_node("invoke_sub_agent", invoke_sub_agent)

libs/langgraph/tests/test_pregel_async.py

@@ -78,10 +78,7 @@ from tests.conftest import (
     awith_store,
 )
 from tests.fake_tracer import FakeTracer
-from tests.memory_assert import (
+from tests.memory_assert import MemorySaverNoPending
-    MemorySaverAssertCheckpointMetadata,
-    MemorySaverNoPending,
-)
 from tests.messages import (
     _AnyIdAIMessage,
     _AnyIdAIMessageChunk,
@@ -5770,11 +5767,11 @@ async def test_checkpoint_metadata() -> None:
     workflow.add_edge("tools", "agent")
 
     # graph w/o interrupt
-    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_1 = InMemorySaver()
     app = workflow.compile(checkpointer=checkpointer_1)
 
     # graph w/ interrupt
-    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_2 = InMemorySaver()
     app_w_interrupt = workflow.compile(
         checkpointer=checkpointer_2, interrupt_before=["tools"]
     )
@@ -7020,6 +7017,8 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None:
         def invoke_sub_agent(state: AgentState):
             return subgraph.invoke(state)
 
+        thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
         parent_agent = (
             StateGraph(AgentState)
             .add_node("invoke_sub_agent", invoke_sub_agent)

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -719,9 +719,7 @@ def create_react_agent(
     def generate_structured_response(
         state: StateSchema, config: RunnableConfig
     ) -> StateSchema:
-        # NOTE: we exclude the last message because there is enough information
+        messages = _get_state_value(state, "messages")
-        # for the LLM to generate the structured response
-        messages = _get_state_value(state, "messages")[:-1]
         structured_response_schema = response_format
         if isinstance(response_format, tuple):
             system_prompt, structured_response_schema = response_format
@@ -736,9 +734,7 @@ def create_react_agent(
     async def agenerate_structured_response(
         state: StateSchema, config: RunnableConfig
     ) -> StateSchema:
-        # NOTE: we exclude the last message because there is enough information
+        messages = _get_state_value(state, "messages")
-        # for the LLM to generate the structured response
-        messages = _get_state_value(state, "messages")[:-1]
         structured_response_schema = response_format
         if isinstance(response_format, tuple):
             system_prompt, structured_response_schema = response_format

libs/prebuilt/poetry.lock

@@ -435,7 +435,7 @@ typing-extensions = ">=4.7"
 
 [[package]]
 name = "langgraph"
-version = "0.3.0"
+version = "0.3.18"
 description = "Building stateful, multi-actor applications with LLMs"
 optional = false
 python-versions = ">=3.9.0,<4.0"
@@ -446,6 +446,7 @@ develop = true
 [package.dependencies]
 langchain-core = ">=0.1,<0.4"
 langgraph-checkpoint = "^2.0.10"
+langgraph-prebuilt = ">=0.1.1,<0.2"
 langgraph-sdk = "^0.1.42"
 
 [package.source]
@@ -454,7 +455,7 @@ url = "../langgraph"
 
 [[package]]
 name = "langgraph-checkpoint"
-version = "2.0.16"
+version = "2.0.21"
 description = "Library with base interfaces for LangGraph checkpoint savers."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -472,7 +473,7 @@ url = "../checkpoint"
 
 [[package]]
 name = "langgraph-checkpoint-postgres"
-version = "2.0.15"
+version = "2.0.19"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0,<4.0"
@@ -481,7 +482,7 @@ files = []
 develop = true
 
 [package.dependencies]
-langgraph-checkpoint = "^2.0.15"
+langgraph-checkpoint = "^2.0.21"
 orjson = ">=3.10.1"
 psycopg = "^3.2.0"
 psycopg-pool = "^3.2.0"
@@ -492,7 +493,7 @@ url = "../checkpoint-postgres"
 
 [[package]]
 name = "langgraph-checkpoint-sqlite"
-version = "2.0.5"
+version = "2.0.6"
 description = "Library with a SQLite implementation of LangGraph checkpoint saver."
 optional = false
 python-versions = "^3.9.0"

libs/prebuilt/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-prebuilt"
-version = "0.1.3"
+version = "0.1.4"
 description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
 authors = []
 license = "MIT"