diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index cf87be337..51b6e1437 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -5,7 +5,7 @@ INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" -CONFIG_KEY_KV = "__pregel_kv" +CONFIG_KEY_STORE = "__pregel_store" CONFIG_KEY_RESUMING = "__pregel_resuming" CONFIG_KEY_TASK_ID = "__pregel_task_id" INTERRUPT = "__interrupt__" @@ -18,7 +18,7 @@ RESERVED = { CONFIG_KEY_SEND, CONFIG_KEY_READ, CONFIG_KEY_CHECKPOINTER, - CONFIG_KEY_KV, + CONFIG_KEY_STORE, CONFIG_KEY_RESUMING, CONFIG_KEY_TASK_ID, INPUT, diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py index 74a7c1ef4..7f647d006 100644 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ b/libs/langgraph/langgraph/managed/shared_value.py @@ -12,7 +12,7 @@ from typing import ( from langchain_core.runnables import RunnableConfig from typing_extensions import NotRequired, Required, Self -from langgraph.constants import CONFIG_KEY_KV +from langgraph.constants import CONFIG_KEY_STORE from langgraph.errors import InvalidUpdateError from langgraph.managed.base import ( ChannelKeyPlaceholder, @@ -56,8 +56,8 @@ class SharedValue(WritableManagedValue[Value, Update]): @contextmanager def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]: with super().enter(config, **kwargs) as value: - if value.kv is not None: - saved = value.kv.list([value.ns]) + if value.store is not None: + saved = value.store.list([value.ns]) value.value = saved[value.ns] yield value @@ -65,8 +65,8 @@ class SharedValue(WritableManagedValue[Value, Update]): @asynccontextmanager async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]: async with super().aenter(config, **kwargs) as value: - if value.kv is not None: - saved = await value.kv.alist([value.ns]) + if value.store is not None: + saved = await value.store.alist([value.ns]) value.value = saved[value.ns] yield value @@ -83,8 +83,8 @@ class SharedValue(WritableManagedValue[Value, Update]): self.scope = scope self.config = config self.value: Value = {} - self.kv: BaseStore = config["configurable"].get(CONFIG_KEY_KV) - if self.kv is None: + self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE) + if self.store is None: self.ns: Optional[str] = None elif scope_value := config["configurable"].get(self.scope): self.ns = f"scoped:{scope}:{key}:{scope_value}" @@ -114,13 +114,13 @@ class SharedValue(WritableManagedValue[Value, Update]): return writes def update(self, values: Sequence[Update]) -> None: - if self.kv is None: + if self.store is None: self._process_update(values) else: - return self.kv.put(self._process_update(values)) + return self.store.put(self._process_update(values)) async def aupdate(self, writes: Sequence[Update]) -> None: - if self.kv is None: + if self.store is None: self._process_update(writes) else: - return await self.kv.aput(self._process_update(writes)) + return await self.store.aput(self._process_update(writes)) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 8079c0dc2..a1eadb697 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -41,9 +41,9 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import ( - CONFIG_KEY_KV, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, + CONFIG_KEY_STORE, ERROR, INPUT, INTERRUPT, @@ -479,7 +479,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): self.managed = self.stack.enter_context( ManagedValuesManager( self.graph.managed_values_dict, - patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}), + patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}), ) ) self.stack.push(self._suppress_interrupt) @@ -567,7 +567,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): self.managed = await self.stack.enter_async_context( AsyncManagedValuesManager( self.graph.managed_values_dict, - patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}), + patch_config(self.config, configurable={CONFIG_KEY_STORE: self.store}), ) ) self.stack.push(self._suppress_interrupt) diff --git a/libs/langgraph/langgraph/store/batch.py b/libs/langgraph/langgraph/store/batch.py index cfd05e548..54eb20d47 100644 --- a/libs/langgraph/langgraph/store/batch.py +++ b/libs/langgraph/langgraph/store/batch.py @@ -13,10 +13,10 @@ class PutOp(NamedTuple): class AsyncBatchedStore(BaseStore): - def __init__(self, kv: BaseStore) -> None: - self.kv = kv + def __init__(self, store: BaseStore) -> None: + self.store = store self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {} - self.task = asyncio.create_task(_run(self.aqueue, self.kv)) + self.task = asyncio.create_task(_run(self.aqueue, self.store)) def __del__(self) -> None: self.task.cancel() @@ -33,7 +33,7 @@ class AsyncBatchedStore(BaseStore): async def _run( - aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseStore + aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], store: BaseStore ) -> None: while True: await asyncio.sleep(0) @@ -46,7 +46,7 @@ async def _run( lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)} if lists: try: - results = await kv.alist( + results = await store.alist( [p for op in lists.values() for p in op.prefixes] ) for fut, op in lists.items(): @@ -57,7 +57,7 @@ async def _run( 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]) + await store.aput([w for op in puts.values() for w in op.writes]) for fut in puts: fut.set_result(None) except Exception as e: diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_store.py similarity index 68% rename from libs/langgraph/tests/test_kv.py rename to libs/langgraph/tests/test_store.py index 0c2862b48..a9fd0d1d4 100644 --- a/libs/langgraph/tests/test_kv.py +++ b/libs/langgraph/tests/test_store.py @@ -7,11 +7,11 @@ from langgraph.store.base import BaseStore from langgraph.store.batch import AsyncBatchedStore -async def test_kv_async_batch(mocker: MockerFixture) -> None: +async def test_async_batch_store(mocker: MockerFixture) -> None: aget = mocker.stub() alist = mocker.stub() - class MockKV(BaseStore): + class MockStore(BaseStore): async def aget( self, pairs: list[tuple[str, str]] ) -> dict[tuple[str, str], dict[str, Any] | None]: @@ -22,21 +22,9 @@ async def test_kv_async_batch(mocker: MockerFixture) -> None: alist(prefixes) return {prefix: {prefix: 1} for prefix in prefixes} - store = AsyncBatchedStore(MockKV()) + store = AsyncBatchedStore(MockStore()) # concurrent calls are batched - results = await asyncio.gather( - store.aget([("a", "b")]), - store.aget([("c", "d")]), - ) - assert results == [ - {("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"]),