langgraph-checkpoint-postgres 1.0.4__tar.gz → 1.0.6__tar.gz

{langgraph_checkpoint_postgres-1.0.4 → langgraph_checkpoint_postgres-1.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.4
+Version: 1.0.6
 Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
 Home-page: https://www.github.com/langchain-ai/langgraph
 License: MIT
@@ -11,7 +11,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: langgraph-checkpoint (>=1.0.1,<2.0.0)
+Requires-Dist: langgraph-checkpoint (>=1.0.8,<2.0.0)
 Requires-Dist: orjson (>=3.10.1)
 Requires-Dist: psycopg (>=3.0.0,<4.0.0)
 Requires-Dist: psycopg-pool (>=3.0.0,<4.0.0)

{langgraph_checkpoint_postgres-1.0.4 → langgraph_checkpoint_postgres-1.0.6}/langgraph/checkpoint/postgres/__init__.py

@@ -154,10 +154,11 @@ class PostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": self._load_blobs(value["channel_values"]),
-                    },
+                    self._load_checkpoint(
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
@@ -232,10 +233,11 @@ class PostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": self._load_blobs(value["channel_values"]),
-                    },
+                    self._load_checkpoint(
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {

{langgraph_checkpoint_postgres-1.0.4 → langgraph_checkpoint_postgres-1.0.6}/langgraph/checkpoint/postgres/aio.py

@@ -1,6 +1,6 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Optional, Union
+from typing import Any, AsyncIterator, Iterator, List, Optional, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
@@ -51,6 +51,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
 
     @classmethod
     @asynccontextmanager
@@ -64,7 +65,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
             pipeline (bool): whether to use AsyncPipeline
 
         Returns:
-            PostgresSaver: A new PostgresSaver instance.
+            AsyncPostgresSaver: A new AsyncPostgresSaver instance.
         """
         async with await AsyncConnection.connect(
             conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
@@ -137,12 +138,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": await asyncio.to_thread(
-                            self._load_blobs, value["channel_values"]
-                        ),
-                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
@@ -196,12 +197,12 @@ class AsyncPostgresSaver(BasePostgresSaver):
                             "checkpoint_id": value["checkpoint_id"],
                         }
                     },
-                    {
-                        **self._load_checkpoint(value["checkpoint"]),
-                        "channel_values": await asyncio.to_thread(
-                            self._load_blobs, value["channel_values"]
-                        ),
-                    },
+                    await asyncio.to_thread(
+                        self._load_checkpoint,
+                        value["checkpoint"],
+                        value["channel_values"],
+                        value["pending_sends"],
+                    ),
                     self._load_metadata(value["metadata"]),
                     {
                         "configurable": {
@@ -329,3 +330,96 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     binary=True, row_factory=dict_row
                 ) as cur:
                     yield cur
+
+    def list(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> Iterator[CheckpointTuple]:
+        """List checkpoints from the database.
+
+        This method retrieves a list of checkpoint tuples from the Postgres database based
+        on the provided config. The checkpoints are ordered by checkpoint ID in descending order (newest first).
+
+        Args:
+            config (Optional[RunnableConfig]): Base configuration for filtering checkpoints.
+            filter (Optional[Dict[str, Any]]): Additional filtering criteria for metadata.
+            before (Optional[RunnableConfig]): If provided, only checkpoints before the specified checkpoint ID are returned. Defaults to None.
+            limit (Optional[int]): Maximum number of checkpoints to return.
+
+        Yields:
+            Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
+        """
+        aiter_ = self.alist(config, filter=filter, before=before, limit=limit)
+        while True:
+            try:
+                yield asyncio.run_coroutine_threadsafe(
+                    anext(aiter_), self.loop
+                ).result()
+            except StopAsyncIteration:
+                break
+
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        """Get a checkpoint tuple from the database.
+
+        This method retrieves a checkpoint tuple from the Postgres database based on the
+        provided config. If the config contains a "checkpoint_id" key, the checkpoint with
+        the matching thread ID and "checkpoint_id" is retrieved. Otherwise, the latest checkpoint
+        for the given thread ID is retrieved.
+
+        Args:
+            config (RunnableConfig): The config to use for retrieving the checkpoint.
+
+        Returns:
+            Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aget_tuple(config), self.loop
+        ).result()
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Save a checkpoint to the database.
+
+        This method saves a checkpoint to the Postgres database. The checkpoint is associated
+        with the provided config and its parent config (if any).
+
+        Args:
+            config (RunnableConfig): The config to associate with the checkpoint.
+            checkpoint (Checkpoint): The checkpoint to save.
+            metadata (CheckpointMetadata): Additional metadata to save with the checkpoint.
+            new_versions (ChannelVersions): New channel versions as of this write.
+
+        Returns:
+            RunnableConfig: Updated configuration after storing the checkpoint.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput(config, checkpoint, metadata, new_versions), self.loop
+        ).result()
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        """Store intermediate writes linked to a checkpoint.
+
+        This method saves intermediate writes associated with a checkpoint to the database.
+
+        Args:
+            config (RunnableConfig): Configuration of the related checkpoint.
+            writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair.
+            task_id (str): Identifier for the task creating the writes.
+        """
+        return asyncio.run_coroutine_threadsafe(
+            self.aput_writes(config, writes, task_id), self.loop
+        ).result()

{langgraph_checkpoint_postgres-1.0.4 → langgraph_checkpoint_postgres-1.0.6}/langgraph/checkpoint/postgres/base.py

@@ -1,4 +1,3 @@
-from base64 import b64decode, b64encode
 from hashlib import md5
 from typing import Any, List, Optional, Tuple
 
@@ -13,7 +12,7 @@ from langgraph.checkpoint.base import (
     get_checkpoint_id,
 )
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
-from langgraph.checkpoint.serde.types import ChannelProtocol
+from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol
 
 MetadataInput = Optional[dict[str, Any]]
 
@@ -58,7 +57,7 @@ MIGRATIONS = [
     "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;",
 ]
 
-SELECT_SQL = """
+SELECT_SQL = f"""
 select
     thread_id,
     checkpoint,
@@ -77,12 +76,20 @@ select
     ) as channel_values,
     (
         select
-        array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob])
+        array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob] order by cw.task_id, cw.idx)
         from checkpoint_writes cw
         where cw.thread_id = checkpoints.thread_id
             and cw.checkpoint_ns = checkpoints.checkpoint_ns
             and cw.checkpoint_id = checkpoints.checkpoint_id
-    ) as pending_writes
+    ) as pending_writes,
+    (
+        select array_agg(array[cw.type::bytea, cw.blob] order by cw.idx)
+        from checkpoint_writes cw
+        where cw.thread_id = checkpoints.thread_id
+            and cw.checkpoint_ns = checkpoints.checkpoint_ns
+            and cw.checkpoint_id = checkpoints.parent_checkpoint_id
+            and cw.channel = '{TASKS}'
+    ) as pending_sends
 from checkpoints """
 
 UPSERT_CHECKPOINT_BLOBS_SQL = """
@@ -116,25 +123,23 @@ class BasePostgresSaver(BaseCheckpointSaver):
 
     jsonplus_serde = JsonPlusSerializer()
 
-    def _load_checkpoint(self, checkpoint: dict[str, Any]) -> Checkpoint:
-        if len(checkpoint["pending_sends"]) == 2 and all(
-            isinstance(a, str) for a in checkpoint["pending_sends"]
-        ):
-            type, bs = checkpoint["pending_sends"]
-            return {
-                **checkpoint,
-                "pending_sends": self.serde.loads_typed((type, b64decode(bs))),
-            }
-
-        return checkpoint
-
-    def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
-        type, bs = self.serde.dumps_typed(checkpoint["pending_sends"])
+    def _load_checkpoint(
+        self,
+        checkpoint: dict[str, Any],
+        channel_values: list[tuple[bytes, bytes, bytes]],
+        pending_sends: list[tuple[bytes, bytes]],
+    ) -> Checkpoint:
         return {
             **checkpoint,
-            "pending_sends": (type, b64encode(bs).decode()),
+            "pending_sends": [
+                self.serde.loads_typed((c.decode(), b)) for c, b in pending_sends or []
+            ],
+            "channel_values": self._load_blobs(channel_values),
         }
 
+    def _dump_checkpoint(self, checkpoint: Checkpoint) -> dict[str, Any]:
+        return {**checkpoint, "pending_sends": []}
+
     def _load_blobs(
         self, blob_values: list[tuple[bytes, bytes, bytes]]
     ) -> dict[str, Any]:
@@ -255,7 +260,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
         if config:
             wheres.append("thread_id = %s ")
             param_values.append(config["configurable"]["thread_id"])
-            if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
+            checkpoint_ns = config["configurable"].get("checkpoint_ns")
+            if checkpoint_ns is not None:
                 wheres.append("checkpoint_ns = %s")
                 param_values.append(checkpoint_ns)
 

{langgraph_checkpoint_postgres-1.0.4 → langgraph_checkpoint_postgres-1.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.4"
+version = "1.0.6"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,13 +10,13 @@ packages = [{ include = "langgraph" }]
 
 [tool.poetry.dependencies]
 python = "^3.9.0,<4.0"
-langgraph-checkpoint = "^1.0.1"
+langgraph-checkpoint = "^1.0.8"
 orjson = ">=3.10.1"
 psycopg = "^3.0.0"
 psycopg-pool = "^3.0.0"
 
 [tool.poetry.group.dev.dependencies]
-ruff = "^0.1.4"
+ruff = "^0.6.2"
 codespell = "^2.2.0"
 pytest = "^7.2.1"
 anyio = "^4.4.0"