From 0d0665a6e39833575278c54f7fd12705ae989d49 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 22 Nov 2024 11:56:16 -0500 Subject: [PATCH 1/8] docs: Add resource allocation (#2511) --- docs/docs/concepts/langgraph_cloud.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/docs/concepts/langgraph_cloud.md b/docs/docs/concepts/langgraph_cloud.md index 371e8bb47..11e5c8afc 100644 --- a/docs/docs/concepts/langgraph_cloud.md +++ b/docs/docs/concepts/langgraph_cloud.md @@ -14,6 +14,13 @@ A **deployment** is an instance of a LangGraph API. A single deployment can have See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment. +## Resource Allocation + +| **Deployment Type** | **CPU** | **Memory** | **Scaling** | +|---------------------|---------|------------|---------------------| +| Development | 1 CPU | 1 GB | Up to 1 container | +| Production | 1 CPU | 2 GB | Up to 10 containers | + ## Revision A revision is an iteration of a [deployment](#deployment). When a new deployment is created, an initial revision is automatically created. To deploy new code changes or update environment variable configurations for a deployment, a new revision must be created. When a revision is created, a new container image is built automatically. @@ -33,6 +40,7 @@ A high-level diagram of a Cloud SaaS deployment. ![diagram](img/langgraph_cloud_architecture.png) + ## Related - [Deployment Options](./deployment_options.md) From 65f515e020682986659eb31cbe4fc5c41269c0f6 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 22 Nov 2024 12:17:36 -0500 Subject: [PATCH 2/8] docs: add helm chart link (#2512) --- docs/docs/concepts/self_hosted.md | 6 +++++- docs/docs/how-tos/deploy-self-hosted.md | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/docs/docs/concepts/self_hosted.md b/docs/docs/concepts/self_hosted.md index ffa26a873..d2f236705 100644 --- a/docs/docs/concepts/self_hosted.md +++ b/docs/docs/concepts/self_hosted.md @@ -7,7 +7,7 @@ ## Versions -There are two versions of the self hosted deployment: [Self-Hosted Enterprise](./deployment_options.md#self-hosted-enterprise) and [Self-Hosted Lite](./deployment_options.md#self-hosted-lite). +There are two versions of the self-hosted deployment: [Self-Hosted Enterprise](./deployment_options.md#self-hosted-enterprise) and [Self-Hosted Lite](./deployment_options.md#self-hosted-lite). ### Self-Hosted Lite @@ -34,6 +34,10 @@ To use the Self-Hosted Enterprise version, you must acquire a license key that y For step-by-step instructions, see [How to set up a self-hosted deployment of LangGraph](../how-tos/deploy-self-hosted.md). +## Helm Chart + +If you would like to deploy LangGraph Cloud on Kubernetes, you can use this [Helm chart](https://github.com/langchain-ai/helm/blob/main/charts/langgraph-cloud/README.md). + ## Related - [How to set up a self-hosted deployment of LangGraph](../how-tos/deploy-self-hosted.md). diff --git a/docs/docs/how-tos/deploy-self-hosted.md b/docs/docs/how-tos/deploy-self-hosted.md index 5e1bf93e4..88d7c72e4 100644 --- a/docs/docs/how-tos/deploy-self-hosted.md +++ b/docs/docs/how-tos/deploy-self-hosted.md @@ -17,6 +17,10 @@ You will need to do the following: 2. Build a docker image with the [LangGraph Server](../concepts/langgraph_server.md) using the [LangGraph CLI](../concepts/langgraph_cli.md). 3. Deploy a web server that will run the docker image and pass in the necessary environment variables. +## Helm Chart + +If you would like to deploy LangGraph Cloud on Kubernetes, you can use this [Helm chart](https://github.com/langchain-ai/helm/blob/main/charts/langgraph-cloud/README.md). + ## Environment Variables You will eventually need to pass in the following environment variables to the LangGraph Deploy server: @@ -70,7 +74,7 @@ If you want to run this quickly without setting up a separate Redis and Postgres * You need to replace `my-image` with the name of the image you built in the previous step (from `langgraph build`). and you should provide appropriate values for `REDIS_URI`, `DATABASE_URI`, and `LANGSMITH_API_KEY`. * If your application requires additional environment variables, you can pass them in a similar way. - * If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise), you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable. + * If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise, you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable. ### Using Docker Compose From b4900341e40a4dde6ebef740a2df262b04064213 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 22 Nov 2024 14:07:54 -0500 Subject: [PATCH 3/8] docs: fix broken link (#2514) We need to check later why CI didn't fail with original PR that broke the link --- docs/docs/how-tos/deploy-self-hosted.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/docs/how-tos/deploy-self-hosted.md b/docs/docs/how-tos/deploy-self-hosted.md index 88d7c72e4..d0d3a3efe 100644 --- a/docs/docs/how-tos/deploy-self-hosted.md +++ b/docs/docs/how-tos/deploy-self-hosted.md @@ -74,7 +74,7 @@ If you want to run this quickly without setting up a separate Redis and Postgres * You need to replace `my-image` with the name of the image you built in the previous step (from `langgraph build`). and you should provide appropriate values for `REDIS_URI`, `DATABASE_URI`, and `LANGSMITH_API_KEY`. * If your application requires additional environment variables, you can pass them in a similar way. - * If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise, you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable. + * If using [Self-Hosted Enterprise](../concepts/deployment_options.md#self-hosted-enterprise), you must provide `LANGGRAPH_CLOUD_LICENSE_KEY` as an additional environment variable. ### Using Docker Compose From 3351d4f6c5c3d5fcb048fd9d6f45bb52a2c57a4d Mon Sep 17 00:00:00 2001 From: Mingqi Hu Date: Sat, 23 Nov 2024 03:17:59 +0800 Subject: [PATCH 4/8] docs: fix typo (#2510) `python-dotenv` not `python-dot-env` Signed-off-by: Mingqi Co-authored-by: Eugene Yurtsev --- docs/docs/tutorials/langgraph-platform/local-server.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/docs/tutorials/langgraph-platform/local-server.md b/docs/docs/tutorials/langgraph-platform/local-server.md index a41333921..5778782e4 100644 --- a/docs/docs/tutorials/langgraph-platform/local-server.md +++ b/docs/docs/tutorials/langgraph-platform/local-server.md @@ -9,7 +9,7 @@ This is a quick start guide to help you get a LangGraph app up and running local ## Install the LangGraph CLI ```bash -pip install "langgraph-cli[inmem]==0.1.58" python-dot-env +pip install "langgraph-cli[inmem]==0.1.58" python-dotenv ``` ## 🌱 Create a LangGraph App From 4f4e7a69818d4f081400711c6382ce83b38d2d16 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Fri, 22 Nov 2024 14:50:31 -0500 Subject: [PATCH 5/8] docs: more fixes for python version (#2515) --- docs/docs/tutorials/langgraph-platform/local-server.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/docs/tutorials/langgraph-platform/local-server.md b/docs/docs/tutorials/langgraph-platform/local-server.md index 5778782e4..9c8e3ea9f 100644 --- a/docs/docs/tutorials/langgraph-platform/local-server.md +++ b/docs/docs/tutorials/langgraph-platform/local-server.md @@ -4,6 +4,7 @@ This is a quick start guide to help you get a LangGraph app up and running local !!! info "Requirements" + - Python >= 3.11 - [LangGraph CLI](https://langchain-ai.github.io/langgraph/cloud/reference/cli/): Requires langchain-cli[inmem] >= 0.1.58 ## Install the LangGraph CLI @@ -32,6 +33,14 @@ Create a new app from the `react-agent` template. This template is a simple agen If you use `langgraph new` without specifying a template, you will be presented with an interactive menu that will allow you to choose from a list of available templates. +## Install Dependencies + +In the root of your new LangGraph app, install the dependencies: + +```shell +pip install . +``` + ## Create a `.env` file You will find a `.env.example` in the root of your new LangGraph app. Create From 486d5412af683c509674f4c3086cb2102c87aa5e Mon Sep 17 00:00:00 2001 From: Talha Munir <103119362+devs-talha@users.noreply.github.com> Date: Sun, 24 Nov 2024 00:37:55 +0500 Subject: [PATCH 6/8] docs: Fix grammatical mistake in introduction.ipynb (#2521) --- docs/docs/tutorials/introduction.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/docs/tutorials/introduction.ipynb b/docs/docs/tutorials/introduction.ipynb index 712fbc28f..119fbf861 100644 --- a/docs/docs/tutorials/introduction.ipynb +++ b/docs/docs/tutorials/introduction.ipynb @@ -2043,7 +2043,7 @@ "\n", "So far, we've relied on a simple state (it's just a list of messages!). You can go far with this simple state, but if you want to define complex behavior without relying on the message list, you can add additional fields to the state. In this section, we will extend our chat bot with a new node to illustrate this.\n", "\n", - "In the examples above, we involved a human deterministically: the graph __always__ interrupted whenever an tool was invoked. Suppose we wanted our chat bot to have the choice of relying on a human.\n", + "In the examples above, we involved a human deterministically: the graph __always__ interrupted whenever a tool was invoked. Suppose we wanted our chat bot to have the choice of relying on a human.\n", "\n", "One way to do this is to create a passthrough \"human\" node, before which the graph will always stop. We will only execute this node if the LLM invokes a \"human\" tool. For our convenience, we will include an \"ask_human\" flag in our graph state that we will flip if the LLM calls this tool.\n", "\n", From 328ef609afc74a91ecb52a05ac06447125c12a91 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 25 Nov 2024 11:59:49 -0800 Subject: [PATCH 7/8] [CLI] Python path (#2531) --- libs/cli/langgraph_cli/cli.py | 8 ++++++++ libs/cli/pyproject.toml | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/libs/cli/langgraph_cli/cli.py b/libs/cli/langgraph_cli/cli.py index 07c894fac..c304252b3 100644 --- a/libs/cli/langgraph_cli/cli.py +++ b/libs/cli/langgraph_cli/cli.py @@ -1,3 +1,4 @@ +import os import pathlib import shutil import sys @@ -598,6 +599,13 @@ def dev( ) from None config_json = langgraph_cli.config.validate_config_file(config) + cwd = os.getcwd() + sys.path.append(cwd) + dependencies = config_json.get("dependencies", []) + for dep in dependencies: + dep_path = pathlib.Path(cwd) / dep + if dep_path.is_dir() and dep_path.exists(): + sys.path.append(str(dep_path)) graphs = config_json.get("graphs", {}) run_server( diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index a271a073b..163bb4a31 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-cli" -version = "0.1.58" +version = "0.1.59" description = "CLI for interacting with LangGraph API" authors = [] license = "MIT" From 98935e1ffd8ab82d5587b90fe154ecf92fa61121 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Mon, 25 Nov 2024 12:19:52 -0800 Subject: [PATCH 8/8] fix: Fix race condition in PostgresSaver (#2494) Signed-off-by: Tyler Ball Co-authored-by: Phoenix Logan Co-authored-by: Tyler Ball <2481463+tyler-ball@users.noreply.github.com> --- .../langgraph/checkpoint/postgres/__init__.py | 26 +- .../checkpoint/postgres/_ainternal.py | 23 + .../checkpoint/postgres/_internal.py | 21 + .../langgraph/checkpoint/postgres/aio.py | 65 ++- .../langgraph/store/postgres/aio.py | 254 +++++---- .../langgraph/store/postgres/base.py | 217 +++++--- libs/checkpoint-postgres/tests/conftest.py | 6 +- .../tests/test_async_store.py | 223 +++----- libs/checkpoint-postgres/tests/test_store.py | 501 +++++++----------- libs/langgraph/Makefile | 7 +- libs/langgraph/tests/conftest.py | 109 +++- libs/langgraph/tests/test_pregel.py | 115 ++-- libs/langgraph/tests/test_pregel_async.py | 90 +++- 13 files changed, 906 insertions(+), 751 deletions(-) create mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py create mode 100644 libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index b8138a945..2107b05a3 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -1,6 +1,6 @@ import threading from contextlib import contextmanager -from typing import Any, Iterator, Optional, Sequence, Union +from typing import Any, Iterator, Optional, Sequence from langchain_core.runnables import RunnableConfig from psycopg import Capabilities, Connection, Cursor, Pipeline @@ -17,21 +17,11 @@ from langgraph.checkpoint.base import ( CheckpointTuple, get_checkpoint_id, ) +from langgraph.checkpoint.postgres import _internal from langgraph.checkpoint.postgres.base import BasePostgresSaver from langgraph.checkpoint.serde.base import SerializerProtocol -Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]] - - -@contextmanager -def _get_connection(conn: Conn) -> Iterator[Connection[DictRow]]: - if isinstance(conn, Connection): - yield conn - elif isinstance(conn, ConnectionPool): - with conn.connection() as conn: - yield conn - else: - raise TypeError(f"Invalid connection type: {type(conn)}") +Conn = _internal.Conn # For backward compatibility class PostgresSaver(BasePostgresSaver): @@ -39,7 +29,7 @@ class PostgresSaver(BasePostgresSaver): def __init__( self, - conn: Conn, + conn: _internal.Conn, pipe: Optional[Pipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: @@ -73,9 +63,9 @@ class PostgresSaver(BasePostgresSaver): ) as conn: if pipeline: with conn.pipeline() as pipe: - yield PostgresSaver(conn, pipe) + yield cls(conn, pipe) else: - yield PostgresSaver(conn) + yield cls(conn) def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -373,7 +363,7 @@ class PostgresSaver(BasePostgresSaver): Will be applied regardless of whether the PostgresSaver instance was initialized with a pipeline. If pipeline mode is not supported, will fall back to using transaction context manager. """ - with _get_connection(self.conn) as conn: + with _internal.get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be @@ -403,4 +393,4 @@ class PostgresSaver(BasePostgresSaver): yield cur -__all__ = ["PostgresSaver", "Conn"] +__all__ = ["PostgresSaver", "BasePostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py new file mode 100644 index 000000000..a0b8b10f5 --- /dev/null +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_ainternal.py @@ -0,0 +1,23 @@ +"""Shared async utility functions for the Postgres checkpoint & storage classes.""" + +from contextlib import asynccontextmanager +from typing import AsyncIterator, Union + +from psycopg import AsyncConnection +from psycopg.rows import DictRow +from psycopg_pool import AsyncConnectionPool + +Conn = Union[AsyncConnection[DictRow], AsyncConnectionPool[AsyncConnection[DictRow]]] + + +@asynccontextmanager +async def get_connection( + conn: Conn, +) -> AsyncIterator[AsyncConnection[DictRow]]: + if isinstance(conn, AsyncConnection): + yield conn + elif isinstance(conn, AsyncConnectionPool): + async with conn.connection() as conn: + yield conn + else: + raise TypeError(f"Invalid connection type: {type(conn)}") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py new file mode 100644 index 000000000..b703262f2 --- /dev/null +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/_internal.py @@ -0,0 +1,21 @@ +"""Shared utility functions for the Postgres checkpoint & storage classes.""" + +from contextlib import contextmanager +from typing import Iterator, Union + +from psycopg import Connection +from psycopg.rows import DictRow +from psycopg_pool import ConnectionPool + +Conn = Union[Connection[DictRow], ConnectionPool[Connection[DictRow]]] + + +@contextmanager +def get_connection(conn: Conn) -> Iterator[Connection[DictRow]]: + if isinstance(conn, Connection): + yield conn + elif isinstance(conn, ConnectionPool): + with conn.connection() as conn: + yield conn + else: + raise TypeError(f"Invalid connection type: {type(conn)}") diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 5b67e4ca9..589520efc 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -1,6 +1,6 @@ import asyncio from contextlib import asynccontextmanager -from typing import Any, AsyncIterator, Iterator, Optional, Sequence, Union +from typing import Any, AsyncIterator, Iterator, Optional, Sequence from langchain_core.runnables import RunnableConfig from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities @@ -17,23 +17,11 @@ from langgraph.checkpoint.base import ( CheckpointTuple, get_checkpoint_id, ) +from langgraph.checkpoint.postgres import _ainternal 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: Conn, -) -> AsyncIterator[AsyncConnection[DictRow]]: - if isinstance(conn, AsyncConnection): - yield conn - elif isinstance(conn, AsyncConnectionPool): - async with conn.connection() as conn: - yield conn - else: - raise TypeError(f"Invalid connection type: {type(conn)}") +Conn = _ainternal.Conn # For backward compatibility class AsyncPostgresSaver(BasePostgresSaver): @@ -41,7 +29,7 @@ class AsyncPostgresSaver(BasePostgresSaver): def __init__( self, - conn: Conn, + conn: _ainternal.Conn, pipe: Optional[AsyncPipeline] = None, serde: Optional[SerializerProtocol] = None, ) -> None: @@ -80,9 +68,9 @@ class AsyncPostgresSaver(BasePostgresSaver): ) as conn: if pipeline: async with conn.pipeline() as pipe: - yield AsyncPostgresSaver(conn=conn, pipe=pipe, serde=serde) + yield cls(conn=conn, pipe=pipe, serde=serde) else: - yield AsyncPostgresSaver(conn=conn, serde=serde) + yield cls(conn=conn, serde=serde) async def setup(self) -> None: """Set up the checkpoint database asynchronously. @@ -157,15 +145,17 @@ class AsyncPostgresSaver(BasePostgresSaver): value["pending_sends"], ), self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": value["thread_id"], - "checkpoint_ns": value["checkpoint_ns"], - "checkpoint_id": value["parent_checkpoint_id"], + ( + { + "configurable": { + "thread_id": value["thread_id"], + "checkpoint_ns": value["checkpoint_ns"], + "checkpoint_id": value["parent_checkpoint_id"], + } } - } - if value["parent_checkpoint_id"] - else None, + if value["parent_checkpoint_id"] + else None + ), await asyncio.to_thread(self._load_writes, value["pending_writes"]), ) @@ -216,15 +206,17 @@ class AsyncPostgresSaver(BasePostgresSaver): value["pending_sends"], ), self._load_metadata(value["metadata"]), - { - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": value["parent_checkpoint_id"], + ( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": value["parent_checkpoint_id"], + } } - } - if value["parent_checkpoint_id"] - else None, + if value["parent_checkpoint_id"] + else None + ), await asyncio.to_thread(self._load_writes, value["pending_writes"]), ) @@ -331,7 +323,7 @@ class AsyncPostgresSaver(BasePostgresSaver): Will be applied regardless of whether the AsyncPostgresSaver instance was initialized with a pipeline. If pipeline mode is not supported, will fall back to using transaction context manager. """ - async with _get_connection(self.conn) as conn: + async with _ainternal.get_connection(self.conn) as conn: if self.pipe: # a connection in pipeline mode can be used concurrently # in multiple threads/coroutines, but only one cursor can be @@ -467,3 +459,6 @@ class AsyncPostgresSaver(BasePostgresSaver): return asyncio.run_coroutine_threadsafe( self.aput_writes(config, writes, task_id), self.loop ).result() + + +__all__ = ["AsyncPostgresSaver", "Conn"] diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index dda7321d0..578523052 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -13,14 +13,17 @@ from typing import ( ) import orjson -from psycopg import AsyncConnection, AsyncCursor +from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities from psycopg.errors import UndefinedTable -from psycopg.rows import dict_row +from psycopg.rows import DictRow, dict_row +from psycopg_pool import AsyncConnectionPool +from langgraph.checkpoint.postgres import _ainternal from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.postgres.base import ( BasePostgresStore, + PoolConfig, Row, _decode_ns_bytes, _group_ops, @@ -30,81 +33,88 @@ from langgraph.store.postgres.base import ( logger = logging.getLogger(__name__) -class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnection]): - __slots__ = ("_deserializer",) +class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]): + __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline") def __init__( self, - conn: AsyncConnection[Any], + conn: _ainternal.Conn, *, + pipe: Optional[AsyncPipeline] = None, deserializer: Optional[ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, ) -> None: + if isinstance(conn, AsyncConnectionPool) and pipe is not None: + raise ValueError( + "Pipeline should be used only with a single AsyncConnection, not AsyncConnectionPool." + ) super().__init__() self._deserializer = deserializer self.conn = conn + self.pipe = pipe + self.lock = asyncio.Lock() self.loop = asyncio.get_running_loop() + self.supports_pipeline = Capabilities().has_pipeline() async def abatch(self, ops: Iterable[Op]) -> list[Result]: grouped_ops, num_ops = _group_ops(ops) results: list[Result] = [None] * num_ops - async with self.conn.pipeline(): - tasks = [] - - if GetOp in grouped_ops: - tasks.append( - self._batch_get_ops( - cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results - ) - ) - - if PutOp in grouped_ops: - tasks.append( - self._batch_put_ops( - cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]) - ) - ) - - if SearchOp in grouped_ops: - tasks.append( - self._batch_search_ops( - cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]), - results, - ) - ) - - if ListNamespacesOp in grouped_ops: - tasks.append( - self._batch_list_namespaces_ops( - cast( - Sequence[tuple[int, ListNamespacesOp]], - grouped_ops[ListNamespacesOp], - ), - results, - ) - ) - - await asyncio.gather(*tasks) + async with _ainternal.get_connection(self.conn) as conn: + if self.pipe: + async with self.pipe: + await self._execute_batch(grouped_ops, results, conn) + else: + await self._execute_batch(grouped_ops, results, conn) return results - def batch(self, ops: Iterable[Op]) -> list[Result]: - return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result() + async def _execute_batch( + self, + grouped_ops: dict, + results: list[Result], + conn: AsyncConnection[DictRow], + ) -> None: + async with self._cursor(conn, pipeline=True) as cur: + if GetOp in grouped_ops: + await self._batch_get_ops( + cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), + results, + cur, + ) + + if SearchOp in grouped_ops: + await self._batch_search_ops( + cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]), + results, + cur, + ) + + if ListNamespacesOp in grouped_ops: + await self._batch_list_namespaces_ops( + cast( + Sequence[tuple[int, ListNamespacesOp]], + grouped_ops[ListNamespacesOp], + ), + results, + cur, + ) + + if PutOp in grouped_ops: + await self._batch_put_ops( + cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), + cur, + ) async def _batch_get_ops( self, get_ops: Sequence[tuple[int, GetOp]], results: list[Result], + cur: AsyncCursor[DictRow], ) -> None: - cursors = [] for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops): - cur = self.conn.cursor(binary=True) await cur.execute(query, params) - cursors.append((cur, namespace, items)) - - for cur, namespace, items in cursors: rows = cast(list[Row], await cur.fetchall()) key_to_row = {row["key"]: row for row in rows} for idx, key in items: @@ -119,26 +129,21 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio async def _batch_put_ops( self, put_ops: Sequence[tuple[int, PutOp]], + cur: AsyncCursor[DictRow], ) -> None: queries = self._get_batch_PUT_queries(put_ops) for query, params in queries: - cur = self.conn.cursor(binary=True) await cur.execute(query, params) async def _batch_search_ops( self, search_ops: Sequence[tuple[int, SearchOp]], results: list[Result], + cur: AsyncCursor[DictRow], ) -> None: queries = self._get_batch_search_queries(search_ops) - cursors: list[tuple[AsyncCursor[Any], int]] = [] - for (query, params), (idx, _) in zip(queries, search_ops): - cur = self.conn.cursor(binary=True) await cur.execute(query, params) - cursors.append((cur, idx)) - - for cur, idx in cursors: rows = cast(list[Row], await cur.fetchall()) items = [ _row_to_item( @@ -152,37 +157,103 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio self, list_ops: Sequence[tuple[int, ListNamespacesOp]], results: list[Result], + cur: AsyncCursor[DictRow], ) -> None: queries = self._get_batch_list_namespaces_queries(list_ops) - cursors: list[tuple[AsyncCursor[Any], int]] = [] for (query, params), (idx, _) in zip(queries, list_ops): - cur = self.conn.cursor(binary=True) await cur.execute(query, params) - cursors.append((cur, idx)) - - for cur, idx in cursors: rows = cast(list[dict], await cur.fetchall()) namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows] results[idx] = namespaces + @asynccontextmanager + async def _cursor( + self, conn: AsyncConnection[DictRow], *, pipeline: bool = False + ) -> AsyncIterator[AsyncCursor[Any]]: + """Create a database cursor as a context manager. + + Args: + conn: The database connection to use + pipeline: whether to use pipeline for the DB operations inside the context manager. + Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline. + If pipeline mode is not supported, will fall back to using transaction context manager. + """ + 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 + async with conn.cursor(binary=True) as cur: + try: + yield cur + finally: + if pipeline: + 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 + if self.supports_pipeline: + async with self.lock, conn.pipeline(), conn.cursor(binary=True) as cur: + yield cur + else: + async with self.lock, conn.transaction(), conn.cursor( + binary=True + ) as cur: + yield cur + else: + async with conn.cursor(binary=True) as cur: + yield cur + + def batch(self, ops: Iterable[Op]) -> list[Result]: + return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result() + @classmethod @asynccontextmanager async def from_conn_string( cls, conn_string: str, + *, + pipeline: bool = False, + pool_config: Optional[PoolConfig] = None, ) -> AsyncIterator["AsyncPostgresStore"]: """Create a new AsyncPostgresStore instance from a connection string. Args: conn_string (str): The Postgres connection info string. + pipeline (bool): Whether to use AsyncPipeline (only for single connections) + pool_config (Optional[PoolConfig]): Configuration for the connection pool. + If provided, will create a connection pool and use it instead of a single connection. + This overrides the `pipeline` argument. Returns: AsyncPostgresStore: A new AsyncPostgresStore instance. """ - async with await AsyncConnection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - yield cls(conn=conn) + if pool_config is not None: + pc = pool_config.copy() + async with cast( + AsyncConnectionPool[AsyncConnection[DictRow]], + AsyncConnectionPool( + conn_string, + min_size=pc.pop("min_size", 1), + max_size=pc.pop("max_size", None), + kwargs={ + "autocommit": True, + "prepare_threshold": 0, + "row_factory": dict_row, + **(pc.pop("kwargs", None) or {}), + }, + **cast(dict, pc), + ), + ) as pool: + yield cls(conn=pool) + else: + 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 cls(conn=conn, pipe=pipe) + else: + yield cls(conn=conn) async def setup(self) -> None: """Set up the store database asynchronously. @@ -191,28 +262,33 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[AsyncConnectio already exist and runs database migrations. It MUST be called directly by the user the first time the store is used. """ - async with self.conn.cursor() as cur: - try: - await cur.execute( - "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1" - ) - row = cast(dict, await cur.fetchone()) - if row is None: - version = -1 - else: - version = row["v"] - except UndefinedTable: - version = -1 - # Create store_migrations table if it doesn't exist - await cur.execute( - """ - CREATE TABLE IF NOT EXISTS store_migrations ( - v INTEGER PRIMARY KEY + async with _ainternal.get_connection(self.conn) as conn: + async with conn.cursor() as cur: + try: + await cur.execute( + "SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1" ) - """ - ) - for v, migration in enumerate( - self.MIGRATIONS[version + 1 :], start=version + 1 - ): - await cur.execute(migration) - await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,)) + row = cast(dict, await cur.fetchone()) + if row is None: + version = -1 + else: + version = row["v"] + except UndefinedTable: + version = -1 + # Create store_migrations table if it doesn't exist + await cur.execute( + """ + CREATE TABLE IF NOT EXISTS store_migrations ( + v INTEGER PRIMARY KEY + ) + """ + ) + for v, migration in enumerate( + self.MIGRATIONS[version + 1 :], start=version + 1 + ): + await cur.execute(migration) + await cur.execute( + "INSERT INTO store_migrations (v) VALUES (%s)", (v,) + ) + if self.pipe: + await self.pipe.sync() diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 8bd2b8279..3bb343b0c 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -1,6 +1,7 @@ import asyncio import json import logging +import threading from collections import defaultdict from contextlib import contextmanager from datetime import datetime @@ -18,12 +19,15 @@ from typing import ( ) import orjson -from psycopg import BaseConnection, Connection, Cursor +from psycopg import Capabilities, 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 typing_extensions import TypedDict +from langgraph.checkpoint.postgres import _ainternal as _ainternal +from langgraph.checkpoint.postgres import _internal as _pg_internal from langgraph.store.base import ( BaseStore, GetOp, @@ -56,7 +60,32 @@ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pa """, ] -C = TypeVar("C", bound=BaseConnection) +C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn]) + + +class PoolConfig(TypedDict, total=False): + """Connection pool settings for PostgreSQL connections. + + Controls connection lifecycle and resource utilization: + - Small pools (1-5) suit low-concurrency workloads + - Larger pools handle concurrent requests but consume more resources + - Setting max_size prevents resource exhaustion under load + """ + + min_size: int + """Minimum number of connections maintained in the pool. Defaults to 1.""" + + max_size: Optional[int] + """Maximum number of connections allowed in the pool. None means unlimited.""" + + kwargs: dict + """Additional connection arguments passed to each connection in the pool. + + Default kwargs set automatically: + - autocommit: True + - prepare_threshold: 0 + - row_factory: dict_row + """ class BasePostgresStore(Generic[C]): @@ -88,9 +117,14 @@ class BasePostgresStore(Generic[C]): self, put_ops: Sequence[tuple[int, PutOp]], ) -> list[tuple[str, Sequence]]: + # Last-write wins + dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {} + for _, op in put_ops: + dedupped_ops[(op.namespace, op.key)] = op + inserts: list[PutOp] = [] deletes: list[PutOp] = [] - for _, op in put_ops: + for op in dedupped_ops.values(): if op.value is None: deletes.append(op) else: @@ -219,13 +253,14 @@ class BasePostgresStore(Generic[C]): return queries -class PostgresStore(BaseStore, BasePostgresStore[Connection]): - __slots__ = ("_deserializer",) +class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): + __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline") def __init__( self, - conn: Connection[Any], + conn: _pg_internal.Conn, *, + pipe: Optional[Pipeline] = None, deserializer: Optional[ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, @@ -233,26 +268,110 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]): super().__init__() self._deserializer = deserializer self.conn = conn + self.pipe = pipe + self.supports_pipeline = Capabilities().has_pipeline() + self.lock = threading.Lock() + + @classmethod + @contextmanager + def from_conn_string( + cls, + conn_string: str, + *, + pipeline: bool = False, + pool_config: Optional[PoolConfig] = None, + ) -> Iterator["PostgresStore"]: + """Create a new PostgresStore instance from a connection string. + + Args: + conn_string (str): The Postgres connection info string. + pipeline (bool): whether to use Pipeline (only for single connections) + pool_config (Optional[PoolArgs]): Configuration for the connection pool. + If provided, will create a connection pool and use it instead of a single connection. + This overrides the `pipeline` argument. + Returns: + PostgresStore: A new PostgresStore instance. + """ + if pool_config is not None: + pc = pool_config.copy() + with cast( + ConnectionPool[Connection[DictRow]], + ConnectionPool( + conn_string, + min_size=pc.pop("min_size", 1), + max_size=pc.pop("max_size", None), + kwargs={ + "autocommit": True, + "prepare_threshold": 0, + "row_factory": dict_row, + **(pc.pop("kwargs", None) or {}), + }, + **cast(dict, pc), + ), + ) as pool: + yield cls(conn=pool) + else: + with Connection.connect( + conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row + ) as conn: + if pipeline: + with conn.pipeline() as pipe: + yield cls(conn, pipe=pipe) + else: + yield cls(conn) + + @contextmanager + def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: + """Create a database cursor as a context manager. + + Args: + pipeline (bool): whether to use pipeline for the DB operations inside the context manager. + Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline. + If pipeline mode is not supported, will fall back to using transaction context manager. + """ + with _pg_internal.get_connection(self.conn) as conn: + 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 conn.cursor(binary=True, row_factory=dict_row) as cur: + yield cur + finally: + if pipeline: + 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 + if self.supports_pipeline: + with self.lock, conn.pipeline(), conn.cursor( + binary=True, row_factory=dict_row + ) as cur: + yield cur + else: + with self.lock, conn.transaction(), conn.cursor( + binary=True, row_factory=dict_row + ) as cur: + yield cur + else: + with conn.cursor(binary=True, row_factory=dict_row) as cur: + yield cur def batch(self, ops: Iterable[Op]) -> list[Result]: grouped_ops, num_ops = _group_ops(ops) results: list[Result] = [None] * num_ops - with self.conn.pipeline(): + with self._cursor(pipeline=True) as cur: if GetOp in grouped_ops: self._batch_get_ops( - cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results - ) - - if PutOp in grouped_ops: - self._batch_put_ops( - cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]) + cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur ) if SearchOp in grouped_ops: self._batch_search_ops( cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]), results, + cur, ) if ListNamespacesOp in grouped_ops: @@ -262,25 +381,23 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]): grouped_ops[ListNamespacesOp], ), results, + cur, + ) + if PutOp in grouped_ops: + self._batch_put_ops( + cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), cur ) return results - async def abatch(self, ops: Iterable[Op]) -> list[Result]: - return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops) - def _batch_get_ops( self, get_ops: Sequence[tuple[int, GetOp]], results: list[Result], + cur: Cursor[DictRow], ) -> None: - cursors = [] for query, params, namespace, items in self._get_batch_GET_ops_queries(get_ops): - cur = self.conn.cursor(binary=True) cur.execute(query, params) - cursors.append((cur, namespace, items)) - - for cur, namespace, items in cursors: rows = cast(list[Row], cur.fetchall()) key_to_row = {row["key"]: row for row in rows} for idx, key in items: @@ -295,70 +412,44 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]): def _batch_put_ops( self, put_ops: Sequence[tuple[int, PutOp]], + cur: Cursor[DictRow], ) -> None: queries = self._get_batch_PUT_queries(put_ops) for query, params in queries: - cur = self.conn.cursor(binary=True) cur.execute(query, params) def _batch_search_ops( self, search_ops: Sequence[tuple[int, SearchOp]], results: list[Result], + cur: Cursor[DictRow], ) -> None: - queries = self._get_batch_search_queries(search_ops) - cursors: list[tuple[Cursor[Any], int]] = [] - - for (query, params), (idx, _) in zip(queries, search_ops): - cur = self.conn.cursor(binary=True) + for (query, params), (idx, _) in zip( + self._get_batch_search_queries(search_ops), search_ops + ): cur.execute(query, params) - cursors.append((cur, idx)) - - for cur, idx in cursors: rows = cast(list[Row], cur.fetchall()) - items = [ + results[idx] = [ _row_to_item( _decode_ns_bytes(row["prefix"]), row, loader=self._deserializer ) for row in rows ] - results[idx] = items def _batch_list_namespaces_ops( self, list_ops: Sequence[tuple[int, ListNamespacesOp]], results: list[Result], + cur: Cursor[DictRow], ) -> None: - queries = self._get_batch_list_namespaces_queries(list_ops) - cursors: list[tuple[Cursor[Any], int]] = [] - for (query, params), (idx, _) in zip(queries, list_ops): - cur = self.conn.cursor(binary=True) + for (query, params), (idx, _) in zip( + self._get_batch_list_namespaces_queries(list_ops), list_ops + ): cur.execute(query, params) - cursors.append((cur, idx)) + results[idx] = [_decode_ns_bytes(row["truncated_prefix"]) for row in cur] - for cur, idx in cursors: - rows = cast(list[dict], cur.fetchall()) - namespaces = [_decode_ns_bytes(row["truncated_prefix"]) for row in rows] - results[idx] = namespaces - - @classmethod - @contextmanager - def from_conn_string( - cls, - conn_string: str, - ) -> Iterator["PostgresStore"]: - """Create a new BasePostgresStore instance from a connection string. - - Args: - conn_string (str): The Postgres connection info string. - - Returns: - BasePostgresStore: A new BasePostgresStore instance. - """ - with Connection.connect( - conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row - ) as conn: - yield cls(conn=conn) + async def abatch(self, ops: Iterable[Op]) -> list[Result]: + return await asyncio.get_running_loop().run_in_executor(None, self.batch, ops) def setup(self) -> None: """Set up the store database. @@ -367,7 +458,7 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]): already exist and runs database migrations. It MUST be called directly by the user the first time the store is used. """ - with self.conn.cursor(binary=True) as cur: + with self._cursor() as cur: try: cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1") row = cast(dict, cur.fetchone()) @@ -376,9 +467,7 @@ class PostgresStore(BaseStore, BasePostgresStore[Connection]): else: version = row["v"] except UndefinedTable: - self.conn.rollback() version = -1 - # Create store_migrations table if it doesn't exist cur.execute( """ CREATE TABLE IF NOT EXISTS store_migrations ( diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index 49061b98c..56d199812 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/libs/checkpoint-postgres/tests/conftest.py @@ -24,6 +24,10 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None: await conn.execute("DELETE FROM checkpoint_blobs") await conn.execute("DELETE FROM checkpoint_writes") await conn.execute("DELETE FROM checkpoint_migrations") - await conn.execute("DELETE FROM store_migrations") + except UndefinedTable: + pass + try: + await conn.execute("DELETE FROM store_migrations") + await conn.execute("DELETE FROM store") except UndefinedTable: pass diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index e7a7b31a1..71aaa4e36 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -1,114 +1,76 @@ # type: ignore +import sys import uuid -from datetime import datetime -from typing import Any -from unittest.mock import AsyncMock, MagicMock +from typing import AsyncIterator import pytest from conftest import DEFAULT_URI # type: ignore +from psycopg import AsyncConnection from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp from langgraph.store.postgres import AsyncPostgresStore -class MockAsyncCursor: - def __init__(self, fetch_result: Any) -> None: - self.fetch_result = fetch_result - self.execute = AsyncMock() - self.fetchall = AsyncMock(return_value=self.fetch_result) +@pytest.fixture(scope="function", params=["default", "pipe", "pool"]) +async def store(request) -> AsyncIterator[AsyncPostgresStore]: + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid.uuid4().hex[:16]}" + uri_parts = DEFAULT_URI.split("/") + uri_base = "/".join(uri_parts[:-1]) + query_params = "" + if "?" in uri_parts[-1]: + db_name, query_params = uri_parts[-1].split("?", 1) + query_params = "?" + query_params -class MockAsyncConnection: - def __init__(self) -> None: - self.cursor = MagicMock() - self.pipeline = MagicMock( - return_value=AsyncMock(__aenter__=AsyncMock(), __aexit__=AsyncMock()) - ) + conn_string = f"{uri_base}/{database}{query_params}" + admin_conn_string = DEFAULT_URI + async with await AsyncConnection.connect( + admin_conn_string, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string(conn_string) as store: + await store.setup() -@pytest.fixture -def mock_connection() -> MockAsyncConnection: - return MockAsyncConnection() - - -@pytest.fixture -async def store(mock_connection: MockAsyncConnection) -> AsyncPostgresStore: - return AsyncPostgresStore(mock_connection) + if request.param == "pipe": + async with AsyncPostgresStore.from_conn_string( + conn_string, pipeline=True + ) as store: + yield store + elif request.param == "pool": + async with AsyncPostgresStore.from_conn_string( + conn_string, pool_config={"min_size": 1, "max_size": 10} + ) as store: + yield store + else: # default + async with AsyncPostgresStore.from_conn_string(conn_string) as store: + yield store + finally: + async with await AsyncConnection.connect( + admin_conn_string, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") async def test_abatch_order(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_get_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_search_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - ] - ) - mock_list_namespaces_cursor = MockAsyncCursor( - [ - {"truncated_prefix": b"\x01test"}, - ] - ) - - failures = [] - - def cursor_side_effect(binary: bool = False) -> Any: - cursor = MagicMock() - - async def execute_side_effect(query: str, *params: Any) -> None: - # My super sophisticated database. - if "SELECT prefix, key," in query: - cursor.fetchall = mock_search_cursor.fetchall - elif "SELECT DISTINCT ON (truncated_prefix)" in query: - cursor.fetchall = mock_list_namespaces_cursor.fetchall - elif "WHERE prefix = %s AND key" in query: - cursor.fetchall = mock_get_cursor.fetchall - elif "INSERT INTO " in query: - pass - else: - e = ValueError(f"Unmatched query: {query}") - failures.append(e) - raise e - - cursor.execute = AsyncMock(side_effect=execute_side_effect) - return cursor - - mock_connection.cursor.side_effect = cursor_side_effect # type: ignore + # Setup test data + await store.aput(("test", "foo"), "key1", {"data": "value1"}) + await store.aput(("test", "bar"), "key2", {"data": "value2"}) ops = [ - GetOp(namespace=("test",), key="key1"), - PutOp(namespace=("test",), key="key2", value={"data": "value2"}), + GetOp(namespace=("test", "foo"), key="key1"), + PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}), SearchOp( namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 ), ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), GetOp(namespace=("test",), key="key3"), ] + results = await store.abatch(ops) - assert not failures assert len(results) == 5 assert isinstance(results[0], Item) assert isinstance(results[0].value, dict) @@ -118,27 +80,29 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None: assert isinstance(results[2], list) assert len(results[2]) == 1 assert isinstance(results[3], list) - assert results[3] == [("test",)] + assert ("test", "foo") in results[3] and ("test", "bar") in results[3] assert results[4] is None ops_reordered = [ SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), - GetOp(namespace=("test",), key="key2"), + GetOp(namespace=("test", "bar"), key="key2"), ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0), PutOp(namespace=("test",), key="key3", value={"data": "value3"}), - GetOp(namespace=("test",), key="key1"), + GetOp(namespace=("test", "foo"), key="key1"), ] results_reordered = await store.abatch(ops_reordered) - assert not failures assert len(results_reordered) == 5 assert isinstance(results_reordered[0], list) - assert len(results_reordered[0]) == 1 + assert len(results_reordered[0]) == 2 assert isinstance(results_reordered[1], Item) assert results_reordered[1].value == {"data": "value2"} assert results_reordered[1].key == "key2" assert isinstance(results_reordered[2], list) - assert results_reordered[2] == [("test",)] + assert ("test", "foo") in results_reordered[2] and ( + "test", + "bar", + ) in results_reordered[2] assert results_reordered[3] is None assert isinstance(results_reordered[4], Item) assert results_reordered[4].value == {"data": "value1"} @@ -146,26 +110,9 @@ async def test_abatch_order(store: AsyncPostgresStore) -> None: async def test_batch_get_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test",), "key1", {"data": "value1"}) + await store.aput(("test",), "key2", {"data": "value2"}) ops = [ GetOp(namespace=("test",), key="key1"), @@ -184,10 +131,6 @@ async def test_batch_get_ops(store: AsyncPostgresStore) -> None: async def test_batch_put_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor([]) - mock_connection.cursor.return_value = mock_cursor - ops = [ PutOp(namespace=("test",), key="key1", value={"data": "value1"}), PutOp(namespace=("test",), key="key2", value={"data": "value2"}), @@ -198,30 +141,16 @@ async def test_batch_put_ops(store: AsyncPostgresStore) -> None: assert len(results) == 3 assert all(result is None for result in results) - assert mock_cursor.execute.call_count == 2 + + # Verify the puts worked + items = await store.asearch(["test"], limit=10) + assert len(items) == 2 # key3 had None value so wasn't stored async def test_batch_search_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test", "foo"), "key1", {"data": "value1"}) + await store.aput(("test", "bar"), "key2", {"data": "value2"}) ops = [ SearchOp( @@ -233,29 +162,23 @@ async def test_batch_search_ops(store: AsyncPostgresStore) -> None: results = await store.abatch(ops) assert len(results) == 2 - assert len(results[0]) == 2 - assert len(results[1]) == 2 + assert len(results[0]) == 1 # Filtered results + assert len(results[1]) == 2 # All results async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockAsyncCursor( - [ - {"truncated_prefix": b"\x01test.namespace1"}, - {"truncated_prefix": b"\x01test.namespace2"}, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + await store.aput(("test", "namespace1"), "key1", {"data": "value1"}) + await store.aput(("test", "namespace2"), "key2", {"data": "value2"}) ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)] results = await store.abatch(ops) assert len(results) == 1 - assert results[0] == [("test", "namespace1"), ("test", "namespace2")] - - -# The following use the actual DB connection + assert len(results[0]) == 2 + assert ("test", "namespace1") in results[0] + assert ("test", "namespace2") in results[0] class TestAsyncPostgresStore: diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index add9fb1c5..645c37e9f 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -1,174 +1,118 @@ # type: ignore -import uuid -from datetime import datetime -from typing import Any -from unittest.mock import MagicMock + +from uuid import uuid4 import pytest from conftest import DEFAULT_URI # type: ignore +from psycopg import Connection -from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp +from langgraph.store.base import ( + GetOp, + Item, + ListNamespacesOp, + MatchCondition, + PutOp, + SearchOp, +) from langgraph.store.postgres import PostgresStore -class MockCursor: - def __init__(self, fetch_result: Any) -> None: - self.fetch_result = fetch_result - self.execute = MagicMock() - self.fetchall = MagicMock(return_value=self.fetch_result) +@pytest.fixture(scope="function", params=["default", "pipe", "pool"]) +def store(request) -> PostgresStore: + database = f"test_{uuid4().hex[:16]}" + uri_parts = DEFAULT_URI.split("/") + uri_base = "/".join(uri_parts[:-1]) + query_params = "" + if "?" in uri_parts[-1]: + db_name, query_params = uri_parts[-1].split("?", 1) + query_params = "?" + query_params + conn_string = f"{uri_base}/{database}{query_params}" + admin_conn_string = DEFAULT_URI -class MockConnection: - def __init__(self) -> None: - self.cursor = MagicMock() - self.pipeline = MagicMock() + with Connection.connect(admin_conn_string, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + with PostgresStore.from_conn_string(conn_string) as store: + store.setup() - -@pytest.fixture -def mock_connection() -> MockConnection: - return MockConnection() - - -@pytest.fixture -def store(mock_connection: MockConnection) -> PostgresStore: - return PostgresStore(mock_connection) + if request.param == "pipe": + with PostgresStore.from_conn_string(conn_string, pipeline=True) as store: + yield store + elif request.param == "pool": + with PostgresStore.from_conn_string( + conn_string, pool_config={"min_size": 1, "max_size": 10} + ) as store: + yield store + else: # default + with PostgresStore.from_conn_string(conn_string) as store: + yield store + finally: + with Connection.connect(admin_conn_string, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") def test_batch_order(store: PostgresStore) -> None: - mock_connection = store.conn - mock_get_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_search_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - ] - ) - mock_list_namespaces_cursor = MockCursor( - [ - {"truncated_prefix": b"\x01test"}, - ] - ) - - failures = [] - - def cursor_side_effect(binary: bool = False) -> Any: - cursor = MagicMock() - - def execute_side_effect(query: str, *params: Any) -> None: - # My super sophisticated database. - if "SELECT prefix, key, value" in query: - cursor.fetchall = mock_search_cursor.fetchall - elif "SELECT DISTINCT ON (truncated_prefix)" in query: - cursor.fetchall = mock_list_namespaces_cursor.fetchall - elif "WHERE prefix = %s AND key" in query: - cursor.fetchall = mock_get_cursor.fetchall - elif "INSERT INTO " in query: - pass - else: - e = ValueError(f"Unmatched query: {query}") - failures.append(e) - raise e - - cursor.execute = MagicMock(side_effect=execute_side_effect) - return cursor - - mock_connection.cursor.side_effect = cursor_side_effect + # Setup test data + store.put(("test", "foo"), "key1", {"data": "value1"}) + store.put(("test", "bar"), "key2", {"data": "value2"}) ops = [ - GetOp(namespace=("test",), key="key1"), - PutOp(namespace=("test",), key="key2", value={"data": "value2"}), + GetOp(namespace=("test", "foo"), key="key1"), + PutOp(namespace=("test", "bar"), key="key2", value={"data": "value2"}), SearchOp( namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 ), ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), GetOp(namespace=("test",), key="key3"), ] + results = store.batch(ops) - assert not failures assert len(results) == 5 assert isinstance(results[0], Item) assert isinstance(results[0].value, dict) assert results[0].value == {"data": "value1"} assert results[0].key == "key1" - assert results[1] is None + assert results[1] is None # Put operation returns None assert isinstance(results[2], list) assert len(results[2]) == 1 assert isinstance(results[3], list) - assert results[3] == [("test",)] - assert results[4] is None + assert len(results[3]) > 0 # Should contain at least our test namespaces + assert results[4] is None # Non-existent key returns None + # Test reordered operations ops_reordered = [ SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), - GetOp(namespace=("test",), key="key2"), + GetOp(namespace=("test", "bar"), key="key2"), ListNamespacesOp(match_conditions=None, max_depth=None, limit=5, offset=0), PutOp(namespace=("test",), key="key3", value={"data": "value3"}), - GetOp(namespace=("test",), key="key1"), + GetOp(namespace=("test", "foo"), key="key1"), ] results_reordered = store.batch(ops_reordered) - assert not failures assert len(results_reordered) == 5 assert isinstance(results_reordered[0], list) - assert len(results_reordered[0]) == 1 + assert len(results_reordered[0]) >= 2 # Should find at least our two test items assert isinstance(results_reordered[1], Item) assert results_reordered[1].value == {"data": "value2"} assert results_reordered[1].key == "key2" assert isinstance(results_reordered[2], list) - assert results_reordered[2] == [("test",)] - assert results_reordered[3] is None + assert len(results_reordered[2]) > 0 + assert results_reordered[3] is None # Put operation returns None assert isinstance(results_reordered[4], Item) assert results_reordered[4].value == {"data": "value1"} assert results_reordered[4].key == "key1" def test_batch_get_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + store.put(("test",), "key1", {"data": "value1"}) + store.put(("test",), "key2", {"data": "value2"}) ops = [ GetOp(namespace=("test",), key="key1"), GetOp(namespace=("test",), key="key2"), - GetOp(namespace=("test",), key="key3"), + GetOp(namespace=("test",), key="key3"), # Non-existent key ] results = store.batch(ops) @@ -182,75 +126,90 @@ def test_batch_get_ops(store: PostgresStore) -> None: def test_batch_put_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor([]) - mock_connection.cursor.return_value = mock_cursor - ops = [ PutOp(namespace=("test",), key="key1", value={"data": "value1"}), PutOp(namespace=("test",), key="key2", value={"data": "value2"}), - PutOp(namespace=("test",), key="key3", value=None), + PutOp(namespace=("test",), key="key3", value=None), # Delete operation ] results = store.batch(ops) - assert len(results) == 3 assert all(result is None for result in results) - assert mock_cursor.execute.call_count == 2 + + # Verify the puts worked + item1 = store.get(("test",), "key1") + item2 = store.get(("test",), "key2") + item3 = store.get(("test",), "key3") + + assert item1 and item1.value == {"data": "value1"} + assert item2 and item2.value == {"data": "value2"} + assert item3 is None def test_batch_search_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - { - "key": "key1", - "value": '{"data": "value1"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.foo", - }, - { - "key": "key2", - "value": '{"data": "value2"}', - "created_at": datetime.now(), - "updated_at": datetime.now(), - "prefix": "test.bar", - }, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data + test_data = [ + (("test", "foo"), "key1", {"data": "value1", "tag": "a"}), + (("test", "bar"), "key2", {"data": "value2", "tag": "a"}), + (("test", "baz"), "key3", {"data": "value3", "tag": "b"}), + ] + for namespace, key, value in test_data: + store.put(namespace, key, value) ops = [ - SearchOp( - namespace_prefix=("test",), filter={"data": "value1"}, limit=10, offset=0 - ), - SearchOp(namespace_prefix=("test",), filter=None, limit=5, offset=0), + SearchOp(namespace_prefix=("test",), filter={"tag": "a"}, limit=10, offset=0), + SearchOp(namespace_prefix=("test",), filter=None, limit=2, offset=0), + SearchOp(namespace_prefix=("test", "foo"), filter=None, limit=10, offset=0), ] results = store.batch(ops) + assert len(results) == 3 - assert len(results) == 2 + # First search should find items with tag "a" assert len(results[0]) == 2 + assert all(item.value["tag"] == "a" for item in results[0]) + + # Second search should return first 2 items assert len(results[1]) == 2 + # Third search should only find items in test/foo namespace + assert len(results[2]) == 1 + assert results[2][0].namespace == ("test", "foo") + def test_batch_list_namespaces_ops(store: PostgresStore) -> None: - mock_connection = store.conn - mock_cursor = MockCursor( - [ - {"truncated_prefix": b"\x01test.namespace1"}, - {"truncated_prefix": b"\x01test.namespace2"}, - ] - ) - mock_connection.cursor.return_value = mock_cursor + # Setup test data with various namespaces + test_data = [ + (("test", "documents", "public"), "doc1", {"content": "public doc"}), + (("test", "documents", "private"), "doc2", {"content": "private doc"}), + (("test", "images", "public"), "img1", {"content": "public image"}), + (("prod", "documents", "public"), "doc3", {"content": "prod doc"}), + ] + for namespace, key, value in test_data: + store.put(namespace, key, value) - ops = [ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0)] + ops = [ + ListNamespacesOp(match_conditions=None, max_depth=None, limit=10, offset=0), + ListNamespacesOp(match_conditions=None, max_depth=2, limit=10, offset=0), + ListNamespacesOp( + match_conditions=[MatchCondition("suffix", "public")], + max_depth=None, + limit=10, + offset=0, + ), + ] results = store.batch(ops) + assert len(results) == 3 - assert len(results) == 1 - assert results[0] == [("test", "namespace1"), ("test", "namespace2")] + # First operation should list all namespaces + assert len(results[0]) == len(test_data) + + # Second operation should only return namespaces up to depth 2 + assert all(len(ns) <= 2 for ns in results[1]) + + # Third operation should only return namespaces ending with "public" + assert all(ns[-1] == "public" for ns in results[2]) class TestPostgresStore: @@ -273,195 +232,111 @@ class TestPostgresStore: assert item.key == item_id assert item.value == item_value - updated_value = { - "title": "Updated Test Document", - "content": "Hello, LangGraph!", - } + # Test update + updated_value = {"title": "Updated Document", "content": "Hello, Updated!"} store.put(namespace, item_id, updated_value) updated_item = store.get(namespace, item_id) assert updated_item.value == updated_value assert updated_item.updated_at > item.updated_at + + # Test get from non-existent namespace different_namespace = ("test", "other_documents") item_in_different_namespace = store.get(different_namespace, item_id) assert item_in_different_namespace is None - new_item_id = "doc2" - new_item_value = {"title": "Another Document", "content": "Greetings!"} - store.put(namespace, new_item_id, new_item_value) - - search_results = store.search(["test"], limit=10) - items = search_results - assert len(items) == 2 - assert any(item.key == item_id for item in items) - assert any(item.key == new_item_id for item in items) - - namespaces = store.list_namespaces(prefix=["test"]) - assert ("test", "documents") in namespaces - + # Test delete store.delete(namespace, item_id) - store.delete(namespace, new_item_id) deleted_item = store.get(namespace, item_id) assert deleted_item is None - deleted_item = store.get(namespace, new_item_id) - assert deleted_item is None - - empty_search_results = store.search(["test"], limit=10) - assert len(empty_search_results) == 0 - def test_list_namespaces(self) -> None: with PostgresStore.from_conn_string(DEFAULT_URI) as store: - test_pref = str(uuid.uuid4()) + # Create test data with various namespaces test_namespaces = [ - (test_pref, "test", "documents", "public", test_pref), - (test_pref, "test", "documents", "private", test_pref), - (test_pref, "test", "images", "public", test_pref), - (test_pref, "test", "images", "private", test_pref), - (test_pref, "prod", "documents", "public", test_pref), - ( - test_pref, - "prod", - "documents", - "some", - "nesting", - "public", - test_pref, - ), - (test_pref, "prod", "documents", "private", test_pref), + ("test", "documents", "public"), + ("test", "documents", "private"), + ("test", "images", "public"), + ("test", "images", "private"), + ("prod", "documents", "public"), + ("prod", "documents", "private"), ] + # Insert test data for namespace in test_namespaces: store.put(namespace, "dummy", {"content": "dummy"}) - prefix_result = store.list_namespaces(prefix=[test_pref, "test"]) - assert len(prefix_result) == 4 - assert all([ns[1] == "test" for ns in prefix_result]) + # Test listing with various filters + all_namespaces = store.list_namespaces() + assert len(all_namespaces) == len(test_namespaces) - specific_prefix_result = store.list_namespaces( - prefix=[test_pref, "test", "documents"] - ) - assert len(specific_prefix_result) == 2 - assert all( - [ns[1:3] == ("test", "documents") for ns in specific_prefix_result] - ) + # Test prefix filtering + test_prefix_namespaces = store.list_namespaces(prefix=["test"]) + assert len(test_prefix_namespaces) == 4 + assert all(ns[0] == "test" for ns in test_prefix_namespaces) - suffix_result = store.list_namespaces(suffix=["public", test_pref]) - assert len(suffix_result) == 4 - assert all(ns[-2] == "public" for ns in suffix_result) + # Test suffix filtering + public_namespaces = store.list_namespaces(suffix=["public"]) + assert len(public_namespaces) == 3 + assert all(ns[-1] == "public" for ns in public_namespaces) - prefix_suffix_result = store.list_namespaces( - prefix=[test_pref, "test"], suffix=["public", test_pref] - ) - assert len(prefix_suffix_result) == 2 - assert all( - ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result - ) + # Test max depth + depth_2_namespaces = store.list_namespaces(max_depth=2) + assert all(len(ns) <= 2 for ns in depth_2_namespaces) - wildcard_prefix_result = store.list_namespaces( - prefix=[test_pref, "*", "documents"] - ) - assert len(wildcard_prefix_result) == 5 - assert all(ns[2] == "documents" for ns in wildcard_prefix_result) - - wildcard_suffix_result = store.list_namespaces( - suffix=["*", "public", test_pref] - ) - assert len(wildcard_suffix_result) == 4 - assert all(ns[-2] == "public" for ns in wildcard_suffix_result) - wildcard_single = store.list_namespaces( - suffix=["some", "*", "public", test_pref] - ) - assert len(wildcard_single) == 1 - assert wildcard_single[0] == ( - test_pref, - "prod", - "documents", - "some", - "nesting", - "public", - test_pref, - ) - - max_depth_result = store.list_namespaces(max_depth=3) - assert all([len(ns) <= 3 for ns in max_depth_result]) - - max_depth_result = store.list_namespaces( - max_depth=4, prefix=[test_pref, "*", "documents"] - ) - assert ( - len(set(tuple(res) for res in max_depth_result)) - == len(max_depth_result) - == 5 - ) - - limit_result = store.list_namespaces(prefix=[test_pref], limit=3) - assert len(limit_result) == 3 - - offset_result = store.list_namespaces(prefix=[test_pref], offset=3) - assert len(offset_result) == len(test_namespaces) - 3 - - empty_prefix_result = store.list_namespaces(prefix=[test_pref]) - assert len(empty_prefix_result) == len(test_namespaces) - assert set(tuple(ns) for ns in empty_prefix_result) == set( - tuple(ns) for ns in test_namespaces - ) + # Test pagination + paginated_namespaces = store.list_namespaces(limit=3) + assert len(paginated_namespaces) == 3 + # Cleanup for namespace in test_namespaces: store.delete(namespace, "dummy") - def test_search(self): + def test_search(self) -> None: with PostgresStore.from_conn_string(DEFAULT_URI) as store: - test_namespaces = [ - ("test_search", "documents", "user1"), - ("test_search", "documents", "user2"), - ("test_search", "reports", "department1"), - ("test_search", "reports", "department2"), - ] - test_items = [ - {"title": "Doc 1", "author": "John Doe", "tags": ["important"]}, - {"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]}, - {"title": "Report A", "author": "John Doe", "tags": ["final"]}, - {"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]}, + # Create test data + test_data = [ + ( + ("test", "docs"), + "doc1", + {"title": "First Doc", "author": "Alice", "tags": ["important"]}, + ), + ( + ("test", "docs"), + "doc2", + {"title": "Second Doc", "author": "Bob", "tags": ["draft"]}, + ), + ( + ("test", "images"), + "img1", + {"title": "Image 1", "author": "Alice", "tags": ["final"]}, + ), ] - for namespace, item in zip(test_namespaces, test_items): - store.put(namespace, f"item_{namespace[-1]}", item) + for namespace, key, value in test_data: + store.put(namespace, key, value) - docs_result = store.search(["test_search", "documents"]) - assert len(docs_result) == 2 - assert all( - [item.namespace[1] == "documents" for item in docs_result] - ), docs_result + # Test basic search + all_items = store.search(["test"]) + assert len(all_items) == 3 - reports_result = store.search(["test_search", "reports"]) - assert len(reports_result) == 2 - assert all(item.namespace[1] == "reports" for item in reports_result) + # Test namespace filtering + docs_items = store.search(["test", "docs"]) + assert len(docs_items) == 2 + assert all(item.namespace == ("test", "docs") for item in docs_items) - limited_result = store.search(["test_search"], limit=2) - assert len(limited_result) == 2 - offset_result = store.search(["test_search"]) - assert len(offset_result) == 4 + # Test value filtering + alice_items = store.search(["test"], filter={"author": "Alice"}) + assert len(alice_items) == 2 + assert all(item.value["author"] == "Alice" for item in alice_items) - offset_result = store.search(["test_search"], offset=2) - assert len(offset_result) == 2 - assert all(item not in limited_result for item in offset_result) + # Test pagination + paginated_items = store.search(["test"], limit=2) + assert len(paginated_items) == 2 - john_doe_result = store.search( - ["test_search"], filter={"author": "John Doe"} - ) - assert len(john_doe_result) == 2 - assert all(item.value["author"] == "John Doe" for item in john_doe_result) + offset_items = store.search(["test"], offset=2) + assert len(offset_items) == 1 - draft_result = store.search(["test_search"], filter={"tags": ["draft"]}) - assert len(draft_result) == 2 - assert all("draft" in item.value["tags"] for item in draft_result) - - page1 = store.search(["test_search"], limit=2, offset=0) - page2 = store.search(["test_search"], limit=2, offset=2) - all_items = page1 + page2 - assert len(all_items) == 4 - assert len(set(item.key for item in all_items)) == 4 - - for namespace in test_namespaces: - store.delete(namespace, f"item_{namespace[-1]}") + # Cleanup + for namespace, key, _ in test_data: + store.delete(namespace, key) diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 2aacf6db8..43d0c7afe 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -48,8 +48,13 @@ test: make stop-postgres; \ exit $$EXIT_CODE +WORKERS ?= auto +XDIST_ARGS := $(if $(WORKERS),-n $(WORKERS) --dist worksteal,) +MAXFAIL ?= +MAXFAIL_ARGS := $(if $(MAXFAIL),--maxfail $(MAXFAIL),) + test_watch: - make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv -x $(XDIST_ARGS) $(MAXFAIL_ARGS) --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index eae7694ff..0381206e3 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -272,6 +272,54 @@ async def _store_postgres_aio(): await conn.execute(f"DROP DATABASE {database}") +@asynccontextmanager +async def _store_postgres_aio_pipe(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database + ) as store: + await store.setup() # Run in its own transaction + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + +@asynccontextmanager +async def _store_postgres_aio_pool(): + if sys.version_info < (3, 10): + pytest.skip("Async Postgres tests require Python 3.10+") + database = f"test_{uuid4().hex[:16]}" + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"CREATE DATABASE {database}") + try: + async with AsyncPostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, + pool_config={"max_size": 10}, + ) as store: + await store.setup() + yield store + finally: + async with await AsyncConnection.connect( + DEFAULT_POSTGRES_URI, autocommit=True + ) as conn: + await conn.execute(f"DROP DATABASE {database}") + + @asynccontextmanager async def _store_duckdb_aio(): async with AsyncDuckDBStore.from_conn_string(":memory:") as store: @@ -296,6 +344,45 @@ def store_postgres(): conn.execute(f"DROP DATABASE {database}") +@pytest.fixture(scope="function") +def store_postgres_pipe(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string(DEFAULT_POSTGRES_URI + database) as store: + store.setup() # Run in its own transaction + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pipeline=True + ) as store: + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + +@pytest.fixture(scope="function") +def store_postgres_pool(): + database = f"test_{uuid4().hex[:16]}" + # create unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"CREATE DATABASE {database}") + try: + # yield store + with PostgresStore.from_conn_string( + DEFAULT_POSTGRES_URI + database, pool_config={"max_size": 10} + ) as store: + store.setup() + yield store + finally: + # drop unique db + with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn: + conn.execute(f"DROP DATABASE {database}") + + @pytest.fixture(scope="function") def store_duckdb(): with DuckDBStore.from_conn_string(":memory:") as store: @@ -317,6 +404,12 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: elif store_name == "postgres_aio": async with _store_postgres_aio() as store: yield store + elif store_name == "postgres_aio_pipe": + async with _store_postgres_aio_pipe() as store: + yield store + elif store_name == "postgres_aio_pool": + async with _store_postgres_aio_pool() as store: + yield store elif store_name == "duckdb_aio": async with _store_duckdb_aio() as store: yield store @@ -342,5 +435,17 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ *ALL_CHECKPOINTERS_ASYNC, None, ] -ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"] -ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"] +ALL_STORES_SYNC = [ + "in_memory", + "postgres", + "postgres_pipe", + "postgres_pool", + "duckdb", +] +ALL_STORES_ASYNC = [ + "in_memory", + "postgres_aio", + "postgres_aio_pipe", + "postgres_aio_pool", + "duckdb_aio", +] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fda881b15..fcee1fa20 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,5 +1,6 @@ import enum import json +import logging import operator import re import time @@ -67,16 +68,9 @@ from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInt from langgraph.graph import END, Graph, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue -from langgraph.prebuilt.chat_agent_executor import ( - create_tool_calling_executor, -) +from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import ( - Channel, - GraphRecursionError, - Pregel, - StateSnapshot, -) +from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore @@ -104,6 +98,8 @@ from tests.messages import ( _AnyIdToolMessage, ) +logger = logging.getLogger(__name__) + # define these objects to avoid importing langchain_core.agents # and therefore avoid relying on core Pydantic version @@ -6628,11 +6624,7 @@ def test_message_graph( from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) - from langchain_core.messages import ( - AIMessage, - BaseMessage, - HumanMessage, - ) + from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool @@ -13937,50 +13929,75 @@ def test_store_injected( doc_id = str(uuid.uuid4()) doc = {"some-key": "this-is-a-val"} - - def node(input: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - store.put( - ("foo", "bar"), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": input["count"], - }, - ) - return {"count": 1} - - builder = StateGraph(State) - builder.add_node("node", node) - builder.add_edge("__start__", "node") - graph = builder.compile(store=the_store, checkpointer=checkpointer) - + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") thread_1 = str(uuid.uuid4()) - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) - assert result == {"count": 1} - returned_doc = the_store.get(("foo", "bar"), doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(the_store.search(("foo", "bar"))) == 1 - - # Check update on existing thread - result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) - assert result == {"count": 2} - returned_doc = the_store.get(("foo", "bar"), doc_id).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1} - assert len(the_store.search(("foo", "bar"))) == 1 - thread_2 = str(uuid.uuid4()) + class Node: + def __init__(self, i: Optional[int] = None): + self.i = i + + def __call__(self, inputs: State, config: RunnableConfig, store: BaseStore): + assert isinstance(store, BaseStore) + store.put( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar"), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} + + builder = StateGraph(State) + builder.add_node("node", Node()) + builder.add_edge("__start__", "node") + N = 500 + M = 1 + if "duckdb" in store_name: + logger.warning( + "DuckDB store implementation has a known issue that does not" + " support concurrent writes, so we're reducing the test scope" + ) + N = M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + + graph = builder.compile(store=the_store, checkpointer=checkpointer) + + results = graph.batch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], + ) + result = results[-1] + assert result == {"count": N + 1} + returned_doc = the_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} + assert len(the_store.search(namespace)) == 1 + # Check results after another turn of the same thread + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_1}}) + assert result == {"count": (N + 1) * 2} + returned_doc = the_store.get(namespace, doc_id).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len(the_store.search(namespace)) == 1 + result = graph.invoke({"count": 0}, {"configurable": {"thread_id": thread_2}}) - assert result == {"count": 1} - returned_doc = the_store.get(("foo", "bar"), doc_id).value + assert result == {"count": N + 1} + returned_doc = the_store.get(namespace, doc_id).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc - assert len(the_store.search(("foo", "bar"))) == 1 # still overwriting the same one + assert len(the_store.search(namespace)) == 1 # still overwriting the same one def test_enum_node_names(): diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a31e444e1..c5812e896 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1,4 +1,5 @@ import asyncio +import logging import operator import random import re @@ -100,6 +101,8 @@ from tests.messages import ( _AnyIdToolMessage, ) +logger = logging.getLogger(__name__) + pytestmark = pytest.mark.anyio @@ -12272,60 +12275,89 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> doc_id = str(uuid.uuid4()) doc = {"some-key": "this-is-a-val"} + uid = uuid.uuid4().hex + namespace = (f"foo-{uid}", "bar") + thread_1 = str(uuid.uuid4()) + thread_2 = str(uuid.uuid4()) - async def node(input: State, config: RunnableConfig, store: BaseStore): - assert isinstance(store, BaseStore) - await store.aput( - ("foo", "bar"), - doc_id, - { - **doc, - "from_thread": config["configurable"]["thread_id"], - "some_val": input["count"], - }, - ) - return {"count": 1} + class Node: + def __init__(self, i: Optional[int] = None): + self.i = i + + async def __call__( + self, inputs: State, config: RunnableConfig, store: BaseStore + ): + assert isinstance(store, BaseStore) + await store.aput( + namespace + if self.i is not None + and config["configurable"]["thread_id"] in (thread_1, thread_2) + else (f"foo_{self.i}", "bar"), + doc_id, + { + **doc, + "from_thread": config["configurable"]["thread_id"], + "some_val": inputs["count"], + }, + ) + return {"count": 1} builder = StateGraph(State) - builder.add_node("node", node) + builder.add_node("node", Node()) builder.add_edge("__start__", "node") + + N = 500 + M = 1 + if "duckdb" in store_name: + logger.warning( + "DuckDB store implementation has a known issue that does not" + " support concurrent writes, so we're reducing the test scope" + ) + N = M = 1 + + for i in range(N): + builder.add_node(f"node_{i}", Node(i)) + builder.add_edge("__start__", f"node_{i}") + async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store( store_name ) as the_store: graph = builder.compile(store=the_store, checkpointer=checkpointer) - thread_1 = str(uuid.uuid4()) - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_1}} + # Test batch operations with multiple threads + results = await graph.abatch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], ) - assert result == {"count": 1} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value + result = results[-1] + assert result == {"count": N + 1} + returned_doc = (await the_store.aget(namespace, doc_id)).value assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len((await the_store.asearch(("foo", "bar")))) == 1 + assert len((await the_store.asearch(namespace))) == 1 - # Check update on existing thread + # Check results after another turn of the same thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_1}} ) - assert result == {"count": 2} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 1} - assert len((await the_store.asearch(("foo", "bar")))) == 1 - - thread_2 = str(uuid.uuid4()) + assert result == {"count": (N + 1) * 2} + returned_doc = (await the_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len((await the_store.asearch(namespace))) == 1 + # Test with a different thread result = await graph.ainvoke( {"count": 0}, {"configurable": {"thread_id": thread_2}} ) - assert result == {"count": 1} - returned_doc = (await the_store.aget(("foo", "bar"), doc_id)).value + assert result == {"count": N + 1} + returned_doc = (await the_store.aget(namespace, doc_id)).value assert returned_doc == { **doc, "from_thread": thread_2, "some_val": 0, } # Overwrites the whole doc assert ( - len((await the_store.asearch(("foo", "bar")))) == 1 + len((await the_store.asearch(namespace))) == 1 ) # still overwriting the same one