From 889b40e7a821c66dd5b41f99e85947881f389ec1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 May 2025 11:05:15 -0700 Subject: [PATCH] Remove Context channel / managed value, Remove SharedValue - Both have been deprecated long ago, have not been present in docs for quite a while - Removing these lets us delete some code paths that were dedicated to this, easing work on distributed runner --- libs/langgraph/Makefile | 2 +- libs/langgraph/langgraph/channels/__init__.py | 2 - libs/langgraph/langgraph/channels/context.py | 5 - libs/langgraph/langgraph/graph/state.py | 14 +- libs/langgraph/langgraph/managed/base.py | 45 +-- libs/langgraph/langgraph/managed/context.py | 111 ------ .../langgraph/managed/shared_value.py | 120 ------ libs/langgraph/langgraph/pregel/__init__.py | 8 +- libs/langgraph/langgraph/pregel/algo.py | 12 +- libs/langgraph/langgraph/pregel/draw.py | 5 +- libs/langgraph/langgraph/pregel/loop.py | 34 +- libs/langgraph/langgraph/pregel/manager.py | 33 +- libs/langgraph/tests/test_large_cases.py | 334 +++++++--------- .../langgraph/tests/test_large_cases_async.py | 358 +++++++----------- libs/langgraph/tests/test_pregel.py | 153 ++------ libs/langgraph/tests/test_pregel_async.py | 178 ++------- libs/langgraph/tests/test_state.py | 59 --- 17 files changed, 375 insertions(+), 1098 deletions(-) delete mode 100644 libs/langgraph/langgraph/channels/context.py delete mode 100644 libs/langgraph/langgraph/managed/context.py delete mode 100644 libs/langgraph/langgraph/managed/shared_value.py diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 1ef067dbf..6e47fa820 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -44,7 +44,7 @@ stop-postgres: docker compose -f tests/compose-postgres.yml down -v start-dev-server: - uv run langgraph dev --config tests/example_app/langgraph.json --no-browser & + LOG_LEVEL=warning uv run langgraph dev --config tests/example_app/langgraph.json --no-browser & @echo "Dev server started." @echo "Dev server PID: $$!" > .devserver.pid diff --git a/libs/langgraph/langgraph/channels/__init__.py b/libs/langgraph/langgraph/channels/__init__.py index 6f9ba2119..cdb193484 100644 --- a/libs/langgraph/langgraph/channels/__init__.py +++ b/libs/langgraph/langgraph/channels/__init__.py @@ -1,6 +1,5 @@ from langgraph.channels.any_value import AnyValue from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -9,7 +8,6 @@ from langgraph.channels.untracked_value import UntrackedValue __all__ = [ "LastValue", "Topic", - "Context", "BinaryOperatorAggregate", "UntrackedValue", "EphemeralValue", diff --git a/libs/langgraph/langgraph/channels/context.py b/libs/langgraph/langgraph/channels/context.py deleted file mode 100644 index 3b4e26805..000000000 --- a/libs/langgraph/langgraph/channels/context.py +++ /dev/null @@ -1,5 +0,0 @@ -from langgraph.managed.context import Context as ContextManagedValue - -Context = ContextManagedValue.of - -__all__ = ["Context"] diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 2a2390b75..65ec2c329 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -66,12 +66,8 @@ from langgraph.graph.graph import ( ) from langgraph.graph.schema_utils import SchemaCoercionMapper from langgraph.managed.base import ( - ChannelKeyPlaceholder, - ChannelTypePlaceholder, - ConfiguredManagedValue, ManagedValueSpec, is_managed_value, - is_writable_managed_value, ) from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.write import ( @@ -727,9 +723,7 @@ class CompiledStateGraph(CompiledGraph): ] else: output_keys = list(self.builder.channels) + [ - k - for k, v in self.builder.managed.items() - if is_writable_managed_value(v) + k for k, v in self.builder.managed.items() ] def _get_updates( @@ -1211,12 +1205,6 @@ def _is_field_managed_value(name: str, typ: type[Any]) -> Optional[ManagedValueS if len(meta) >= 1: decoration = get_origin(meta[-1]) or meta[-1] if is_managed_value(decoration): - if isinstance(decoration, ConfiguredManagedValue): - for k, v in decoration.kwargs.items(): - if v is ChannelKeyPlaceholder: - decoration.kwargs[k] = name - if v is ChannelTypePlaceholder: - decoration.kwargs[k] = typ.__origin__ return decoration return None diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index a3f305fbe..078f71e7c 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/libs/langgraph/langgraph/managed/base.py @@ -1,13 +1,11 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator from contextlib import asynccontextmanager, contextmanager from inspect import isclass from typing import ( Any, Generic, - NamedTuple, TypeVar, - Union, ) from typing_extensions import Self, TypeGuard @@ -54,48 +52,11 @@ class ManagedValue(ABC, Generic[V]): def __call__(self) -> V: ... -class WritableManagedValue(Generic[V, U], ManagedValue[V], ABC): - @abstractmethod - def update(self, writes: Sequence[U]) -> None: ... - - @abstractmethod - async def aupdate(self, writes: Sequence[U]) -> None: ... - - -class ConfiguredManagedValue(NamedTuple): - cls: type[ManagedValue] - kwargs: dict[str, Any] - - -ManagedValueSpec = Union[type[ManagedValue], ConfiguredManagedValue] +ManagedValueSpec = type[ManagedValue] def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]: - return (isclass(value) and issubclass(value, ManagedValue)) or isinstance( - value, ConfiguredManagedValue - ) - - -def is_readonly_managed_value(value: Any) -> TypeGuard[type[ManagedValue]]: - return ( - isclass(value) - and issubclass(value, ManagedValue) - and not issubclass(value, WritableManagedValue) - ) or ( - isinstance(value, ConfiguredManagedValue) - and not issubclass(value.cls, WritableManagedValue) - ) - - -def is_writable_managed_value(value: Any) -> TypeGuard[type[WritableManagedValue]]: - return (isclass(value) and issubclass(value, WritableManagedValue)) or ( - isinstance(value, ConfiguredManagedValue) - and issubclass(value.cls, WritableManagedValue) - ) - - -ChannelKeyPlaceholder = object() -ChannelTypePlaceholder = object() + return isclass(value) and issubclass(value, ManagedValue) ManagedValueMapping = dict[str, ManagedValue] diff --git a/libs/langgraph/langgraph/managed/context.py b/libs/langgraph/langgraph/managed/context.py deleted file mode 100644 index 1352254d5..000000000 --- a/libs/langgraph/langgraph/managed/context.py +++ /dev/null @@ -1,111 +0,0 @@ -from collections.abc import AsyncIterator, Iterator -from contextlib import ( - AbstractAsyncContextManager, - AbstractContextManager, - asynccontextmanager, - contextmanager, -) -from inspect import signature -from typing import ( - Any, - Callable, - Generic, - Optional, - Union, -) - -from typing_extensions import Self - -from langgraph.managed.base import ConfiguredManagedValue, ManagedValue, V -from langgraph.types import LoopProtocol - - -class Context(ManagedValue[V], Generic[V]): - runtime = True - - value: V - - @staticmethod - def of( - ctx: Union[ - None, - Callable[..., AbstractContextManager[V]], - type[AbstractContextManager[V]], - Callable[..., AbstractAsyncContextManager[V]], - type[AbstractAsyncContextManager[V]], - ] = None, - actx: Optional[ - Union[ - Callable[..., AbstractAsyncContextManager[V]], - type[AbstractAsyncContextManager[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, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - with super().enter(loop, **kwargs) as self: - if self.ctx is None: - raise ValueError( - "Synchronous context manager not found. Please initialize Context value with a sync context manager, or invoke your graph asynchronously." - ) - ctx = ( - self.ctx(loop.config) # type: ignore[call-arg] - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - with ctx as v: # type: ignore[union-attr] - self.value = v - yield self - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - async with super().aenter(loop, **kwargs) as self: - if self.actx is not None: - ctx = ( - self.actx(loop.config) # type: ignore[call-arg] - if signature(self.actx).parameters.get("config") - else self.actx() - ) - elif self.ctx is not None: - ctx = ( - self.ctx(loop.config) # type: ignore - if signature(self.ctx).parameters.get("config") - else self.ctx() - ) - else: - raise ValueError( - "Asynchronous context manager not found. Please initialize Context value with an async context manager, or invoke your graph synchronously." - ) - if hasattr(ctx, "__aenter__"): - async with ctx as v: - self.value = v - yield self - elif hasattr(ctx, "__enter__") and hasattr(ctx, "__exit__"): - with ctx as v: - self.value = v - yield self - else: - raise ValueError( - "Context manager must have either __enter__ or __aenter__ method." - ) - - def __init__( - self, - loop: LoopProtocol, - *, - ctx: Union[ - None, type[AbstractContextManager[V]], type[AbstractAsyncContextManager[V]] - ] = None, - actx: Optional[type[AbstractAsyncContextManager[V]]] = None, - ) -> None: - self.ctx = ctx - self.actx = actx - - def __call__(self) -> V: - return self.value diff --git a/libs/langgraph/langgraph/managed/shared_value.py b/libs/langgraph/langgraph/managed/shared_value.py deleted file mode 100644 index 39f4684cc..000000000 --- a/libs/langgraph/langgraph/managed/shared_value.py +++ /dev/null @@ -1,120 +0,0 @@ -import collections.abc -from collections.abc import AsyncIterator, Iterator, Sequence -from contextlib import asynccontextmanager, contextmanager -from typing import ( - Any, - Optional, -) - -from typing_extensions import NotRequired, Required, Self - -from langgraph.constants import CONF -from langgraph.errors import InvalidUpdateError -from langgraph.managed.base import ( - ChannelKeyPlaceholder, - ChannelTypePlaceholder, - ConfiguredManagedValue, - WritableManagedValue, -) -from langgraph.store.base import PutOp -from langgraph.types import LoopProtocol - -V = dict[str, Any] - - -Value = dict[str, V] -Update = dict[str, Optional[V]] - - -# Adapted from typing_extensions -def _strip_extras(t): # type: ignore[no-untyped-def] - """Strips Annotated, Required and NotRequired from a given type.""" - if hasattr(t, "__origin__"): - return _strip_extras(t.__origin__) - if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired): - return _strip_extras(t.__args__[0]) - - return t - - -class SharedValue(WritableManagedValue[Value, Update]): - @staticmethod - def on(scope: str) -> ConfiguredManagedValue: - return ConfiguredManagedValue( - SharedValue, - { - "scope": scope, - "key": ChannelKeyPlaceholder, - "typ": ChannelTypePlaceholder, - }, - ) - - @classmethod - @contextmanager - def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]: - with super().enter(loop, **kwargs) as value: - if loop.store is not None: - saved = loop.store.search(value.ns) - value.value = {it.key: it.value for it in saved} - yield value - - @classmethod - @asynccontextmanager - async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]: - async with super().aenter(loop, **kwargs) as value: - if loop.store is not None: - saved = await loop.store.asearch(value.ns) - value.value = {it.key: it.value for it in saved} - yield value - - def __init__( - self, loop: LoopProtocol, *, typ: type[Any], scope: str, key: str - ) -> None: - super().__init__(loop) - if typ := _strip_extras(typ): - if typ not in ( - dict, - collections.abc.Mapping, - collections.abc.MutableMapping, - ): - raise ValueError("SharedValue must be a dict") - self.scope = scope - self.value: Value = {} - if self.loop.store is None: - pass - elif scope_value := self.loop.config[CONF].get(self.scope): - self.ns = ("scoped", scope, key, scope_value) - else: - raise ValueError( - f"Scope {scope} for shared state key not in config.configurable" - ) - - def __call__(self) -> Value: - return self.value - - def _process_update(self, values: Sequence[Update]) -> list[PutOp]: - writes: list[PutOp] = [] - for vv in values: - for k, v in vv.items(): - if v is None: - if k in self.value: - del self.value[k] - writes.append(PutOp(self.ns, k, None)) - elif not isinstance(v, dict): - raise InvalidUpdateError("Received a non-dict value") - else: - self.value[k] = v - writes.append(PutOp(self.ns, k, v)) - return writes - - def update(self, values: Sequence[Update]) -> None: - if self.loop.store is None: - self._process_update(values) - else: - return self.loop.store.batch(self._process_update(values)) - - async def aupdate(self, writes: Sequence[Update]) -> None: - if self.loop.store is None: - self._process_update(writes) - else: - return await self.loop.store.abatch(self._process_update(writes)) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 3d2467ece..b8de81f53 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -902,7 +902,6 @@ class Pregel(PregelProtocol): step=saved.metadata.get("step", -1) + 1, stop=saved.metadata.get("step", -1) + 2, ), - skip_context=True, ) as (channels, managed): # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1024,7 +1023,6 @@ class Pregel(PregelProtocol): step=saved.metadata.get("step", -1) + 1, stop=saved.metadata.get("step", -1) + 2, ), - skip_context=True, ) as ( channels, managed, @@ -1700,14 +1698,13 @@ class Pregel(PregelProtocol): if saved and channel_writes: checkpointer.put_writes(checkpoint_config, channel_writes, task_id) # apply to checkpoint and save - mv_writes, _ = apply_writes( + apply_writes( checkpoint, channels, run_tasks, checkpointer.get_next_version, self.trigger_to_nodes, ) - assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = checkpointer.put( checkpoint_config, @@ -2125,14 +2122,13 @@ class Pregel(PregelProtocol): checkpoint_config, channel_writes, task_id ) # apply to checkpoint and save - mv_writes, _ = apply_writes( + apply_writes( checkpoint, channels, run_tasks, checkpointer.get_next_version, self.trigger_to_nodes, ) - assert not mv_writes, "Can't write to SharedValues from update_state" checkpoint = create_checkpoint(checkpoint, channels, step + 1) # save checkpoint, after applying writes next_config = await checkpointer.aput( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index f2826ee89..f8bd96ce8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -222,7 +222,7 @@ def apply_writes( tasks: Iterable[WritesProtocol], get_next_version: Optional[GetNextVersion], trigger_to_nodes: Mapping[str, Sequence[str]], -) -> tuple[dict[str, list[Any]], set[str]]: +) -> set[str]: """Apply writes from a set of tasks (usually the tasks from a Pregel step) to the checkpoint and channels, and return managed values writes to be applied externally. @@ -234,8 +234,7 @@ def apply_writes( get_next_version: Optional function to determine the next version of a channel. Returns: - A tuple containing the managed values writes to be applied externally, and - the set of channels that were updated in this step. + Set of channels that were updated in this step. """ # sort tasks on path, to ensure deterministic order for update application # any path parts after the 3rd are ignored for sorting @@ -280,7 +279,6 @@ def apply_writes( # Group writes by channel pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) - pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): @@ -290,7 +288,9 @@ def apply_writes( elif chan in channels: pending_writes_by_channel[chan].append(val) else: - pending_writes_by_managed[chan].append(val) + logger.warning( + f"Task {task.name} with path {task.path} wrote to unknown channel {chan}, ignoring it." + ) # Find the highest version of all channels if checkpoint["channel_versions"]: @@ -341,7 +341,7 @@ def apply_writes( updated_channels.add(chan) # Return managed values writes to be applied externally - return pending_writes_by_managed, updated_channels + return updated_channels def has_next_tasks( diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py index bcae60d98..2518130a5 100644 --- a/libs/langgraph/langgraph/pregel/draw.py +++ b/libs/langgraph/langgraph/pregel/draw.py @@ -60,7 +60,6 @@ def draw_graph( specs, checkpoint, LoopProtocol(step=step, stop=-1, config=config), - skip_context=True, ) as (channels, managed): static_seen: set[Any] = set() sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {} @@ -72,7 +71,7 @@ def draw_graph( } # apply input writes input_writes = list(map_input(input_channels, {})) - _, updated_channels = apply_writes( + updated_channels = apply_writes( checkpoint, channels, [ @@ -149,7 +148,7 @@ def draw_graph( for trigger, cond, label in triggers: trigger_to_sources[trigger].add((src, cond, label)) # apply writes - _, updated_channels = apply_writes( + updated_channels = apply_writes( checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes ) # prepare next tasks diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index d4c075a99..88eb06886 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -76,7 +76,6 @@ from langgraph.errors import ( from langgraph.managed.base import ( ManagedValueMapping, ManagedValueSpec, - WritableManagedValue, ) from langgraph.pregel.algo import ( Call, @@ -490,16 +489,13 @@ class PregelLoop(LoopProtocol): ), ) # all tasks have finished - mv_writes, updated_channels = apply_writes( + updated_channels = apply_writes( self.checkpoint, self.channels, self.tasks.values(), self.checkpointer_get_next_version, self.trigger_to_nodes, ) - # apply writes to managed values - for key, values in mv_writes.items(): - self._update_mv(key, values) # validate input if requested if self.input is INPUT_SHOULD_VALIDATE: self.input = INPUT_DONE @@ -700,15 +696,13 @@ class PregelLoop(LoopProtocol): if null_writes := [ w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID ]: - mv_writes, _ = apply_writes( + apply_writes( self.checkpoint, self.channels, [PregelTaskWrites((), INPUT, null_writes, [])], self.checkpointer_get_next_version, self.trigger_to_nodes, ) - for key, values in mv_writes.items(): - self._update_mv(key, values) # proceed past previous checkpoint if is_resuming: self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) @@ -750,7 +744,7 @@ class PregelLoop(LoopProtocol): manager=None, ) # apply input writes - mv_writes, updated_channels = apply_writes( + updated_channels = apply_writes( self.checkpoint, self.channels, [ @@ -760,7 +754,6 @@ class PregelLoop(LoopProtocol): self.checkpointer_get_next_version, self.trigger_to_nodes, ) - assert not mv_writes, "Can't write to SharedValues in graph input" # save input checkpoint self._put_checkpoint({"source": "input", "writes": dict(input_writes)}) # set flag @@ -869,9 +862,6 @@ class PregelLoop(LoopProtocol): # increment step self.step += 1 - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - raise NotImplementedError - def _suppress_interrupt( self, exc_type: Optional[type[BaseException]], @@ -891,15 +881,13 @@ class PregelLoop(LoopProtocol): and self.checkpoint_pending_writes and any(task.writes for task in self.tasks.values()) ): - mv_writes, updated_channels = apply_writes( + updated_channels = apply_writes( self.checkpoint, self.channels, self.tasks.values(), self.checkpointer_get_next_version, self.trigger_to_nodes, ) - for key, values in mv_writes.items(): - self._update_mv(key, values) if not updated_channels.isdisjoint( (self.output_keys,) if isinstance(self.output_keys, str) @@ -1066,13 +1054,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): config, checkpoint, metadata, new_versions ) - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - managed_value = self.managed.get(key) - if managed_value is None: - return - - return self.submit(cast(WritableManagedValue, managed_value).update, values) - def match_cached_writes(self) -> Sequence[PregelExecutableTask]: if self.cache is None: return () @@ -1263,13 +1244,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): config, checkpoint, metadata, new_versions ) - def _update_mv(self, key: str, values: Sequence[Any]) -> None: - managed_value = self.managed.get(key) - if managed_value is None: - return - - return self.submit(cast(WritableManagedValue, managed_value).aupdate, values) - async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]: if self.cache is None: return [] diff --git a/libs/langgraph/langgraph/pregel/manager.py b/libs/langgraph/langgraph/pregel/manager.py index 2d790720d..bdb583974 100644 --- a/libs/langgraph/langgraph/pregel/manager.py +++ b/libs/langgraph/langgraph/pregel/manager.py @@ -7,11 +7,9 @@ from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint from langgraph.constants import MISSING from langgraph.managed.base import ( - ConfiguredManagedValue, ManagedValueMapping, ManagedValueSpec, ) -from langgraph.managed.context import Context from langgraph.types import LoopProtocol @@ -20,8 +18,6 @@ def ChannelsManager( specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], checkpoint: Checkpoint, loop: LoopProtocol, - *, - skip_context: bool = False, ) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" channel_specs: dict[str, BaseChannel] = {} @@ -29,10 +25,6 @@ def ChannelsManager( for k, v in specs.items(): 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: @@ -43,11 +35,7 @@ def ChannelsManager( }, ManagedValueMapping( { - key: stack.enter_context( - value.cls.enter(loop, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.enter(loop) - ) + key: stack.enter_context(value.enter(loop)) for key, value in managed_specs.items() } ), @@ -59,8 +47,6 @@ async def AsyncChannelsManager( specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]], checkpoint: Checkpoint, loop: LoopProtocol, - *, - skip_context: bool = False, ) -> AsyncIterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]: """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" channel_specs: dict[str, BaseChannel] = {} @@ -68,22 +54,12 @@ async def AsyncChannelsManager( for k, v in specs.items(): 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: # managed: create enter tasks with reference to spec, await them if tasks := { - asyncio.create_task( - stack.enter_async_context( - value.cls.aenter(loop, **value.kwargs) - if isinstance(value, ConfiguredManagedValue) - else value.aenter(loop) - ) - ): key + asyncio.create_task(stack.enter_async_context(value.aenter(loop))): key for key, value in managed_specs.items() }: done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED) @@ -98,8 +74,3 @@ async def AsyncChannelsManager( # managed: build mapping from spec to result ManagedValueMapping({tasks[task]: task.result() for task in done}), ) - - -@contextmanager -def noop_context() -> Iterator[None]: - yield None diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index c78dd30bc..e0b41f48a 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2,19 +2,15 @@ import json import operator import re import time -from collections.abc import Iterator -from contextlib import contextmanager from dataclasses import replace from typing import Annotated, Any, Literal, Optional, Union, cast -import httpx import pytest from langchain_core.runnables import RunnableConfig, RunnableMap, RunnablePick from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict -from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import BaseCheckpointSaver @@ -23,7 +19,6 @@ from langgraph.errors import NodeInterrupt from langgraph.graph import StateGraph from langgraph.graph.graph import Graph from langgraph.graph.message import MessageGraph, MessagesState, add_messages -from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, Pregel @@ -1440,39 +1435,14 @@ def test_conditional_state_graph( checkpointer: BaseCheckpointSaver = request.getfixturevalue( f"checkpointer_{checkpointer_name}" ) - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() class AgentState(TypedDict, total=False): input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - session: Annotated[httpx.Client, Context(make_httpx_client)] class ToolState(TypedDict, total=False): agent_outcome: Union[AgentAction, AgentFinish] - session: Annotated[httpx.Client, Context(make_httpx_client)] # Assemble the tools @tool() @@ -1514,7 +1484,6 @@ def test_conditional_state_graph( # Define tool execution logic def execute_tools(data: ToolState) -> dict: # check session in data - assert isinstance(data["session"], httpx.Client) assert "input" not in data assert "intermediate_steps" not in data # execute the tool @@ -1527,7 +1496,6 @@ def test_conditional_state_graph( # Define decision-making logic def should_continue(data: AgentState) -> str: # check session in data - assert isinstance(data["session"], httpx.Client) # Logic to decide whether to continue in the loop or exit if isinstance(data["agent_outcome"], AgentFinish): return "exit" @@ -1556,88 +1524,86 @@ def test_conditional_state_graph( assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot assert app.get_graph().draw_mermaid(with_styles=False) == snapshot - with assert_ctx_once(): - assert app.invoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], + assert app.invoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } - with assert_ctx_once(): - assert [*app.stream({"input": "what is weather in sf"})] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] + assert [*app.stream({"input": "what is weather in sf"})] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] # test state get/update methods with interrupt_after @@ -1647,22 +1613,20 @@ def test_conditional_state_graph( ) config = {"configurable": {"thread_id": "1"}} - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - {"__interrupt__": ()}, - ] + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -1704,17 +1668,16 @@ def test_conditional_state_graph( interrupts=(), ) - with assert_ctx_once(): - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -1758,44 +1721,42 @@ def test_conditional_state_graph( interrupts=(), ) - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - {"__interrupt__": ()}, - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] - with assert_ctx_once(): - app_w_interrupt.update_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) assert app_w_interrupt.get_state(config) == StateSnapshot( values={ @@ -2749,7 +2710,6 @@ 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: @@ -2794,7 +2754,6 @@ def test_state_graph_packets( ) def agent(data: AgentState) -> AgentState: - assert isinstance(data["session"], httpx.Client) return { "messages": model.invoke(data["messages"]), "something_extra": "hi there", @@ -2802,7 +2761,6 @@ def test_state_graph_packets( # Define decision-making logic def should_continue(data: dict) -> str: - assert isinstance(data["session"], httpx.Client) assert data["something_extra"] == "hi there", ( "nodes can pass extra data to their cond edges, which isn't saved in state" ) @@ -6258,28 +6216,12 @@ def test_start_branch_then( class State(TypedDict): my_key: Annotated[str, operator.add] market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} + return {"my_key": " slow"} def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} + return {"my_key": " fast"} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two_slow", tool_two_slow) diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 4007636a7..91e4886ef 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2,27 +2,21 @@ import asyncio import operator import re import sys -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager from typing import ( Annotated, - Any, Literal, Optional, Union, cast, ) -import httpx import pytest from langchain_core.messages import ToolCall from langchain_core.runnables import RunnableConfig, RunnablePick -from pydantic import BaseModel from pytest_mock import MockerFixture from syrupy import SnapshotAssertion from typing_extensions import TypedDict -from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import BaseCheckpointSaver @@ -30,7 +24,6 @@ from langgraph.constants import END, PULL, PUSH, START from langgraph.graph.graph import Graph from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph -from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, Pregel @@ -1529,51 +1522,16 @@ async def test_conditional_graph(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph_state( - mocker: MockerFixture, checkpointer_name: str -) -> None: +async def test_conditional_graph_state(checkpointer_name: str) -> None: from langchain_core.agents import AgentAction, AgentFinish from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate from langchain_core.tools import tool - setup = mocker.Mock() - teardown = mocker.Mock() - - @asynccontextmanager - async def assert_ctx_once() -> AsyncIterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - class MyPydanticContextModel(BaseModel, arbitrary_types_allowed=True): - session: httpx.AsyncClient - something_else: str - - @asynccontextmanager - async def make_context( - config: RunnableConfig, - ) -> AsyncIterator[MyPydanticContextModel]: - assert isinstance(config, dict) - setup() - session = httpx.AsyncClient() - try: - yield MyPydanticContextModel(session=session, something_else="hello") - finally: - await session.aclose() - teardown() - class AgentState(TypedDict): input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] - context: Annotated[MyPydanticContextModel, Context(make_context)] # Assemble the tools @tool() @@ -1614,8 +1572,6 @@ async def test_conditional_graph_state( # Define tool execution logic def execute_tools(data: AgentState) -> dict: - # check we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) # execute the tool agent_action: AgentAction = data.pop("agent_outcome") observation = {t.name: t for t in tools}[agent_action.tool].invoke( @@ -1625,8 +1581,6 @@ async def test_conditional_graph_state( # Define decision-making logic def should_continue(data: AgentState) -> str: - # check we have httpx session in AgentState - assert isinstance(data["context"], MyPydanticContextModel) # Logic to decide whether to continue in the loop or exit if isinstance(data["agent_outcome"], AgentFinish): return "exit" @@ -1649,91 +1603,88 @@ async def test_conditional_graph_state( app = workflow.compile() - async with assert_ctx_once(): - assert await app.ainvoke({"input": "what is weather in sf"}) == { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], + assert await app.ainvoke({"input": "what is weather in sf"}) == { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } - async with assert_ctx_once(): - assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] + assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], - ], - } - }, - { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] - async with assert_ctx_once(): - patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] + patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] patch_paths = {op["path"] for log in patches for op in log.ops} # Check that agent (one of the nodes) has its output streamed to the logs @@ -1774,24 +1725,23 @@ async def test_conditional_graph_state( ) config = {"configurable": {"thread_id": "1"}} - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - {"__interrupt__": ()}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ @@ -1837,17 +1787,16 @@ async def test_conditional_graph_state( interrupts=(), ) - async with assert_ctx_once(): - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ @@ -1893,44 +1842,42 @@ async def test_conditional_graph_state( interrupts=(), ) - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - {"__interrupt__": ()}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] - async with assert_ctx_once(): - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ @@ -2540,7 +2487,6 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] - session: Annotated[httpx.AsyncClient, Context(httpx.AsyncClient)] @tool() def search_api(query: str) -> str: @@ -2585,7 +2531,6 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: # Define decision-making logic def should_continue(data: AgentState) -> str: - assert isinstance(data["session"], httpx.AsyncClient) # 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] @@ -4077,29 +4022,12 @@ async def test_start_branch_then(checkpointer_name: str) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str - shared: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - other: Annotated[dict[str, dict[str, Any]], SharedValue.on("assistant_id")] - - def assert_shared_value(data: State, config: RunnableConfig) -> State: - assert "shared" in data - if thread_id := config["configurable"].get("thread_id"): - if thread_id == "1": - # this is the first thread, so should not see a value - assert data["shared"] == {} - return {"shared": {"1": {"hello": "world"}}, "other": {"2": {1: 2}}} - elif thread_id == "2": - # this should get value saved by thread 1 - assert data["shared"] == {"1": {"hello": "world"}} - elif thread_id == "3": - # this is a different assistant, so should not see previous value - assert data["shared"] == {} - return {} def tool_two_slow(data: State, config: RunnableConfig) -> State: - return {"my_key": " slow", **assert_shared_value(data, config)} + return {"my_key": " slow"} def tool_two_fast(data: State, config: RunnableConfig) -> State: - return {"my_key": " fast", **assert_shared_value(data, config)} + return {"my_key": " fast"} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two_slow", tool_two_slow) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index d5c02a911..0f66c63be 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9,14 +9,12 @@ import time import uuid import warnings from collections import Counter, deque -from collections.abc import Generator, Iterator, Sequence +from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager from dataclasses import dataclass, field from random import randrange from typing import Annotated, Any, Literal, Optional, Union, get_type_hints -import httpx import pytest from langchain_core.language_models import GenericFakeChatModel from langchain_core.runnables import ( @@ -34,7 +32,6 @@ from typing_extensions import TypedDict from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate -from langgraph.channels.context import Context from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic @@ -2013,52 +2010,6 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: Pregel(nodes={"one": one, "two": two}) -def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: - setup = mocker.Mock() - cleanup = mocker.Mock() - - @contextmanager - def an_int() -> Generator[int, None, None]: - setup() - try: - yield 5 - finally: - cleanup() - - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).batch - | Channel.write_to("output").batch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "ctx": Context(an_int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels=["inbox", "output"], - stream_channels=["inbox", "output"], - ) - - assert setup.call_count == 0 - assert cleanup.call_count == 0 - for i, chunk in enumerate(app.stream(2)): - assert setup.call_count == 1, "Expected setup to be called once" - if i == 0: - assert chunk == {"inbox": [3]} - elif i == 1: - assert chunk == {"output": 4} - else: - assert False, "Expected only two chunks" - assert cleanup.call_count == 1, "Expected cleanup to be called once" - - def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: def left(data: str) -> str: return data + "->left" @@ -3139,34 +3090,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion, - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str, ) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - setup = mocker.Mock() - teardown = mocker.Mock() - - @contextmanager - def assert_ctx_once() -> Iterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @contextmanager - def make_httpx_client() -> Iterator[httpx.Client]: - setup() - with httpx.Client() as client: - try: - yield client - finally: - teardown() def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -3187,7 +3114,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( inner: Annotated[InnerObject, lambda x, y: y] answer: Optional[str] = None docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.Client, Context(make_httpx_client)] class StateUpdate(BaseModel): query: Optional[str] = None @@ -3251,25 +3177,21 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( assert app.get_input_schema().model_json_schema() == snapshot assert app.get_output_schema().model_json_schema() == snapshot - with pytest.raises(ValidationError), assert_ctx_once(): + with pytest.raises(ValidationError): app.invoke({"query": {}}) - with assert_ctx_once(): - assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } + assert app.invoke({"query": "what is weather in sf", "inner": {"yo": 1}}) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } - with assert_ctx_once(): - assert [ - *app.stream({"query": "what is weather in sf", "inner": {"yo": 1}}) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [*app.stream({"query": "what is weather in sf", "inner": {"yo": 1}})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] app_w_interrupt = workflow.compile( checkpointer=checkpointer, @@ -3277,35 +3199,32 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( ) config = {"configurable": {"thread_id": "1"}} - with assert_ctx_once(): - assert [ - c - for c in app_w_interrupt.stream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + assert [ + c + for c in app_w_interrupt.stream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - with assert_ctx_once(): - assert [c for c in app_w_interrupt.stream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - with assert_ctx_once(): - assert app_w_interrupt.update_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } + assert app_w_interrupt.update_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", } + } @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 739de4312..4cbdc465e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8,8 +8,6 @@ import random import sys import uuid from collections import Counter, deque -from collections.abc import AsyncGenerator, AsyncIterator, Generator -from contextlib import asynccontextmanager, contextmanager from dataclasses import replace from time import perf_counter from typing import ( @@ -21,7 +19,6 @@ from typing import ( ) from uuid import UUID -import httpx import pytest from langchain_core.language_models import GenericFakeChatModel from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough @@ -34,7 +31,6 @@ from typing_extensions import TypedDict from langgraph.cache.base import BaseCache from langgraph.channels.base import BaseChannel 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.checkpoint.base import ( @@ -4313,75 +4309,6 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: assert await app.ainvoke(2) is None -async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: - setup_sync = mocker.Mock() - cleanup_sync = mocker.Mock() - setup_async = mocker.Mock() - cleanup_async = mocker.Mock() - - @contextmanager - def an_int() -> Generator[int, None, None]: - setup_sync() - try: - yield 5 - finally: - cleanup_sync() - - @asynccontextmanager - async def an_int_async() -> AsyncGenerator[int, None]: - setup_async() - try: - yield 5 - finally: - cleanup_async() - - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = ( - Channel.subscribe_to("inbox") - | RunnableLambda(add_one).abatch - | Channel.write_to("output").abatch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "input": LastValue(int), - "output": LastValue(int), - "inbox": Topic(int), - "ctx": Context(an_int, an_int_async), - }, - input_channels="input", - output_channels=["inbox", "output"], - stream_channels=["inbox", "output"], - ) - - async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: - i = 0 - async for chunk in aiter: - yield i, chunk - i += 1 - - assert setup_sync.call_count == 0 - assert cleanup_sync.call_count == 0 - assert setup_async.call_count == 0 - assert cleanup_async.call_count == 0 - async for i, chunk in aenumerate(app.astream(2)): - assert setup_sync.call_count == 0, "Sync context manager should not be used" - assert cleanup_sync.call_count == 0, "Sync context manager should not be used" - assert setup_async.call_count == 1, "Expected setup to be called once" - if i == 0: - assert chunk == {"inbox": [3]} - elif i == 1: - assert chunk == {"output": 4} - else: - assert False, "Expected only two chunks" - assert setup_sync.call_count == 0 - assert cleanup_sync.call_count == 0 - assert setup_async.call_count == 1, "Expected setup to be called once" - assert cleanup_async.call_count == 1, "Expected cleanup to be called once" - - async def test_conditional_entrypoint_graph() -> None: async def left(data: str) -> str: return data + "->left" @@ -4959,32 +4886,8 @@ async def test_nested_pydantic_models() -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( - snapshot: SnapshotAssertion, mocker: MockerFixture, checkpointer_name: str + checkpointer_name: str, ) -> None: - setup = mocker.Mock() - teardown = mocker.Mock() - - @asynccontextmanager - async def assert_ctx_once() -> AsyncIterator[None]: - assert setup.call_count == 0 - assert teardown.call_count == 0 - try: - yield - finally: - assert setup.call_count == 1 - assert teardown.call_count == 1 - setup.reset_mock() - teardown.reset_mock() - - @asynccontextmanager - async def make_httpx_client() -> AsyncIterator[httpx.AsyncClient]: - setup() - async with httpx.AsyncClient() as client: - try: - yield client - finally: - teardown() - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -5000,7 +4903,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( query: str answer: Optional[str] = None docs: Annotated[list[str], sorted_add] - client: Annotated[httpx.AsyncClient, Context(make_httpx_client)] class Input(BaseModel): query: str @@ -5053,24 +4955,21 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( app = workflow.compile() - async with assert_ctx_once(): - with pytest.raises(ValidationError): - await app.ainvoke({"query": {}}) + with pytest.raises(ValidationError): + await app.ainvoke({"query": {}}) - async with assert_ctx_once(): - assert await app.ainvoke({"query": "what is weather in sf"}) == { - "docs": ["doc1", "doc2", "doc3", "doc4"], - "answer": "doc1,doc2,doc3,doc4", - } + assert await app.ainvoke({"query": "what is weather in sf"}) == { + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + } - async with assert_ctx_once(): - assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] async with awith_checkpointer(checkpointer_name) as checkpointer: app_w_interrupt = workflow.compile( @@ -5079,24 +4978,22 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( ) config = {"configurable": {"thread_id": "1"}} - async with assert_ctx_once(): - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config + ) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - async with assert_ctx_once(): - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ @@ -5135,16 +5032,15 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( interrupts=(), ) - async with assert_ctx_once(): - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", } + } @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 011dab41e..84616c4f1 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -11,7 +11,6 @@ from pydantic import BaseModel from typing_extensions import NotRequired, Required, TypedDict from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema -from langgraph.managed.shared_value import SharedValue class State(BaseModel): @@ -153,9 +152,6 @@ def test_state_schema_optional_values(total_: bool): class State(InputState): # this would be ignored val4: dict - some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field( - default="foo" - ) builder = StateGraph(State, input=InputState, output=OutputState) builder.add_node("n", lambda x: x) @@ -219,9 +215,6 @@ def test_state_schema_default_values(kw_only_: bool): val11: Annotated[list[str], "annotated list"] = field( default_factory=lambda: ["a", "b"] ) - some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field( - default="foo" - ) builder = StateGraph(InputState) builder.add_node("n", lambda x: x) @@ -247,58 +240,6 @@ def test_state_schema_default_values(kw_only_: bool): ) -def test_raises_invalid_managed(): - class BadInputState(TypedDict): - some_thing: str - some_input_channel: Annotated[str, SharedValue.on("assistant_id")] - - class InputState(TypedDict): - some_thing: str - some_input_channel: str - - class BadOutputState(TypedDict): - some_thing: str - some_output_channel: Annotated[str, SharedValue.on("assistant_id")] - - class OutputState(TypedDict): - some_thing: str - some_output_channel: str - - class State(TypedDict): - some_thing: str - some_channel: Annotated[str, SharedValue.on("assistant_id")] - - # All OK - StateGraph(State, input=InputState, output=OutputState) - StateGraph(State) - StateGraph(State, input=State, output=State) - StateGraph(State, input=InputState) - StateGraph(State, input=InputState) - - bad_input_examples = [ - (State, BadInputState, OutputState), - (State, BadInputState, BadOutputState), - (State, BadInputState, State), - (State, BadInputState, None), - ] - for _state, _inp, _outp in bad_input_examples: - with pytest.raises( - ValueError, - match="Invalid managed channels detected in BadInputState: some_input_channel. Managed channels are not permitted in Input/Output schema.", - ): - StateGraph(_state, input=_inp, output=_outp) - bad_output_examples = [ - (State, InputState, BadOutputState), - (State, None, BadOutputState), - ] - for _state, _inp, _outp in bad_output_examples: - with pytest.raises( - ValueError, - match="Invalid managed channels detected in BadOutputState: some_output_channel. Managed channels are not permitted in Input/Output schema.", - ): - StateGraph(_state, input=_inp, output=_outp) - - def test__get_node_name() -> None: # default runnable name assert _get_node_name(RunnableLambda(func=lambda x: x)) == "RunnableLambda"