langgraph-checkpoint-postgres 1.0.5__tar.gz → 1.0.7__tar.gz

{langgraph_checkpoint_postgres-1.0.5 → langgraph_checkpoint_postgres-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.5
+Version: 1.0.7
 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.5 → langgraph_checkpoint_postgres-1.0.7}/langgraph/checkpoint/postgres/__init__.py

@@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -337,9 +338,14 @@ class PostgresSaver(BasePostgresSaver):
             writes (List[Tuple[str, Any]]): List of writes to store.
             task_id (str): Identifier for the task creating the writes.
         """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
         with self._cursor(pipeline=True) as cur:
             cur.executemany(
-                self.UPSERT_CHECKPOINT_WRITES_SQL,
+                query,
                 self._dump_writes(
                     config["configurable"]["thread_id"],
                     config["configurable"]["checkpoint_ns"],

{langgraph_checkpoint_postgres-1.0.5 → langgraph_checkpoint_postgres-1.0.7}/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
@@ -10,6 +10,7 @@ from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -51,11 +52,16 @@ class AsyncPostgresSaver(BasePostgresSaver):
         self.conn = conn
         self.pipe = pipe
         self.lock = asyncio.Lock()
+        self.loop = asyncio.get_running_loop()
 
     @classmethod
     @asynccontextmanager
     async def from_conn_string(
-        cls, conn_string: str, *, pipeline: bool = False
+        cls,
+        conn_string: str,
+        *,
+        pipeline: bool = False,
+        serde: Optional[SerializerProtocol] = None,
     ) -> AsyncIterator["AsyncPostgresSaver"]:
         """Create a new PostgresSaver instance from a connection string.
 
@@ -71,9 +77,9 @@ class AsyncPostgresSaver(BasePostgresSaver):
         ) as conn:
             if pipeline:
                 async with conn.pipeline() as pipe:
-                    yield AsyncPostgresSaver(conn, pipe)
+                    yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde)
             else:
-                yield AsyncPostgresSaver(conn)
+                yield AsyncPostgresSaver(conn=conn, serde=serde)
 
     async def setup(self) -> None:
         """Set up the checkpoint database asynchronously.
@@ -291,18 +297,21 @@ class AsyncPostgresSaver(BasePostgresSaver):
             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.
         """
+        query = (
+            self.UPSERT_CHECKPOINT_WRITES_SQL
+            if all(w[0] in WRITES_IDX_MAP for w in writes)
+            else self.INSERT_CHECKPOINT_WRITES_SQL
+        )
+        params = await asyncio.to_thread(
+            self._dump_writes,
+            config["configurable"]["thread_id"],
+            config["configurable"]["checkpoint_ns"],
+            config["configurable"]["checkpoint_id"],
+            task_id,
+            writes,
+        )
         async with self._cursor(pipeline=True) as cur:
-            await cur.executemany(
-                self.UPSERT_CHECKPOINT_WRITES_SQL,
-                await asyncio.to_thread(
-                    self._dump_writes,
-                    config["configurable"]["thread_id"],
-                    config["configurable"]["checkpoint_ns"],
-                    config["configurable"]["checkpoint_id"],
-                    task_id,
-                    writes,
-                ),
-            )
+            await cur.executemany(query, params)
 
     @asynccontextmanager
     async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
@@ -329,3 +338,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.5 → langgraph_checkpoint_postgres-1.0.7}/langgraph/checkpoint/postgres/base.py

@@ -1,4 +1,4 @@
-from hashlib import md5
+import random
 from typing import Any, List, Optional, Tuple
 
 from langchain_core.runnables import RunnableConfig
@@ -8,7 +8,6 @@ from langgraph.checkpoint.base import (
     WRITES_IDX_MAP,
     BaseCheckpointSaver,
     Checkpoint,
-    EmptyChannelError,
     get_checkpoint_id,
 )
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
@@ -108,6 +107,15 @@ UPSERT_CHECKPOINTS_SQL = """
 """
 
 UPSERT_CHECKPOINT_WRITES_SQL = """
+    INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
+    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO UPDATE SET
+        channel = EXCLUDED.channel,
+        type = EXCLUDED.type,
+        blob = EXCLUDED.blob;
+"""
+
+INSERT_CHECKPOINT_WRITES_SQL = """
     INSERT INTO checkpoint_writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, blob)
     VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
     ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
@@ -120,6 +128,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
     UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
     UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
     UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+    INSERT_CHECKPOINT_WRITES_SQL = INSERT_CHECKPOINT_WRITES_SQL
 
     jsonplus_serde = JsonPlusSerializer()
 
@@ -217,13 +226,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
         return self.jsonplus_serde.loads(self.jsonplus_serde.dumps(metadata))
 
     def _dump_metadata(self, metadata) -> str:
-        serialized_metadata_type, serialized_metadata = self.jsonplus_serde.dumps_typed(
-            metadata
-        )
-        if serialized_metadata_type != "json":
-            raise TypeError(
-                f"Failed to properly serialize metadata -- expected 'json', got '{serialized_metadata_type}'"
-            )
+        serialized_metadata = self.jsonplus_serde.dumps(metadata)
         return serialized_metadata.decode()
 
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
@@ -234,11 +237,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
         else:
             current_v = int(current.split(".")[0])
         next_v = current_v + 1
-        try:
-            next_h = md5(self.serde.dumps_typed(channel.checkpoint())[1]).hexdigest()
-        except EmptyChannelError:
-            next_h = ""
-        return f"{next_v:032}.{next_h}"
+        next_h = random.random()
+        return f"{next_v:032}.{next_h:016}"
 
     def _search_where(
         self,

{langgraph_checkpoint_postgres-1.0.5 → langgraph_checkpoint_postgres-1.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.5"
+version = "1.0.7"
 description = "Library with a Postgres implementation of LangGraph checkpoint saver."
 authors = []
 license = "MIT"
@@ -10,7 +10,7 @@ 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"