From 630d9c79edca85ef0417afd78727487c973e25e2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 20 Aug 2024 13:55:28 -0700 Subject: [PATCH] Split out async batch to sep file --- libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/langgraph/kv/base.py | 84 +----------------------- libs/langgraph/langgraph/kv/batch.py | 85 +++++++++++++++++++++++++ libs/langgraph/tests/test_kv.py | 31 ++++++--- 4 files changed, 110 insertions(+), 92 deletions(-) create mode 100644 libs/langgraph/langgraph/kv/batch.py diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 261fbf1d5..c116e6d47 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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, diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/kv/base.py index 570783ec1..37e07ee84 100644 --- a/libs/langgraph/langgraph/kv/base.py +++ b/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) diff --git a/libs/langgraph/langgraph/kv/batch.py b/libs/langgraph/langgraph/kv/batch.py new file mode 100644 index 000000000..fa6cfc59b --- /dev/null +++ b/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) diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py index 2e24ee07d..5af46ef9c 100644 --- a/libs/langgraph/tests/test_kv.py +++ b/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"],), + ]