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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: langgraph-checkpoint-postgres
-Version: 1.0.6
+Version: 1.0.8
 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.8,<2.0.0)
+Requires-Dist: langgraph-checkpoint (>=1.0.11,<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.6 → langgraph_checkpoint_postgres-1.0.8}/langgraph/checkpoint/postgres/__init__.py

@@ -1,15 +1,16 @@
 import threading
 from contextlib import contextmanager
-from typing import Any, Iterator, List, Optional, Union
+from typing import Any, Iterator, Optional, Sequence, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import Connection, Cursor, Pipeline
 from psycopg.errors import UndefinedTable
-from psycopg.rows import dict_row
+from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import ConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -21,9 +22,11 @@ from langgraph.checkpoint.postgres.base import (
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
+Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]]
+
 
 @contextmanager
-def _get_connection(conn: Union[Connection, ConnectionPool]) -> Iterator[Connection]:
+def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]:
     if isinstance(conn, Connection):
         yield conn
     elif isinstance(conn, ConnectionPool):
@@ -38,7 +41,7 @@ class PostgresSaver(BasePostgresSaver):
 
     def __init__(
         self,
-        conn: Union[Connection, ConnectionPool],
+        conn: Conn,
         pipe: Optional[Pipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -84,9 +87,13 @@ class PostgresSaver(BasePostgresSaver):
         """
         with self._cursor() as cur:
             try:
-                version = cur.execute(
+                row = cur.execute(
                     "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
-                ).fetchone()["v"]
+                ).fetchone()
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
             except UndefinedTable:
                 version = -1
             for v, migration in zip(
@@ -211,7 +218,7 @@ class PostgresSaver(BasePostgresSaver):
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         if checkpoint_id:
-            args = (thread_id, checkpoint_ns, checkpoint_id)
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
             where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
         else:
             args = (thread_id, checkpoint_ns)
@@ -305,7 +312,7 @@ class PostgresSaver(BasePostgresSaver):
                 self._dump_blobs(
                     thread_id,
                     checkpoint_ns,
-                    copy.pop("channel_values"),
+                    copy.pop("channel_values"),  # type: ignore[misc]
                     new_versions,
                 ),
             )
@@ -325,7 +332,7 @@ class PostgresSaver(BasePostgresSaver):
     def put_writes(
         self,
         config: RunnableConfig,
-        writes: List[tuple[str, Any]],
+        writes: Sequence[tuple[str, Any]],
         task_id: str,
     ) -> None:
         """Store intermediate writes linked to a checkpoint.
@@ -337,9 +344,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"],
@@ -350,7 +362,7 @@ class PostgresSaver(BasePostgresSaver):
             )
 
     @contextmanager
-    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
         with _get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently

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

@@ -1,15 +1,16 @@
 import asyncio
 from contextlib import asynccontextmanager
-from typing import Any, AsyncIterator, Iterator, List, Optional, Union
+from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Union
 
 from langchain_core.runnables import RunnableConfig
 from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
 from psycopg.errors import UndefinedTable
-from psycopg.rows import dict_row
+from psycopg.rows import DictRow, dict_row
 from psycopg.types.json import Jsonb
 from psycopg_pool import AsyncConnectionPool
 
 from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
     ChannelVersions,
     Checkpoint,
     CheckpointMetadata,
@@ -19,11 +20,13 @@ from langgraph.checkpoint.base import (
 from langgraph.checkpoint.postgres.base import BasePostgresSaver
 from langgraph.checkpoint.serde.base import SerializerProtocol
 
+Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]]
+
 
 @asynccontextmanager
 async def _get_connection(
-    conn: Union[AsyncConnection, AsyncConnectionPool],
-) -> AsyncIterator[AsyncConnection]:
+    conn: Conn,
+) -> AsyncIterator[AsyncConnection[DictRow]]:
     if isinstance(conn, AsyncConnection):
         yield conn
     elif isinstance(conn, AsyncConnectionPool):
@@ -38,7 +41,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
 
     def __init__(
         self,
-        conn: Union[AsyncConnection, AsyncConnectionPool],
+        conn: Conn,
         pipe: Optional[AsyncPipeline] = None,
         serde: Optional[SerializerProtocol] = None,
     ) -> None:
@@ -56,7 +59,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
     @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.
 
@@ -72,9 +79,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.
@@ -88,7 +95,11 @@ class AsyncPostgresSaver(BasePostgresSaver):
                 results = await cur.execute(
                     "SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
                 )
-                version = (await results.fetchone())["v"]
+                row = await results.fetchone()
+                if row is None:
+                    version = -1
+                else:
+                    version = row["v"]
             except UndefinedTable:
                 version = -1
             for v, migration in zip(
@@ -175,7 +186,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
         checkpoint_id = get_checkpoint_id(config)
         checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
         if checkpoint_id:
-            args = (thread_id, checkpoint_ns, checkpoint_id)
+            args: tuple[Any, ...] = (thread_id, checkpoint_ns, checkpoint_id)
             where = "WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s"
         else:
             args = (thread_id, checkpoint_ns)
@@ -260,7 +271,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
                     self._dump_blobs,
                     thread_id,
                     checkpoint_ns,
-                    copy.pop("channel_values"),
+                    copy.pop("channel_values"),  # type: ignore[misc]
                     new_versions,
                 ),
             )
@@ -280,7 +291,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
     async def aput_writes(
         self,
         config: RunnableConfig,
-        writes: list[tuple[str, Any]],
+        writes: Sequence[tuple[str, Any]],
         task_id: str,
     ) -> None:
         """Store intermediate writes linked to a checkpoint asynchronously.
@@ -292,21 +303,26 @@ 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]:
+    async def _cursor(
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
         async with _get_connection(self.conn) as conn:
             if self.pipe:
                 # a connection in pipeline mode can be used concurrently
@@ -357,7 +373,8 @@ class AsyncPostgresSaver(BasePostgresSaver):
         while True:
             try:
                 yield asyncio.run_coroutine_threadsafe(
-                    anext(aiter_), self.loop
+                    anext(aiter_),
+                    self.loop,
                 ).result()
             except StopAsyncIteration:
                 break
@@ -376,6 +393,18 @@ class AsyncPostgresSaver(BasePostgresSaver):
         Returns:
             Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
         """
+        try:
+            # check if we are in the main thread, only bg threads can block
+            # we don't check in other methods to avoid the overhead
+            if asyncio.get_running_loop() is self.loop:
+                raise asyncio.InvalidStateError(
+                    "Synchronous calls to AsyncPostgresSaver are only allowed from a "
+                    "different thread. From the main thread, use the async interface."
+                    "For example, use `await checkpointer.aget_tuple(...)` or `await "
+                    "graph.ainvoke(...)`."
+                )
+        except RuntimeError:
+            pass
         return asyncio.run_coroutine_threadsafe(
             self.aget_tuple(config), self.loop
         ).result()
@@ -408,7 +437,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
     def put_writes(
         self,
         config: RunnableConfig,
-        writes: List[tuple[str, Any]],
+        writes: Sequence[tuple[str, Any]],
         task_id: str,
     ) -> None:
         """Store intermediate writes linked to a checkpoint.

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

@@ -1,5 +1,5 @@
-from hashlib import md5
-from typing import Any, List, Optional, Tuple
+import random
+from typing import Any, List, Optional, Sequence, Tuple, cast
 
 from langchain_core.runnables import RunnableConfig
 from psycopg.types.json import Jsonb
@@ -7,8 +7,9 @@ from psycopg.types.json import Jsonb
 from langgraph.checkpoint.base import (
     WRITES_IDX_MAP,
     BaseCheckpointSaver,
+    ChannelVersions,
     Checkpoint,
-    EmptyChannelError,
+    CheckpointMetadata,
     get_checkpoint_id,
 )
 from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
@@ -108,18 +109,28 @@ 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
 """
 
 
-class BasePostgresSaver(BaseCheckpointSaver):
+class BasePostgresSaver(BaseCheckpointSaver[str]):
     SELECT_SQL = SELECT_SQL
     MIGRATIONS = MIGRATIONS
     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()
 
@@ -156,8 +167,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
         thread_id: str,
         checkpoint_ns: str,
         values: dict[str, Any],
-        versions: dict[str, str],
-    ) -> list[tuple[str, str, str, str, str, bytes]]:
+        versions: ChannelVersions,
+    ) -> list[tuple[str, str, str, str, str, Optional[bytes]]]:
         if not versions:
             return []
 
@@ -166,7 +177,7 @@ class BasePostgresSaver(BaseCheckpointSaver):
                 thread_id,
                 checkpoint_ns,
                 k,
-                ver,
+                cast(str, ver),
                 *(
                     self.serde.dumps_typed(values[k])
                     if k in values
@@ -198,8 +209,8 @@ class BasePostgresSaver(BaseCheckpointSaver):
         checkpoint_ns: str,
         checkpoint_id: str,
         task_id: str,
-        writes: list[tuple[str, Any]],
-    ) -> list[tuple[str, str, str, int, str, str, bytes]]:
+        writes: Sequence[tuple[str, Any]],
+    ) -> list[tuple[str, str, str, str, int, str, str, bytes]]:
         return [
             (
                 thread_id,
@@ -213,17 +224,11 @@ class BasePostgresSaver(BaseCheckpointSaver):
             for idx, (channel, value) in enumerate(writes)
         ]
 
-    def _load_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
+    def _load_metadata(self, metadata: dict[str, Any]) -> CheckpointMetadata:
         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}'"
-            )
+    def _dump_metadata(self, metadata: CheckpointMetadata) -> str:
+        serialized_metadata = self.jsonplus_serde.dumps(metadata)
         return serialized_metadata.decode()
 
     def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
@@ -234,11 +239,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.6 → langgraph_checkpoint_postgres-1.0.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langgraph-checkpoint-postgres"
-version = "1.0.6"
+version = "1.0.8"
 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.8"
+langgraph-checkpoint = "^1.0.11"
 orjson = ">=3.10.1"
 psycopg = "^3.0.0"
 psycopg-pool = "^3.0.0"
@@ -51,3 +51,13 @@ lint.select = [
   "I",  # isort
 ]
 lint.ignore = ["E501", "B008", "UP007", "UP006"]
+
+[tool.mypy]
+# https://mypy.readthedocs.io/en/stable/config_file.html
+disallow_untyped_defs = "True"
+explicit_package_bases = "True"
+warn_no_return = "False"
+warn_unused_ignores = "True"
+warn_redundant_casts = "True"
+allow_redefinition = "True"
+disable_error_code = "typeddict-item, return-value"