langgraph-checkpoint-postgres 1.0.0__tar.gz

langgraph_checkpoint_postgres-1.0.0/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.1
+Name: langgraph-checkpoint-postgres
+Version: 1.0.0
+Summary: Library with a Postgres implementation of LangGraph checkpoint saver.
+Home-page: https://www.github.com/langchain-ai/langgraph
+License: MIT
+Requires-Python: >=3.9.0,<4.0.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+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: orjson (>=3.10.1)
+Requires-Dist: psycopg[binary] (>=3.1.19)
+Project-URL: Repository, https://www.github.com/langchain-ai/langgraph
+Description-Content-Type: text/markdown
+
+# LangGraph Checkpoint Postgres
+
+Implementation of LangGraph CheckpointSaver that uses Postgres.
+
+## Usage
+
+```python
+from langgraph.checkpoint.postgres import PostgresSaver
+
+write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+read_config = {"configurable": {"thread_id": "1"}}
+
+DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+        "current_tasks": {}
+    }
+
+    # store checkpoint
+    checkpointer.put(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    checkpointer.get(read_config)
+
+    # list checkpoints
+    list(checkpointer.list(read_config))
+```
+
+### Async
+
+```python
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+        "current_tasks": {}
+    }
+
+    # store checkpoint
+    await checkpointer.aput(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    await checkpointer.aget(read_config)
+
+    # list checkpoints
+    [c async for c in checkpointer.alist(read_config)]
+```

langgraph_checkpoint_postgres-1.0.0/README.md

@@ -0,0 +1,93 @@
+# LangGraph Checkpoint Postgres
+
+Implementation of LangGraph CheckpointSaver that uses Postgres.
+
+## Usage
+
+```python
+from langgraph.checkpoint.postgres import PostgresSaver
+
+write_config = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
+read_config = {"configurable": {"thread_id": "1"}}
+
+DB_URI = "postgres://postgres:postgres@localhost:5432/postgres?sslmode=disable"
+with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+        "current_tasks": {}
+    }
+
+    # store checkpoint
+    checkpointer.put(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    checkpointer.get(read_config)
+
+    # list checkpoints
+    list(checkpointer.list(read_config))
+```
+
+### Async
+
+```python
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+async with AsyncPostgresSaver.from_conn_string(DB_URI) as checkpointer:
+    checkpoint = {
+        "v": 1,
+        "ts": "2024-07-31T20:14:19.804150+00:00",
+        "id": "1ef4f797-8335-6428-8001-8a1503f9b875",
+        "channel_values": {
+            "my_key": "meow",
+            "node": "node"
+        },
+        "channel_versions": {
+            "__start__": 2,
+            "my_key": 3,
+            "start:node": 3,
+            "node": 3
+        },
+        "versions_seen": {
+            "__input__": {},
+            "__start__": {
+            "__start__": 1
+            },
+            "node": {
+            "start:node": 2
+            }
+        },
+        "pending_sends": [],
+        "current_tasks": {}
+    }
+
+    # store checkpoint
+    await checkpointer.aput(write_config, checkpoint, {}, {})
+
+    # load checkpoint
+    await checkpointer.aget(read_config)
+
+    # list checkpoints
+    [c async for c in checkpointer.alist(read_config)]
+```

langgraph_checkpoint_postgres-1.0.0/langgraph/checkpoint/postgres/__init__.py

@@ -0,0 +1,255 @@
+import threading
+from contextlib import contextmanager
+from typing import Any, Iterator, List, Optional
+
+from langchain_core.runnables import RunnableConfig
+from psycopg import Connection, Cursor, Pipeline
+from psycopg.rows import dict_row
+from psycopg.types.json import Jsonb
+
+from langgraph.checkpoint.base import (
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.postgres.base import (
+    BasePostgresSaver,
+)
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class PostgresSaver(BasePostgresSaver):
+    lock: threading.Lock
+
+    is_setup: bool
+
+    def __init__(
+        self,
+        conn: Connection,
+        pipe: Optional[Pipeline] = None,
+        serde: Optional[SerializerProtocol] = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = threading.Lock()
+        self.is_setup = False
+
+    @classmethod
+    @contextmanager
+    def from_conn_string(
+        cls, conn_string: str, *, pipeline: bool = False
+    ) -> Iterator["PostgresSaver"]:
+        """Create a new PostgresSaver instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+            pipeline (bool): whether to use Pipeline
+
+        Returns:
+            PostgresSaver: A new PostgresSaver instance.
+        """
+        with Connection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                with conn.pipeline() as pipe:
+                    yield PostgresSaver(conn, pipe)
+            else:
+                yield PostgresSaver(conn)
+
+    def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the SQLite database if they don't
+        already exist. It is called automatically when needed and should not be called
+        directly by the user.
+        """
+        if self.is_setup:
+            return
+        with self.lock:
+            create_table_queries = [
+                self.CREATE_CHECKPOINTS_SQL,
+                self.CREATE_CHECKPOINT_BLOBS_SQL,
+                self.CREATE_CHECKPOINT_WRITES_SQL,
+            ]
+            with self.conn.cursor(binary=True) as cur:
+                for query in create_table_queries:
+                    cur.execute(query)
+
+            if self.pipe:
+                self.pipe.sync()
+
+            self.is_setup = True
+
+    def list(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> Iterator[CheckpointTuple]:
+        self.setup()
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+        # if we change this to use .stream() we need to make sure to close the cursor
+        for value in self.conn.execute(query, args, binary=True):
+            yield CheckpointTuple(
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["checkpoint_id"],
+                    }
+                },
+                {
+                    **self._load_checkpoint(value["checkpoint"]),
+                    "channel_values": self._load_blobs(value["channel_values"]),
+                },
+                value["metadata"],
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["parent_checkpoint_id"],
+                    }
+                }
+                if value["parent_checkpoint_id"]
+                else None,
+            )
+
+    def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        self.setup()
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args = (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)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
+
+        with self._cursor() as cur:
+            cur = self.conn.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+
+            for value in cur:
+                return CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    {
+                        **self._load_checkpoint(value["checkpoint"]),
+                        "channel_values": self._load_blobs(value["channel_values"]),
+                    },
+                    value["metadata"],
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
+                    }
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    self._load_writes(value["pending_writes"]),
+                )
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop(
+            "checkpoint_id", configurable.pop("thread_ts", None)
+        )
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        with self._cursor(pipeline=True) as cur:
+            cur.executemany(
+                self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                self._dump_blobs(
+                    thread_id,
+                    checkpoint_ns,
+                    copy.pop("channel_values"),
+                    copy["channel_versions"],
+                    new_versions,
+                ),
+            )
+            cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    Jsonb(self._dump_checkpoint(copy)),
+                    Jsonb(metadata),
+                ),
+            )
+        return next_config
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: List[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        with self._cursor() as cur:
+            cur.executemany(
+                self.UPSERT_CHECKPOINT_WRITES_SQL,
+                self._dump_writes(
+                    config["configurable"]["thread_id"],
+                    config["configurable"]["checkpoint_ns"],
+                    config["configurable"]["checkpoint_id"],
+                    task_id,
+                    writes,
+                ),
+            )
+
+    @contextmanager
+    def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor]:
+        self.setup()
+        if self.pipe:
+            # a connection in pipeline mode can be used concurrently
+            # in multiple threads/coroutines, but only one cursor can be
+            # used at a time
+            try:
+                with self.conn.cursor(binary=True) as cur:
+                    yield cur
+            finally:
+                self.pipe.sync()
+        elif pipeline:
+            # a connection not in pipeline mode can only be used by one
+            # thread/coroutine at a time, so we acquire a lock
+            with self.lock, self.conn.pipeline(), self.conn.cursor(binary=True) as cur:
+                yield cur
+        else:
+            with self.lock, self.conn.cursor(binary=True) as cur:
+                yield cur

langgraph_checkpoint_postgres-1.0.0/langgraph/checkpoint/postgres/aio.py

@@ -0,0 +1,262 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Optional
+
+from langchain_core.runnables import RunnableConfig
+from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline
+from psycopg.rows import dict_row
+from psycopg.types.json import Jsonb
+
+from langgraph.checkpoint.base import (
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.postgres.base import BasePostgresSaver
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class AsyncPostgresSaver(BasePostgresSaver):
+    lock: asyncio.Lock
+
+    is_setup: bool
+
+    def __init__(
+        self,
+        conn: AsyncConnection,
+        pipe: Optional[AsyncPipeline] = None,
+        serde: Optional[SerializerProtocol] = None,
+    ) -> None:
+        super().__init__(serde=serde)
+        self.conn = conn
+        self.pipe = pipe
+        self.lock = asyncio.Lock()
+        self.is_setup = False
+
+    @classmethod
+    @asynccontextmanager
+    async def from_conn_string(
+        cls, conn_string: str, *, pipeline: bool = False
+    ) -> AsyncIterator["AsyncPostgresSaver"]:
+        """Create a new PostgresSaver instance from a connection string.
+
+        Args:
+            conn_string (str): The Postgres connection info string.
+            pipeline (bool): whether to use AsyncPipeline
+
+        Returns:
+            PostgresSaver: A new PostgresSaver instance.
+        """
+        async with await AsyncConnection.connect(
+            conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        ) as conn:
+            if pipeline:
+                async with conn.pipeline() as pipe:
+                    yield AsyncPostgresSaver(conn, pipe)
+            else:
+                yield AsyncPostgresSaver(conn)
+
+    async def setup(self) -> None:
+        """Set up the checkpoint database asynchronously.
+
+        This method creates the necessary tables in the SQLite database if they don't
+        already exist. It is called automatically when needed and should not be called
+        directly by the user.
+        """
+        if self.is_setup:
+            return
+        async with self.lock:
+            create_table_queries = [
+                self.CREATE_CHECKPOINTS_SQL,
+                self.CREATE_CHECKPOINT_BLOBS_SQL,
+                self.CREATE_CHECKPOINT_WRITES_SQL,
+            ]
+            async with self.conn.cursor() as cur:
+                for query in create_table_queries:
+                    await cur.execute(query)
+
+            if self.pipe:
+                await self.pipe.sync()
+
+            self.is_setup = True
+
+    async def alist(
+        self,
+        config: Optional[RunnableConfig],
+        *,
+        filter: Optional[dict[str, Any]] = None,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        await self.setup()
+        where, args = self._search_where(config, filter, before)
+        query = self.SELECT_SQL + where + " ORDER BY checkpoint_id DESC"
+        if limit:
+            query += f" LIMIT {limit}"
+        # if we change this to use .stream() we need to make sure to close the cursor
+        async for value in await self.conn.execute(query, args, binary=True):
+            yield CheckpointTuple(
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["checkpoint_id"],
+                    }
+                },
+                {
+                    **self._load_checkpoint(value["checkpoint"]),
+                    "channel_values": await asyncio.to_thread(
+                        self._load_blobs, value["channel_values"]
+                    ),
+                },
+                value["metadata"],
+                {
+                    "configurable": {
+                        "thread_id": value["thread_id"],
+                        "checkpoint_ns": value["checkpoint_ns"],
+                        "checkpoint_id": value["parent_checkpoint_id"],
+                    }
+                }
+                if value["parent_checkpoint_id"]
+                else None,
+            )
+
+    async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
+        await self.setup()
+        thread_id = config["configurable"]["thread_id"]
+        checkpoint_id = get_checkpoint_id(config)
+        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        if checkpoint_id:
+            args = (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)
+            where = "WHERE thread_id = %s AND checkpoint_ns = %s ORDER BY checkpoint_id DESC LIMIT 1"
+
+        async with self._cursor() as cur:
+            cur = await self.conn.execute(
+                self.SELECT_SQL + where,
+                args,
+                binary=True,
+            )
+
+            async for value in cur:
+                return CheckpointTuple(
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": value["checkpoint_id"],
+                        }
+                    },
+                    {
+                        **self._load_checkpoint(value["checkpoint"]),
+                        "channel_values": await asyncio.to_thread(
+                            self._load_blobs, value["channel_values"]
+                        ),
+                    },
+                    value["metadata"],
+                    {
+                        "configurable": {
+                            "thread_id": thread_id,
+                            "checkpoint_ns": checkpoint_ns,
+                            "checkpoint_id": value["parent_checkpoint_id"],
+                        }
+                    }
+                    if value["parent_checkpoint_id"]
+                    else None,
+                    await asyncio.to_thread(self._load_writes, value["pending_writes"]),
+                )
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        await self.setup()
+        configurable = config["configurable"].copy()
+        thread_id = configurable.pop("thread_id")
+        checkpoint_ns = configurable.pop("checkpoint_ns")
+        checkpoint_id = configurable.pop(
+            "checkpoint_id", configurable.pop("thread_ts", None)
+        )
+
+        copy = checkpoint.copy()
+        next_config = {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+        async with self._cursor(pipeline=True) as cur:
+            await cur.executemany(
+                self.UPSERT_CHECKPOINT_BLOBS_SQL,
+                await asyncio.to_thread(
+                    self._dump_blobs,
+                    thread_id,
+                    checkpoint_ns,
+                    copy.pop("channel_values"),
+                    copy["channel_versions"],
+                    new_versions,
+                ),
+            )
+            await cur.execute(
+                self.UPSERT_CHECKPOINTS_SQL,
+                (
+                    thread_id,
+                    checkpoint_ns,
+                    checkpoint["id"],
+                    checkpoint_id,
+                    Jsonb(self._dump_checkpoint(copy)),
+                    Jsonb(metadata),
+                ),
+            )
+        return next_config
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: list[tuple[str, Any]],
+        task_id: str,
+    ) -> None:
+        async with self._cursor() 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,
+                ),
+            )
+
+    @asynccontextmanager
+    async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor]:
+        await self.setup()
+        if self.pipe:
+            # a connection in pipeline mode can be used concurrently
+            # in multiple threads/coroutines, but only one cursor can be
+            # used at a time
+            try:
+                async with self.conn.cursor(binary=True) as cur:
+                    yield cur
+            finally:
+                await self.pipe.sync()
+        elif pipeline:
+            # a connection not in pipeline mode can only be used by one
+            # thread/coroutine at a time, so we acquire a lock
+            async with self.lock, self.conn.pipeline(), self.conn.cursor(
+                binary=True
+            ) as cur:
+                yield cur
+        else:
+            async with self.lock, self.conn.cursor(binary=True) as cur:
+                yield cur

langgraph_checkpoint_postgres-1.0.0/langgraph/checkpoint/postgres/base.py

@@ -0,0 +1,256 @@
+from base64 import b64decode, b64encode
+from hashlib import md5
+from typing import Any, List, Optional, Tuple
+
+from langchain_core.runnables import RunnableConfig
+from psycopg.types.json import Jsonb
+
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    EmptyChannelError,
+    get_checkpoint_id,
+)
+from langgraph.checkpoint.serde.types import ChannelProtocol
+
+MetadataInput = Optional[dict[str, Any]]
+
+SELECT_SQL = """
+select
+    thread_id,
+    checkpoint,
+    checkpoint_ns,
+    checkpoint_id,
+    parent_checkpoint_id,
+    metadata,
+    (
+        select array_agg(array[bl.channel::bytea, bl.type::bytea, bl.blob])
+        from jsonb_each_text(checkpoint -> 'channel_versions')
+        inner join checkpoint_blobs bl
+            on bl.thread_id = checkpoints.thread_id
+            and bl.checkpoint_ns = checkpoints.checkpoint_ns
+            and bl.channel = jsonb_each_text.key
+            and bl.version = jsonb_each_text.value
+    ) as channel_values,
+    (
+        select
+        array_agg(array[cw.task_id::text::bytea, cw.channel::bytea, cw.type::bytea, cw.blob])
+        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
+from checkpoints """
+
+CREATE_CHECKPOINTS_SQL = """
+    CREATE TABLE IF NOT EXISTS checkpoints (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    checkpoint_id TEXT NOT NULL,
+    parent_checkpoint_id TEXT,
+    type TEXT,
+    checkpoint JSONB NOT NULL,
+    metadata JSONB NOT NULL DEFAULT '{}',
+    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
+);
+"""
+CREATE_CHECKPOINT_BLOBS_SQL = """
+    CREATE TABLE IF NOT EXISTS checkpoint_blobs (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    channel TEXT NOT NULL,
+    version TEXT NOT NULL,
+    type TEXT NOT NULL,
+    blob BYTEA NOT NULL,
+    PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
+);"""
+
+CREATE_CHECKPOINT_WRITES_SQL = """
+    CREATE TABLE IF NOT EXISTS checkpoint_writes (
+    thread_id TEXT NOT NULL,
+    checkpoint_ns TEXT NOT NULL DEFAULT '',
+    checkpoint_id TEXT NOT NULL,
+    task_id TEXT NOT NULL,
+    idx INTEGER NOT NULL,
+    channel TEXT NOT NULL,
+    type TEXT,
+    blob BYTEA NOT NULL,
+    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
+);
+"""
+
+UPSERT_CHECKPOINT_BLOBS_SQL = """
+    INSERT INTO checkpoint_blobs (thread_id, checkpoint_ns, channel, version, type, blob)
+    VALUES (%s, %s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, channel, version) DO NOTHING
+"""
+
+UPSERT_CHECKPOINTS_SQL = """
+    INSERT INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, checkpoint, metadata)
+    VALUES (%s, %s, %s, %s, %s, %s)
+    ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id)
+    DO UPDATE SET
+        checkpoint = EXCLUDED.checkpoint,
+        metadata = EXCLUDED.metadata;
+"""
+
+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 NOTHING
+"""
+
+
+class BasePostgresSaver(BaseCheckpointSaver):
+    SELECT_SQL = SELECT_SQL
+    CREATE_CHECKPOINTS_SQL = CREATE_CHECKPOINTS_SQL
+    CREATE_CHECKPOINT_BLOBS_SQL = CREATE_CHECKPOINT_BLOBS_SQL
+    CREATE_CHECKPOINT_WRITES_SQL = CREATE_CHECKPOINT_WRITES_SQL
+    UPSERT_CHECKPOINT_BLOBS_SQL = UPSERT_CHECKPOINT_BLOBS_SQL
+    UPSERT_CHECKPOINTS_SQL = UPSERT_CHECKPOINTS_SQL
+    UPSERT_CHECKPOINT_WRITES_SQL = UPSERT_CHECKPOINT_WRITES_SQL
+
+    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"])
+        return {
+            **checkpoint,
+            "pending_sends": (type, b64encode(bs).decode()),
+        }
+
+    def _load_blobs(
+        self, blob_values: list[tuple[bytes, bytes, bytes]]
+    ) -> dict[str, Any]:
+        if not blob_values:
+            return {}
+        return {
+            k.decode(): self.serde.loads_typed((t.decode(), v))
+            for k, t, v in blob_values
+        }
+
+    def _dump_blobs(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        values: dict[str, Any],
+        versions: dict[str, str],
+        new_versions: Optional[dict[str, str]],
+    ) -> list[tuple[str, str, str, str, str, bytes]]:
+        if not versions:
+            return []
+
+        if new_versions:
+            versions = new_versions
+
+        return [
+            (
+                thread_id,
+                checkpoint_ns,
+                k,
+                ver,
+                *self.serde.dumps_typed(values[k]),
+            )
+            for k, ver in versions.items()
+            if k in values
+        ]
+
+    def _load_writes(
+        self, writes: list[tuple[bytes, bytes, bytes, bytes]]
+    ) -> list[tuple[str, str, Any]]:
+        return (
+            [
+                (
+                    tid.decode(),
+                    channel.decode(),
+                    self.serde.loads_typed((t.decode(), v)),
+                )
+                for tid, channel, t, v in writes
+            ]
+            if writes
+            else []
+        )
+
+    def _dump_writes(
+        self,
+        thread_id: str,
+        checkpoint_ns: str,
+        checkpoint_id: str,
+        task_id: str,
+        writes: list[tuple[str, Any]],
+    ) -> list[tuple[str, str, str, int, str, str, bytes]]:
+        return [
+            (
+                thread_id,
+                checkpoint_ns,
+                checkpoint_id,
+                task_id,
+                idx,
+                channel,
+                *self.serde.dumps_typed(value),
+            )
+            for idx, (channel, value) in enumerate(writes)
+        ]
+
+    def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
+        if current is None:
+            current_v = 0
+        elif isinstance(current, int):
+            current_v = current
+        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}"
+
+    def _search_where(
+        self,
+        config: Optional[RunnableConfig],
+        filter: MetadataInput,
+        before: Optional[RunnableConfig] = None,
+    ) -> Tuple[str, List[Any]]:
+        """Return WHERE clause predicates for alist() given config, filter, cursor.
+
+        This method returns a tuple of a string and a tuple of values. The string
+        is the parametered WHERE clause predicate (including the WHERE keyword):
+        "WHERE column1 = $1 AND column2 IS $2". The list of values contains the
+        values for each of the corresponding parameters.
+        """
+        wheres = []
+        param_values = []
+
+        # construct predicate for config filter
+        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:
+            wheres.append("metadata @> %s ")
+            param_values.append(Jsonb(filter))
+
+        # construct predicate for `before`
+        if before is not None:
+            wheres.append("checkpoint_id < %s ")
+            param_values.append(get_checkpoint_id(before))
+
+        return (
+            "WHERE " + " AND ".join(wheres) if wheres else "",
+            param_values,
+        )

langgraph_checkpoint_postgres-1.0.0/pyproject.toml

@@ -0,0 +1,52 @@
+[tool.poetry]
+name = "langgraph-checkpoint-postgres"
+version = "1.0.0"
+description = "Library with a Postgres implementation of LangGraph checkpoint saver."
+authors = []
+license = "MIT"
+readme = "README.md"
+repository = "https://www.github.com/langchain-ai/langgraph"
+packages = [{ include = "langgraph" }]
+
+[tool.poetry.dependencies]
+python = "^3.9.0,<4.0"
+langgraph-checkpoint = "^1.0.1"
+orjson = ">=3.10.1"
+psycopg = {extras = ["binary"], version = ">=3.1.19"}
+
+[tool.poetry.group.dev.dependencies]
+ruff = "^0.1.4"
+codespell = "^2.2.0"
+pytest = "^7.2.1"
+anyio = "^4.4.0"
+pytest-asyncio = "^0.21.1"
+pytest-mock = "^3.11.1"
+pytest-watch = "^4.2.0"
+mypy = "^1.10.0"
+psycopg-pool = "^3.2.2"
+langgraph-checkpoint = {path = "../checkpoint", develop = true}
+
+[tool.pytest.ini_options]
+# --strict-markers will raise errors on unknown marks.
+# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
+#
+# https://docs.pytest.org/en/7.1.x/reference/reference.html
+# --strict-config       any warnings encountered while parsing the `pytest`
+#                       section of the configuration file raise errors.
+addopts = "--strict-markers --strict-config --durations=5 -vv"
+asyncio_mode = "auto"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.ruff]
+lint.select = [
+  "E",  # pycodestyle
+  "F",  # Pyflakes
+  "UP", # pyupgrade
+  "B",  # flake8-bugbear
+  "I",  # isort
+]
+lint.ignore = ["E501", "B008", "UP007", "UP006"]