Split out async batch to sep file

libs/langgraph/langgraph/graph/state.py

@@ -381,9 +381,9 @@ class StateGraph(Graph):
 
     def compile(
         self,
+        checkpointer: Optional[BaseCheckpointSaver] = None,
         *,
         kv: Optional[BaseKV] = None,
-        checkpointer: Optional[BaseCheckpointSaver] = None,
         interrupt_before: Optional[Union[All, Sequence[str]]] = None,
         interrupt_after: Optional[Union[All, Sequence[str]]] = None,
         debug: bool = False,

libs/langgraph/langgraph/kv/base.py

@@ -1,5 +1,4 @@
-import asyncio
-from typing import Any, List, NamedTuple, Optional, Union
+from typing import Any, List, Optional
 
 V = dict[str, Any]
 
@@ -30,84 +29,3 @@ class BaseKV:
     async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
         # list[(namespace, key, value | none)] -> None
         raise NotImplementedError
-
-
-class GetOp(NamedTuple):
-    pairs: List[tuple[str, str]]
-
-
-class ListOp(NamedTuple):
-    prefixes: List[str]
-
-
-class PutOp(NamedTuple):
-    writes: List[tuple[str, str, Optional[V]]]
-
-
-class KeyValueStore(BaseKV):
-    def __init__(self, kv: BaseKV) -> None:
-        self.kv = kv
-        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
-        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
-
-    def __del__(self) -> None:
-        self.task.cancel()
-
-    async def aget(
-        self, pairs: List[tuple[str, str]]
-    ) -> dict[tuple[str, str], Optional[V]]:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = GetOp(pairs)
-        return await fut
-
-    async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = ListOp(prefixes)
-        return await fut
-
-    async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
-        fut = asyncio.get_running_loop().create_future()
-        self.aqueue[fut] = PutOp(writes)
-        return await fut
-
-
-async def _run(
-    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
-) -> None:
-    while True:
-        await asyncio.sleep(0)
-        if not aqueue:
-            continue
-        # this could use a lock, if we want thread safety
-        taken = aqueue.copy()
-        aqueue.clear()
-        # action each operation
-        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
-        if gets:
-            try:
-                results = await kv.aget([p for op in gets.values() for p in op.pairs])
-                for fut, op in gets.items():
-                    fut.set_result({k: results.get(k) for k in op.pairs})
-            except Exception as e:
-                for fut in gets:
-                    fut.set_exception(e)
-        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
-        if lists:
-            try:
-                results = await kv.alist(
-                    [p for op in lists.values() for p in op.prefixes]
-                )
-                for fut, op in lists.items():
-                    fut.set_result({k: results.get(k) for k in op.prefixes})
-            except Exception as e:
-                for fut in lists:
-                    fut.set_exception(e)
-        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
-        if puts:
-            try:
-                await kv.aput([w for op in puts.values() for w in op.writes])
-                for fut in puts:
-                    fut.set_result(None)
-            except Exception as e:
-                for fut in puts:
-                    fut.set_exception(e)

libs/langgraph/langgraph/kv/batch.py

@@ -0,0 +1,85 @@
+import asyncio
+from typing import NamedTuple, Optional, Union
+
+from langgraph.kv.base import BaseKV, V
+
+
+class GetOp(NamedTuple):
+    pairs: list[tuple[str, str]]
+
+
+class ListOp(NamedTuple):
+    prefixes: list[str]
+
+
+class PutOp(NamedTuple):
+    writes: list[tuple[str, str, Optional[V]]]
+
+
+class AsyncBatchedKV(BaseKV):
+    def __init__(self, kv: BaseKV) -> None:
+        self.kv = kv
+        self.aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]] = {}
+        self.task = asyncio.create_task(_run(self.aqueue, self.kv))
+
+    def __del__(self) -> None:
+        self.task.cancel()
+
+    async def aget(
+        self, pairs: list[tuple[str, str]]
+    ) -> dict[tuple[str, str], Optional[V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = GetOp(pairs)
+        return await fut
+
+    async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = ListOp(prefixes)
+        return await fut
+
+    async def aput(self, writes: list[tuple[str, str, Optional[V]]]) -> None:
+        fut = asyncio.get_running_loop().create_future()
+        self.aqueue[fut] = PutOp(writes)
+        return await fut
+
+
+async def _run(
+    aqueue: dict[asyncio.Future, Union[GetOp, ListOp, PutOp]], kv: BaseKV
+) -> None:
+    while True:
+        await asyncio.sleep(0)
+        if not aqueue:
+            continue
+        # this could use a lock, if we want thread safety
+        taken = aqueue.copy()
+        aqueue.clear()
+        # action each operation
+        gets = {f: o for f, o in taken.items() if isinstance(o, GetOp)}
+        if gets:
+            try:
+                results = await kv.aget([p for op in gets.values() for p in op.pairs])
+                for fut, op in gets.items():
+                    fut.set_result({k: results.get(k) for k in op.pairs})
+            except Exception as e:
+                for fut in gets:
+                    fut.set_exception(e)
+        lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
+        if lists:
+            try:
+                results = await kv.alist(
+                    [p for op in lists.values() for p in op.prefixes]
+                )
+                for fut, op in lists.items():
+                    fut.set_result({k: results.get(k) for k in op.prefixes})
+            except Exception as e:
+                for fut in lists:
+                    fut.set_exception(e)
+        puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
+        if puts:
+            try:
+                await kv.aput([w for op in puts.values() for w in op.writes])
+                for fut in puts:
+                    fut.set_result(None)
+            except Exception as e:
+                for fut in puts:
+                    fut.set_exception(e)

libs/langgraph/tests/test_kv.py

@@ -1,22 +1,28 @@
 import asyncio
-from typing import Any, List
+from typing import Any
 
 from pytest_mock import MockerFixture
 
-from langgraph.kv.base import BaseKV, KeyValueStore
+from langgraph.kv.base import BaseKV
+from langgraph.kv.batch import AsyncBatchedKV
 
 
-async def test_kv_queue(mocker: MockerFixture) -> None:
+async def test_kv_async_batch(mocker: MockerFixture) -> None:
     aget = mocker.stub()
+    alist = mocker.stub()
 
     class MockKV(BaseKV):
         async def aget(
-            self, pairs: List[tuple[str, str]]
+            self, pairs: list[tuple[str, str]]
         ) -> dict[tuple[str, str], dict[str, Any] | None]:
             aget(pairs)
-            return {pair: {0: pair[0], 1: pair[1]} for pair in pairs}
+            return {pair: 1 for pair in pairs}
 
-    store = KeyValueStore(MockKV())
+        async def alist(self, prefixes: list[str]) -> dict[str, dict[str, Any]]:
+            alist(prefixes)
+            return {prefix: {prefix: 1} for prefix in prefixes}
+
+    store = AsyncBatchedKV(MockKV())
 
     # concurrent calls are batched
     results = await asyncio.gather(
@@ -24,9 +30,18 @@ async def test_kv_queue(mocker: MockerFixture) -> None:
         store.aget([("c", "d")]),
     )
     assert results == [
-        {("a", "b"): {0: "a", 1: "b"}},
-        {("c", "d"): {0: "c", 1: "d"}},
+        {("a", "b"): 1},
+        {("c", "d"): 1},
     ]
     assert [c.args for c in aget.call_args_list] == [
         ([("a", "b"), ("c", "d")],),
     ]
+
+    results = await asyncio.gather(
+        store.alist(["a", "b"]),
+        store.alist(["c", "d"]),
+    )
+    assert results == [{"a": {"a": 1}, "b": {"b": 1}}, {"c": {"c": 1}, "d": {"d": 1}}]
+    assert [c.args for c in alist.call_args_list] == [
+        (["a", "b", "c", "d"],),
+    ]