From c9e6ee6da70c3880df5ae54fd8f54e5c546a5666 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 21 Aug 2024 09:34:12 -0700 Subject: [PATCH] Rename --- libs/langgraph/langgraph/graph/state.py | 6 +++--- libs/langgraph/langgraph/managed/shared_value.py | 4 ++-- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- libs/langgraph/langgraph/pregel/loop.py | 14 +++++++------- libs/langgraph/langgraph/{kv => store}/__init__.py | 0 libs/langgraph/langgraph/{kv => store}/base.py | 2 +- libs/langgraph/langgraph/{kv => store}/batch.py | 8 ++++---- libs/langgraph/langgraph/{kv => store}/memory.py | 4 ++-- libs/langgraph/tests/test_kv.py | 8 ++++---- libs/langgraph/tests/test_pregel.py | 4 ++-- libs/langgraph/tests/test_pregel_async.py | 4 ++-- 11 files changed, 29 insertions(+), 29 deletions(-) rename libs/langgraph/langgraph/{kv => store}/__init__.py (100%) rename libs/langgraph/langgraph/{kv => store}/base.py (97%) rename libs/langgraph/langgraph/{kv => store}/batch.py (93%) rename libs/langgraph/langgraph/{kv => store}/memory.py (91%) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 26770bd20..8c85b7760 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -40,7 +40,6 @@ from langgraph.graph.graph import ( Graph, Send, ) -from langgraph.kv.base import BaseMemory from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, @@ -52,6 +51,7 @@ from langgraph.managed.base import ( from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry +from langgraph.store.base import BaseStore from langgraph.utils import RunnableCallable, coerce_to_runnable logger = logging.getLogger(__name__) @@ -383,7 +383,7 @@ class StateGraph(Graph): self, checkpointer: Optional[BaseCheckpointSaver] = None, *, - kv: Optional[BaseMemory] = None, + store: Optional[BaseStore] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: bool = False, @@ -452,7 +452,7 @@ class StateGraph(Graph): interrupt_after_nodes=interrupt_after, auto_validate=False, debug=debug, - kv=kv, + store=store, ) compiled.attach_node(START, None) diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py index a1659bb8b..74a7c1ef4 100644 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ b/libs/langgraph/langgraph/managed/shared_value.py @@ -14,13 +14,13 @@ from typing_extensions import NotRequired, Required, Self from langgraph.constants import CONFIG_KEY_KV from langgraph.errors import InvalidUpdateError -from langgraph.kv.base import BaseMemory from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, ConfiguredManagedValue, WritableManagedValue, ) +from langgraph.store.base import BaseStore V = dict[str, Any] @@ -83,7 +83,7 @@ class SharedValue(WritableManagedValue[Value, Update]): self.scope = scope self.config = config self.value: Value = {} - self.kv: BaseMemory = config["configurable"].get(CONFIG_KEY_KV) + self.kv: BaseStore = config["configurable"].get(CONFIG_KEY_KV) if self.kv is None: self.ns: Optional[str] = None elif scope_value := config["configurable"].get(self.scope): diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a5049d36f..5e058fd16 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -73,7 +73,6 @@ from langgraph.constants import ( Interrupt, ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError -from langgraph.kv.base import BaseMemory from langgraph.managed.base import ( AsyncManagedValuesManager, ManagedValuesManager, @@ -109,6 +108,7 @@ from langgraph.pregel.types import ( from langgraph.pregel.utils import get_new_channel_versions from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry +from langgraph.store.base import BaseStore WriteValue = Union[ Runnable[Input, Output], @@ -224,7 +224,7 @@ class Pregel( checkpointer: Optional[BaseCheckpointSaver] = None """Checkpointer used to save and load graph state. Defaults to None.""" - kv: Optional[BaseMemory] = None + store: Optional[BaseStore] = None """Key-value store to use. Defaults to None.""" retry_policy: Optional[RetryPolicy] = None diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 41234e9c7..8079c0dc2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -50,8 +50,6 @@ from langgraph.constants import ( Interrupt, ) from langgraph.errors import EmptyInputError, GraphInterrupt -from langgraph.kv.base import BaseMemory -from langgraph.kv.batch import AsyncBatchedKV from langgraph.managed.base import ( AsyncManagedValuesManager, ManagedValueMapping, @@ -74,6 +72,8 @@ from langgraph.pregel.executor import ( from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single from langgraph.pregel.types import PregelExecutableTask from langgraph.pregel.utils import get_new_channel_versions +from langgraph.store.base import BaseStore +from langgraph.store.batch import AsyncBatchedStore if TYPE_CHECKING: from langgraph.pregel import Pregel @@ -105,7 +105,7 @@ class PregelLoop: ] ] graph: "Pregel" - kv: Optional[BaseMemory] + store: Optional[BaseStore] submit: Submit channels: Mapping[str, BaseChannel] managed: ManagedValueMapping @@ -427,7 +427,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): graph: "Pregel", ) -> None: super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) - self.kv = graph.kv + self.store = graph.store self.stack = ExitStack() if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version @@ -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.kv}), + patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}), ) ) self.stack.push(self._suppress_interrupt) @@ -511,7 +511,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): graph: "Pregel", ) -> None: super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) - self.kv = AsyncBatchedKV(graph.kv) if graph.kv else None + self.store = AsyncBatchedStore(graph.store) if graph.store else None self.stack = AsyncExitStack() if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version @@ -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.kv}), + patch_config(self.config, configurable={CONFIG_KEY_KV: self.store}), ) ) self.stack.push(self._suppress_interrupt) diff --git a/libs/langgraph/langgraph/kv/__init__.py b/libs/langgraph/langgraph/store/__init__.py similarity index 100% rename from libs/langgraph/langgraph/kv/__init__.py rename to libs/langgraph/langgraph/store/__init__.py diff --git a/libs/langgraph/langgraph/kv/base.py b/libs/langgraph/langgraph/store/base.py similarity index 97% rename from libs/langgraph/langgraph/kv/base.py rename to libs/langgraph/langgraph/store/base.py index 0c5791eb6..7f0030f56 100644 --- a/libs/langgraph/langgraph/kv/base.py +++ b/libs/langgraph/langgraph/store/base.py @@ -3,7 +3,7 @@ from typing import Any, List, Optional V = dict[str, Any] -class BaseMemory: +class BaseStore: def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]: # list[namespace] -> dict[namespace, list[value]] raise NotImplementedError diff --git a/libs/langgraph/langgraph/kv/batch.py b/libs/langgraph/langgraph/store/batch.py similarity index 93% rename from libs/langgraph/langgraph/kv/batch.py rename to libs/langgraph/langgraph/store/batch.py index 971d50df4..cfd05e548 100644 --- a/libs/langgraph/langgraph/kv/batch.py +++ b/libs/langgraph/langgraph/store/batch.py @@ -1,7 +1,7 @@ import asyncio from typing import NamedTuple, Optional, Union -from langgraph.kv.base import BaseMemory, V +from langgraph.store.base import BaseStore, V class ListOp(NamedTuple): @@ -12,8 +12,8 @@ class PutOp(NamedTuple): writes: list[tuple[str, str, Optional[V]]] -class AsyncBatchedKV(BaseMemory): - def __init__(self, kv: BaseMemory) -> None: +class AsyncBatchedStore(BaseStore): + def __init__(self, kv: BaseStore) -> None: self.kv = kv self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {} self.task = asyncio.create_task(_run(self.aqueue, self.kv)) @@ -33,7 +33,7 @@ class AsyncBatchedKV(BaseMemory): async def _run( - aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseMemory + aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], kv: BaseStore ) -> None: while True: await asyncio.sleep(0) diff --git a/libs/langgraph/langgraph/kv/memory.py b/libs/langgraph/langgraph/store/memory.py similarity index 91% rename from libs/langgraph/langgraph/kv/memory.py rename to libs/langgraph/langgraph/store/memory.py index 4b0e92a2d..48fa2884f 100644 --- a/libs/langgraph/langgraph/kv/memory.py +++ b/libs/langgraph/langgraph/store/memory.py @@ -1,10 +1,10 @@ from collections import defaultdict from typing import List, Optional -from langgraph.kv.base import BaseMemory, V +from langgraph.store.base import BaseStore, V -class MemoryKV(BaseMemory): +class MemoryStore(BaseStore): def __init__(self) -> None: self.data: dict[str, dict[str, V]] = defaultdict(dict) diff --git a/libs/langgraph/tests/test_kv.py b/libs/langgraph/tests/test_kv.py index 68294f42c..0c2862b48 100644 --- a/libs/langgraph/tests/test_kv.py +++ b/libs/langgraph/tests/test_kv.py @@ -3,15 +3,15 @@ from typing import Any from pytest_mock import MockerFixture -from langgraph.kv.base import BaseMemory -from langgraph.kv.batch import AsyncBatchedKV +from langgraph.store.base import BaseStore +from langgraph.store.batch import AsyncBatchedStore async def test_kv_async_batch(mocker: MockerFixture) -> None: aget = mocker.stub() alist = mocker.stub() - class MockKV(BaseMemory): + class MockKV(BaseStore): async def aget( self, pairs: list[tuple[str, str]] ) -> dict[tuple[str, str], dict[str, Any] | None]: @@ -22,7 +22,7 @@ async def test_kv_async_batch(mocker: MockerFixture) -> None: alist(prefixes) return {prefix: {prefix: 1} for prefix in prefixes} - store = AsyncBatchedKV(MockKV()) + store = AsyncBatchedStore(MockKV()) # concurrent calls are batched results = await asyncio.gather( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fef936c80..8cf6aafa8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -58,7 +58,6 @@ from langgraph.graph import END, Graph from langgraph.graph.graph import START from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph -from langgraph.kv.memory import MemoryKV from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, @@ -67,6 +66,7 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask +from langgraph.store.memory import MemoryStore from tests.any_str import AnyStr, ExceptionLike from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -6250,7 +6250,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: with SqliteSaver.from_conn_string(":memory:") as saver: tool_two = tool_two_graph.compile( - kv=MemoryKV(), + store=MemoryStore(), checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"], ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b383a4d3f..d7a532489 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -52,7 +52,6 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START from langgraph.graph.message import MessageGraph, add_messages -from langgraph.kv.memory import MemoryKV from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, @@ -62,6 +61,7 @@ from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask +from langgraph.store.memory import MemoryStore from tests.any_str import AnyStr, ExceptionLike from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, @@ -4823,7 +4823,7 @@ async def test_start_branch_then() -> None: async with AsyncSqliteSaver.from_conn_string(":memory:") as saver: tool_two = tool_two_graph.compile( - kv=MemoryKV(), + store=MemoryStore(), checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"], )