From bc86757e734e37df548e30c6066dedc1f0bc6e48 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 Aug 2024 12:46:31 -0700 Subject: [PATCH] lib: Context values never stored in checkpoints - Convert Context to a ManagedValue - Add shim for old Context constructor - Add `runtime` flag for managed values, which, prior to serialization, replaces the value with a placeholder, and replaces it back with the actual value on resuming from checkpoint --- libs/langgraph/langgraph/channels/context.py | 125 +----------------- libs/langgraph/langgraph/constants.py | 2 + libs/langgraph/langgraph/graph/state.py | 33 +++-- libs/langgraph/langgraph/managed/base.py | 38 +++++- libs/langgraph/langgraph/managed/context.py | 85 ++++++++++++ .../langgraph/managed/shared_value.py | 1 - libs/langgraph/langgraph/pregel/__init__.py | 51 +++++-- libs/langgraph/langgraph/pregel/algo.py | 30 ++++- libs/langgraph/langgraph/pregel/manager.py | 44 +++--- libs/langgraph/langgraph/pregel/read.py | 10 +- libs/langgraph/tests/test_channels.py | 84 +----------- libs/langgraph/tests/test_pregel.py | 16 ++- libs/langgraph/tests/test_pregel_async.py | 13 +- 13 files changed, 264 insertions(+), 268 deletions(-) create mode 100644 libs/langgraph/langgraph/managed/context.py diff --git a/libs/langgraph/langgraph/channels/context.py b/libs/langgraph/langgraph/channels/context.py index b48260b40..3b4e26805 100644 --- a/libs/langgraph/langgraph/channels/context.py +++ b/libs/langgraph/langgraph/channels/context.py @@ -1,124 +1,5 @@ -from contextlib import asynccontextmanager, contextmanager -from inspect import signature -from typing import ( - Any, - AsyncContextManager, - AsyncGenerator, - ContextManager, - Generator, - Generic, - Optional, - Sequence, - Type, - Union, -) +from langgraph.managed.context import Context as ContextManagedValue -from langchain_core.runnables import RunnableConfig -from typing_extensions import Self +Context = ContextManagedValue.of -from langgraph.channels.base import BaseChannel, Value -from langgraph.errors import EmptyChannelError, InvalidUpdateError - - -class Context(Generic[Value], BaseChannel[Value, None, None]): - """Exposes the value of a context manager, for the duration of an invocation. - Context manager is entered before the first step, and exited after the last step. - Optionally, provide an equivalent async context manager, which will be used - instead for async invocations. - - ```python - import httpx - - client = Channels.Context(httpx.Client, httpx.AsyncClient) - ``` - """ - - value: Value - - def __init__( - self, - ctx: Union[ - None, Type[ContextManager[Value]], Type[AsyncContextManager[Value]] - ] = None, - actx: Optional[Type[AsyncContextManager[Value]]] = None, - ) -> None: - if ctx is None and actx is None: - raise ValueError("Must provide either sync or async context manager.") - self.ctx = ctx - self.actx = actx - - def __eq__(self, value: object) -> bool: - return ( - isinstance(value, Context) - and value.ctx == self.ctx - and value.actx == self.actx - ) - - @property - def ValueType(self) -> Any: - """The type of the value stored in the channel.""" - return None - - @property - def UpdateType(self) -> Type[None]: - """The type of the update received by the channel.""" - return None - - def checkpoint(self) -> None: - raise EmptyChannelError() - - @contextmanager - def from_checkpoint( - self, checkpoint: None, config: RunnableConfig - ) -> Generator[Self, None, None]: - if self.ctx is None: - raise ValueError("Cannot enter sync context manager.") - - empty = self.__class__(ctx=self.ctx, actx=self.actx) - ctx = ( - self.ctx(config) - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - with ctx as value: - empty.value = value - yield empty - - @asynccontextmanager - async def afrom_checkpoint( - self, checkpoint: None, config: RunnableConfig - ) -> AsyncGenerator[Self, None]: - empty = self.__class__(ctx=self.ctx, actx=self.actx) - if self.actx is not None: - ctx = ( - self.actx(config) - if signature(self.actx).parameters.get("config") - else self.actx() - ) - else: - ctx = ( - self.ctx(config) - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - if hasattr(ctx, "__aenter__"): - async with ctx as value: - empty.value = value - yield empty - else: - with ctx as value: - empty.value = value - yield empty - - def update(self, values: Sequence[None]) -> bool: - if values: - raise InvalidUpdateError( - f"At key '{self.key}': Context channel does not accept writes." - ) - return False - - def get(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() +__all__ = ["Context"] diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 21844a565..e2621f4f6 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -11,6 +11,7 @@ CONFIG_KEY_TASK_ID = "__pregel_task_id" INTERRUPT = "__interrupt__" ERROR = "__error__" TASKS = "__pregel_tasks" +RUNTIME_PLACEHOLDER = "__pregel_runtime_placeholder__" RESERVED = { INTERRUPT, ERROR, @@ -22,6 +23,7 @@ RESERVED = { CONFIG_KEY_RESUMING, CONFIG_KEY_TASK_ID, INPUT, + RUNTIME_PLACEHOLDER, } TAG_HIDDEN = "langsmith:hidden" diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 7228c5fa2..94de80284 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -6,6 +6,7 @@ from functools import partial from inspect import isclass, isfunction, signature from typing import ( Any, + Callable, NamedTuple, Optional, Sequence, @@ -25,7 +26,6 @@ from langchain_core.runnables.utils import ( from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitForNames from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue @@ -425,16 +425,14 @@ class StateGraph(Graph): else [ key for key, val in self.schemas[self.output].items() - if not isinstance(val, Context) and not is_managed_value(val) + if not is_managed_value(val) ] ) stream_channels = ( "__root__" if len(self.channels) == 1 and "__root__" in self.channels else [ - key - for key, val in self.channels.items() - if not isinstance(val, Context) and not is_managed_value(val) + key for key, val in self.channels.items() if not is_managed_value(val) ] ) @@ -502,7 +500,6 @@ class CompiledStateGraph(CompiledGraph): k: (self.channels[k].UpdateType, None) for k in self.builder.schemas[self.builder.input] if isinstance(self.channels[k], BaseChannel) - and not isinstance(self.channels[k], Context) }, ) @@ -523,7 +520,7 @@ class CompiledStateGraph(CompiledGraph): output_keys = [ k for k, v in self.builder.schemas[self.builder.input].items() - if not isinstance(v, Context) and not is_managed_value(v) + if not is_managed_value(v) ] else: output_keys = list(self.builder.channels) + [ @@ -650,7 +647,14 @@ class CompiledStateGraph(CompiledGraph): return ChannelWrite(writes, tags=[TAG_HIDDEN]) # attach branch publisher - self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.builder)) + schema = ( + self.builder.nodes[start].input + if start in self.builder.nodes + else self.builder.schema + ) + self.nodes[start] |= branch.run( + branch_writer, _get_state_reader(self.builder, schema) + ) # attach branch subscribers ends = ( @@ -676,16 +680,17 @@ class CompiledStateGraph(CompiledGraph): ) -def _get_state_reader(graph: StateGraph) -> ChannelRead: - state_keys = list(graph.channels) +def _get_state_reader( + builder: StateGraph, schema: Type[Any] +) -> Callable[[RunnableConfig], Any]: + state_keys = list(builder.channels) + select = list(builder.schemas[schema]) return partial( ChannelRead.do_read, - channel=state_keys[0] if state_keys == ["__root__"] else state_keys, + select=select[0] if select == ["__root__"] else select, fresh=True, # coerce state dict to schema class (eg. pydantic model) - mapper=( - None if state_keys == ["__root__"] else partial(_coerce_state, graph.schema) - ), + mapper=(None if state_keys == ["__root__"] else partial(_coerce_state, schema)), ) diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index bebca1be2..c5516ca1e 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/libs/langgraph/langgraph/managed/base.py @@ -16,11 +16,16 @@ from typing import ( from langchain_core.runnables import RunnableConfig from typing_extensions import Self, TypeGuard +from langgraph.constants import RUNTIME_PLACEHOLDER + V = TypeVar("V") U = TypeVar("U") class ManagedValue(ABC, Generic[V]): + runtime: bool = False + """Whether the managed value is always created at runtime, ie. never stored.""" + def __init__(self, config: RunnableConfig) -> None: self.config = config @@ -74,8 +79,6 @@ class ConfiguredManagedValue(NamedTuple): ManagedValueSpec = Union[Type[ManagedValue], ConfiguredManagedValue] -ManagedValueMapping = dict[str, ManagedValue] - def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]: return (isclass(value) and issubclass(value, ManagedValue)) or isinstance( @@ -103,3 +106,34 @@ def is_writable_managed_value(value: Any) -> TypeGuard[Type[WritableManagedValue ChannelKeyPlaceholder = object() ChannelTypePlaceholder = object() + + +class ManagedValueMapping(dict[str, ManagedValue]): + def replace_runtime_values(self, step: int, values: Union[dict[str, Any], Any]): + print("replace_runtime_values", values) + if isinstance(values, dict): + for key, value in values.items(): + for chan, mv in self.items(): + print("chan", chan, "mv", mv, "v", mv(step), "value", value) + print(mv, mv.runtime, mv(step) is value) + if mv.runtime and mv(step) is value: + values[key] = {RUNTIME_PLACEHOLDER: chan} + elif hasattr(values, "__dir__") and callable(values.__dir__): + for key in dir(values): + value = getattr(values, key) + for chan, mv in self.items(): + if mv.runtime and mv(step) is value: + setattr(values, key, {RUNTIME_PLACEHOLDER: chan}) + + def replace_runtime_placeholders( + self, step: int, values: Union[dict[str, Any], Any] + ): + if isinstance(values, dict): + for key, value in values.items(): + if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value: + values[key] = self[value[RUNTIME_PLACEHOLDER]](step) + elif hasattr(values, "__dir__") and callable(values.__dir__): + for key in dir(values): + value = getattr(values, key) + if isinstance(value, dict) and RUNTIME_PLACEHOLDER in value: + setattr(values, key, self[value[RUNTIME_PLACEHOLDER]](step)) diff --git a/libs/langgraph/langgraph/managed/context.py b/libs/langgraph/langgraph/managed/context.py new file mode 100644 index 000000000..381257a87 --- /dev/null +++ b/libs/langgraph/langgraph/managed/context.py @@ -0,0 +1,85 @@ +from contextlib import asynccontextmanager, contextmanager +from inspect import signature +from typing import ( + Any, + AsyncContextManager, + AsyncIterator, + ContextManager, + Iterator, + Optional, + Self, + Type, + Union, +) + +from langchain_core.runnables import RunnableConfig + +from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V + + +class Context(ManagedValue): + runtime = True + + value: V + + @staticmethod + def of( + ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None, + actx: Optional[Type[AsyncContextManager[V]]] = None, + ) -> ConfiguredManagedValue: + if ctx is None and actx is None: + raise ValueError("Must provide either sync or async context manager.") + return ConfiguredManagedValue(Context, {"ctx": ctx, "actx": actx}) + + @classmethod + @contextmanager + def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]: + with super().enter(config, **kwargs) as self: + if self.ctx is None: + raise ValueError("Cannot enter sync context manager.") + ctx = ( + self.ctx(config) + if signature(self.ctx).parameters.get("config") + else self.ctx() + ) + with ctx as v: + self.value = v + yield self + + @classmethod + @asynccontextmanager + async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]: + async with super().aenter(config, **kwargs) as self: + if self.actx is not None: + ctx = ( + self.actx(config) + if signature(self.actx).parameters.get("config") + else self.actx() + ) + else: + ctx = ( + self.ctx(config) + if signature(self.ctx).parameters.get("config") + else self.ctx() + ) + if hasattr(ctx, "__aenter__"): + async with ctx as v: + self.value = v + yield self + else: + with ctx as v: + self.value = v + yield self + + def __init__( + self, + config: RunnableConfig, + *, + ctx: Union[None, Type[ContextManager[V]], Type[AsyncContextManager[V]]] = None, + actx: Optional[Type[AsyncContextManager[V]]] = None, + ) -> None: + self.ctx = ctx + self.actx = actx + + def __call__(self, step: int) -> V: + return self.value diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py index 7bb6e23b7..f5e0561bd 100644 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ b/libs/langgraph/langgraph/managed/shared_value.py @@ -81,7 +81,6 @@ class SharedValue(WritableManagedValue[Value, Update]): ): raise ValueError("SharedValue must be a dict") self.scope = scope - self.config = config self.value: Value = {} self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE) if self.store is None: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c06f2bce3..73db1bfb8 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -51,7 +51,6 @@ from typing_extensions import Self from langgraph.channels.base import ( BaseChannel, ) -from langgraph.channels.context import Context from langgraph.checkpoint.base import ( BaseCheckpointSaver, copy_checkpoint, @@ -68,7 +67,12 @@ from langgraph.constants import ( ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ManagedValueSpec -from langgraph.pregel.algo import apply_writes, local_read, prepare_next_tasks +from langgraph.pregel.algo import ( + apply_writes, + local_read, + local_write, + prepare_next_tasks, +) from langgraph.pregel.debug import ( print_step_checkpoint, print_step_tasks, @@ -330,10 +334,7 @@ class Pregel( @property def stream_channels_asis(self) -> Union[str, Sequence[str]]: return self.stream_channels or [ - k - for k in self.channels - if isinstance(self.channels[k], BaseChannel) - and not isinstance(self.channels[k], Context) + k for k in self.channels if isinstance(self.channels[k], BaseChannel) ] def get_state(self, config: RunnableConfig) -> StateSnapshot: @@ -564,7 +565,7 @@ class Pregel( # update channels with ChannelsManager(self.channels, checkpoint, config) as ( channels, - _, + managed, ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() @@ -588,9 +589,22 @@ class Pregel( run_name=self.name + "UpdateState", configurable={ # deque.extend is thread-safe - CONFIG_KEY_SEND: task.writes.extend, + CONFIG_KEY_SEND: partial( + local_write, + step + 1, + task.writes.extend, + self.nodes, + channels, + managed, + ), CONFIG_KEY_READ: partial( - local_read, checkpoint, channels, task, config + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, ), }, ), @@ -682,7 +696,7 @@ class Pregel( # update channels, acting as the chosen node async with AsyncChannelsManager(self.channels, checkpoint, config) as ( channels, - _, + managed, ): # create task to run all writers of the chosen node writers = self.nodes[as_node].get_writers() @@ -706,9 +720,22 @@ class Pregel( run_name=self.name + "UpdateState", configurable={ # deque.extend is thread-safe - CONFIG_KEY_SEND: task.writes.extend, + CONFIG_KEY_SEND: partial( + local_write, + step + 1, + task.writes.extend, + self.nodes, + channels, + managed, + ), CONFIG_KEY_READ: partial( - local_read, checkpoint, channels, task, config + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, ), }, ), diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 1d3f57040..a1dd5d536 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -24,7 +24,6 @@ from langchain_core.runnables.config import ( ) from langgraph.channels.base import BaseChannel -from langgraph.channels.context import Context from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -95,28 +94,37 @@ def should_interrupt( def local_read( + step: int, checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], + managed: ManagedValueMapping, task: WritesProtocol, config: RunnableConfig, select: Union[list[str], str], fresh: bool = False, ) -> Union[dict[str, Any], Any]: + if isinstance(select, str): + managed_keys = [] + else: + managed_keys = [k for k in select if k in managed] + select = [k for k in select if k not in managed] if fresh: new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1) - context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)} with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as ( channels, _, ): - all_channels = {**channels, **context_channels} - apply_writes(new_checkpoint, all_channels, [task], None) - return read_channels(all_channels, select) + apply_writes(new_checkpoint, channels, [task], None) + values = read_channels(channels, select) else: - return read_channels(channels, select) + values = read_channels(channels, select) + if managed_keys: + values.update({k: managed[k](step) for k in managed_keys}) + return values def local_write( + step: int, commit: Callable[[Sequence[tuple[str, Any]]], None], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], @@ -131,6 +139,9 @@ def local_write( ) if value.node not in processes: raise InvalidUpdateError(f"Invalid node name {value.node} in packet") + # replace any runtime values with placeholders + managed.replace_runtime_values(step, value.arg) + print("after replace", value) elif chan not in channels and chan not in managed: logger.warning(f"Skipping write for channel '{chan}' which has no readers") commit(writes) @@ -293,6 +304,7 @@ def prepare_next_tasks( if for_execution: proc = processes[packet.node] if node := proc.get_node(): + managed.replace_runtime_placeholders(step, packet.arg) writes = deque() tasks.append( PregelExecutableTask( @@ -317,6 +329,7 @@ def prepare_next_tasks( # deque.extend is thread-safe CONFIG_KEY_SEND: partial( local_write, + step, writes.extend, processes, channels, @@ -324,8 +337,10 @@ def prepare_next_tasks( ), CONFIG_KEY_READ: partial( local_read, + step, checkpoint, channels, + managed, PregelTaskWrites(packet.node, writes, triggers), config, ), @@ -412,6 +427,7 @@ def prepare_next_tasks( # deque.extend is thread-safe CONFIG_KEY_SEND: partial( local_write, + step, writes.extend, processes, channels, @@ -419,8 +435,10 @@ def prepare_next_tasks( ), CONFIG_KEY_READ: partial( local_read, + step, checkpoint, channels, + managed, PregelTaskWrites(name, writes, triggers), config, ), diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py index 849395c50..ccfc8de8a 100644 --- a/libs/langgraph/langgraph/pregel/manager.py +++ b/libs/langgraph/langgraph/pregel/manager.py @@ -5,8 +5,6 @@ from typing import AsyncIterator, Iterator, Mapping, Optional, Union from langchain_core.runnables import RunnableConfig, patch_config from langgraph.channels.base import BaseChannel -from langgraph.channels.context import Context -from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import Checkpoint from langgraph.constants import CONFIG_KEY_STORE from langgraph.managed.base import ( @@ -14,6 +12,7 @@ from langgraph.managed.base import ( ManagedValueMapping, ManagedValueSpec, ) +from langgraph.managed.context import Context from langgraph.store.base import BaseStore @@ -31,10 +30,12 @@ def ChannelsManager( channel_specs: Mapping[str, BaseChannel] = {} managed_specs: Mapping[str, ManagedValueSpec] = {} for k, v in specs.items(): - if skip_context and isinstance(v, Context): - channel_specs[k] = LastValue(None) - elif isinstance(v, BaseChannel): + if isinstance(v, BaseChannel): channel_specs[k] = v + elif ( + skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context + ): + managed_specs[k] = Context.of(noop_context) else: managed_specs[k] = v with ExitStack() as stack: @@ -45,14 +46,16 @@ def ChannelsManager( ) for k, v in channel_specs.items() }, - { - key: stack.enter_context( - value.cls.enter(config_for_managed, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.enter(config_for_managed) - ) - for key, value in managed_specs.items() - }, + ManagedValueMapping( + { + key: stack.enter_context( + value.cls.enter(config_for_managed, **value.kwargs) + if isinstance(value, ConfiguredManagedValue) + else value.enter(config_for_managed) + ) + for key, value in managed_specs.items() + } + ), ) @@ -70,10 +73,12 @@ async def AsyncChannelsManager( channel_specs: Mapping[str, BaseChannel] = {} managed_specs: Mapping[str, ManagedValueSpec] = {} for k, v in specs.items(): - if skip_context and isinstance(v, Context): - channel_specs[k] = LastValue(None) - elif isinstance(v, BaseChannel): + if isinstance(v, BaseChannel): channel_specs[k] = v + elif ( + skip_context and isinstance(v, ConfiguredManagedValue) and v.cls is Context + ): + managed_specs[k] = Context.of(noop_context) else: managed_specs[k] = v async with AsyncExitStack() as stack: @@ -102,5 +107,10 @@ async def AsyncChannelsManager( for k, v in channel_specs.items() }, # managed: build mapping from spec to result - {tasks[task]: task.result() for task in done}, + ManagedValueMapping({tasks[task]: task.result() for task in done}), ) + + +@contextmanager +def noop_context() -> Iterator[None]: + yield None diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index b5c971e4a..163665e2e 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -67,19 +67,19 @@ class ChannelRead(RunnableCallable): def _read(self, _: Any, config: RunnableConfig) -> Any: return self.do_read( - config, channel=self.channel, fresh=self.fresh, mapper=self.mapper + config, select=self.channel, fresh=self.fresh, mapper=self.mapper ) async def _aread(self, _: Any, config: RunnableConfig) -> Any: return self.do_read( - config, channel=self.channel, fresh=self.fresh, mapper=self.mapper + config, select=self.channel, fresh=self.fresh, mapper=self.mapper ) @staticmethod def do_read( config: RunnableConfig, *, - channel: Union[str, list[str]], + select: Union[str, list[str]], fresh: bool = False, mapper: Optional[Callable[[Any], Any]] = None, ) -> Any: @@ -91,9 +91,9 @@ class ChannelRead(RunnableCallable): "Make sure to call in the context of a Pregel process" ) if mapper: - return mapper(read(channel, fresh)) + return mapper(read(select, fresh)) else: - return read(channel, fresh) + return read(select, fresh) DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index b7d936878..3e0c924ae 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -1,14 +1,9 @@ import operator -from contextlib import asynccontextmanager, contextmanager -from typing import AsyncGenerator, Generator, Sequence, Union +from typing import Sequence, Union -import httpx import pytest -from langchain_core.runnables import RunnableConfig -from pytest_mock import MockerFixture from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -257,80 +252,3 @@ async def test_binop_async() -> None: checkpoint, {} ) as channel: assert channel.get() == 10 - - -def test_ctx_manager(mocker: MockerFixture) -> None: - setup = mocker.Mock() - cleanup = mocker.Mock() - - @contextmanager - def an_int() -> Generator[int, None, None]: - setup() - try: - yield 5 - finally: - cleanup() - - with Context(an_int, None).from_checkpoint(None, {}) as channel: - assert setup.call_count == 1 - assert cleanup.call_count == 0 - - assert channel.ValueType is None - assert channel.UpdateType is None - - assert channel.get() == 5 - - with pytest.raises(InvalidUpdateError): - channel.update([5]) # type: ignore - - assert setup.call_count == 1 - assert cleanup.call_count == 1 - - -def test_ctx_manager_ctx(mocker: MockerFixture) -> None: - with Context(httpx.Client).from_checkpoint(None, {}) as channel: - assert channel.ValueType is None - assert channel.UpdateType is None - - assert isinstance(channel.get(), httpx.Client) - - with pytest.raises(InvalidUpdateError): - channel.update([5]) # type: ignore - - with pytest.raises(EmptyChannelError): - channel.checkpoint() - - -async def test_ctx_manager_async(mocker: MockerFixture) -> None: - setup = mocker.Mock() - cleanup = mocker.Mock() - - @contextmanager - def an_int_sync(config: RunnableConfig) -> Generator[int, None, None]: - try: - yield 5 - finally: - pass - - @asynccontextmanager - async def an_int() -> AsyncGenerator[int, None]: - setup() - try: - yield 5 - finally: - cleanup() - - async with Context(an_int_sync, an_int).afrom_checkpoint(None, {}) as channel: - assert setup.call_count == 1 - assert cleanup.call_count == 0 - - assert channel.ValueType is None - assert channel.UpdateType is None - - assert channel.get() == 5 - - with pytest.raises(InvalidUpdateError): - channel.update([5]) # type: ignore - - assert setup.call_count == 1 - assert cleanup.call_count == 1 diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 7325aa41d..ad13e52c0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4042,7 +4042,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ["memory", "sqlite", "postgres", "postgres_pipe"], ) def test_state_graph_packets( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture ) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -4062,6 +4062,7 @@ def test_state_graph_packets( class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] + session: Annotated[httpx.Client, Context(httpx.Client)] @tool() def search_api(query: str) -> str: @@ -4117,11 +4118,19 @@ def test_state_graph_packets( ), "nodes can pass extra data to their cond edges, which isn't saved in state" # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [Send("tools", tool_call) for tool_call in tool_calls] + return [ + Send("tools", {"call": tool_call, "my_session": data["session"]}) + for tool_call in tool_calls + ] else: return END - def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: + class ToolInput(TypedDict): + call: ToolCall + my_session: httpx.Client + + def tools_node(input: ToolInput, config: RunnableConfig) -> AgentState: + tool_call = input["call"] time.sleep(tool_call["args"].get("idx", 0) / 10) output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) return { @@ -7492,7 +7501,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1( return {"answer": ",".join(data.docs)} def decider(data: State) -> str: - print("decider", data) assert isinstance(data, State) return "retriever_two" diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a35bf7b00..6d0436bb5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -3892,6 +3892,7 @@ async def test_state_graph_packets() -> None: class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] + session: Annotated[httpx.Client, Context(httpx.Client)] @tool() def search_api(query: str) -> str: @@ -3938,11 +3939,19 @@ async def test_state_graph_packets() -> None: def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [Send("tools", tool_call) for tool_call in tool_calls] + return [ + Send("tools", {"call": tool_call, "my_session": data["session"]}) + for tool_call in tool_calls + ] else: return END - async def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: + class ToolInput(TypedDict): + call: ToolCall + my_session: httpx.Client + + async def tools_node(input: ToolInput, config: RunnableConfig) -> AgentState: + tool_call = input["call"] await asyncio.sleep(tool_call["args"].get("idx", 0) / 10) output = await tools_by_name[tool_call["name"]].ainvoke( tool_call["args"], config