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

lib: Context values never stored in checkpoints

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