Wfh/scaling

libs/checkpoint-postgres/langgraph/store/postgres/aio.py

@@ -1,4 +1,4 @@
 import asyncio
 import logging
 from collections.abc import AsyncIterator, Iterable, Sequence
 from contextlib import asynccontextmanager
@@ -19,7 +19,9 @@
     Result,
     SearchOp,
 )
-from langgraph.store.base.batch import AsyncBatchedBaseStore
+from langgraph.store.base.batch import (
+    BatchedBaseStore,
+)
 from langgraph.store.postgres.base import (
     _PLACEHOLDER,
     BasePostgresStore,
@@ -36,7 +38,7 @@
 logger = logging.getLogger(__name__)
 
 
-class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
+class AsyncPostgresStore(BatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
     """Asynchronous Postgres-backed store with optional vector search using pgvector.
 
     !!! example "Examples"

libs/checkpoint-postgres/langgraph/store/postgres/base.py

@@ -30,7 +30,6 @@
 from langgraph.checkpoint.postgres import _ainternal as _ainternal
 from langgraph.checkpoint.postgres import _internal as _pg_internal
 from langgraph.store.base import (
-    BaseStore,
     GetOp,
     IndexConfig,
     Item,
@@ -44,6 +43,7 @@
     get_text_at_path,
     tokenize_path,
 )
+from langgraph.store.base.batch import SyncBatchedBaseStore
 
 if TYPE_CHECKING:
     from langchain_core.embeddings import Embeddings
@@ -533,7 +533,7 @@ def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, lis
             raise ValueError(f"Unsupported operator: {op}")
 
 
-class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
+class PostgresStore(SyncBatchedBaseStore, BasePostgresStore[_pg_internal.Conn]):
     """Postgres-backed store with optional vector search using pgvector.
 
     !!! example "Examples"

libs/checkpoint-postgres/tests/test_async_store.py

@@ -1,8 +1,10 @@
 # type: ignore
+import asyncio
 import itertools
 import sys
 import uuid
 from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from typing import Any, Optional
 
@@ -63,6 +65,96 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
             await conn.execute(f"DROP DATABASE {database}")
 
 
+def test_large_batches(store: AsyncPostgresStore) -> None:
+    N = 1000
+    M = 10
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for m in range(M):
+            for i in range(N):
+                _ = [
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.get,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                    ),
+                    executor.submit(
+                        store.list_namespaces,
+                        prefix=None,
+                        max_depth=m + 1,
+                    ),
+                    executor.submit(
+                        store.search,
+                        ("test",),
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        None,
+                    ),
+                ]
+
+
+async def test_large_batches_async(store: AsyncPostgresStore) -> None:
+    N = 1000
+    M = 10
+    coros = []
+    for m in range(M):
+        for i in range(N):
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.aget(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+            coros.append(
+                store.alist_namespaces(
+                    prefix=None,
+                    max_depth=m + 1,
+                )
+            )
+            coros.append(
+                store.asearch(
+                    ("test",),
+                )
+            )
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.adelete(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+
+    await asyncio.gather(*coros)
+
+
 async def test_abatch_order(store: AsyncPostgresStore) -> None:
     # Setup test data
     await store.aput(("test", "foo"), "key1", {"data": "value1"})

libs/checkpoint-postgres/tests/test_store.py

@@ -1,5 +1,7 @@
 # type: ignore
 
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from typing import Any, Optional
 from uuid import uuid4
@@ -17,11 +19,7 @@
     SearchOp,
 )
 from langgraph.store.postgres import PostgresStore
-from tests.conftest import (
-    DEFAULT_URI,
-    VECTOR_TYPES,
-    CharacterEmbeddings,
-)
+from tests.conftest import DEFAULT_URI, VECTOR_TYPES, CharacterEmbeddings
 
 
 @pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -59,6 +57,96 @@ def store(request) -> PostgresStore:
             conn.execute(f"DROP DATABASE {database}")
 
 
+def test_large_batches(store: PostgresStore) -> None:
+    N = 1000
+    M = 10
+
+    with ThreadPoolExecutor(max_workers=10) as executor:
+        for m in range(M):
+            for i in range(N):
+                _ = [
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.get,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                    ),
+                    executor.submit(
+                        store.list_namespaces,
+                        prefix=None,
+                        max_depth=m + 1,
+                    ),
+                    executor.submit(
+                        store.search,
+                        ("test",),
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        value={"foo": "bar" + str(i)},
+                    ),
+                    executor.submit(
+                        store.put,
+                        ("test", "foo", "bar", "baz", str(m % 2)),
+                        f"key{i}",
+                        None,
+                    ),
+                ]
+
+
+async def test_large_batches_async(store: PostgresStore) -> None:
+    N = 1000
+    M = 10
+    coros = []
+    for m in range(M):
+        for i in range(N):
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.aget(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+            coros.append(
+                store.alist_namespaces(
+                    prefix=None,
+                    max_depth=m + 1,
+                )
+            )
+            coros.append(
+                store.asearch(
+                    ("test",),
+                )
+            )
+            coros.append(
+                store.aput(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                    value={"foo": "bar" + str(i)},
+                )
+            )
+            coros.append(
+                store.adelete(
+                    ("test", "foo", "bar", "baz", str(m % 2)),
+                    f"key{i}",
+                )
+            )
+
+    await asyncio.gather(*coros)
+
+
 def test_batch_order(store: PostgresStore) -> None:
     # Setup test data
     store.put(("test", "foo"), "key1", {"data": "value1"})

libs/checkpoint/langgraph/store/base/__init__.py

@@ -808,6 +808,8 @@ def list_namespaces(
             # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
             ```
         """
+        if max_depth is not None and max_depth <= 0:
+            raise ValueError("If provided, max_depth must be greater than 0")
         match_conditions = []
         if prefix:
             match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
@@ -1004,6 +1006,8 @@ async def alist_namespaces(
             # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
             ```
         """
+        if max_depth is not None and max_depth <= 0:
+            raise ValueError("If provided, max_depth must be greater than 0")
         match_conditions = []
         if prefix:
             match_conditions.append(MatchCondition(match_type="prefix", path=prefix))