WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-01 04:39:01 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove Context channel / managed value, Remove SharedValue (#4857)

Browse Source
This commit is contained in:
Nuno Campos
2025-05-28 14:02:12 -07:00
committed by GitHub
parent efb282a197 889b40e7a8
commit 2e4645ce61
17 changed files with 375 additions and 1098 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/Makefile
+1 -1
View File
@@ -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
libs/langgraph/langgraph/channels/__init__.py
-2
View File
@@ -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",
libs/langgraph/langgraph/channels/context.py
-5
View File
@@ -1,5 +0,0 @@
from langgraph.managed.context import Context as ContextManagedValue
Context = ContextManagedValue.of
__all__ = ["Context"]
libs/langgraph/langgraph/graph/state.py
+1 -13
View File
@@ -65,12 +65,8 @@ from langgraph.graph.graph import (
Send,
)
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 (
@@ -726,9 +722,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(
@@ -1207,12 +1201,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
libs/langgraph/langgraph/managed/base.py
+3 -42
View File
@@ -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]
libs/langgraph/langgraph/managed/context.py
-111
View File
@@ -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
libs/langgraph/langgraph/managed/shared_value.py
-120
View File
@@ -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))
libs/langgraph/langgraph/pregel/__init__.py
+2 -6
View File
@@ -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(
libs/langgraph/langgraph/pregel/algo.py
+6 -6
View File
@@ -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(
libs/langgraph/langgraph/pregel/draw.py
+2 -3
View File
@@ -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
libs/langgraph/langgraph/pregel/loop.py
+4 -30
View File
@@ -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 []
libs/langgraph/langgraph/pregel/manager.py
+2 -31
View File
@@ -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
libs/langgraph/tests/test_large_cases.py
+138 -196
View File
@@ -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)
libs/langgraph/tests/test_large_cases_async.py
+143 -215
View File
@@ -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)
libs/langgraph/tests/test_pregel.py
+36 -117
View File
@@ -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)
libs/langgraph/tests/test_pregel_async.py
+37 -141
View File
@@ -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)
libs/langgraph/tests/test_state.py
-59
View File
@@ -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"