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

Add support for using Context channel in StateGraph (#761)

Browse Source 
* Add support for using Context channel in StateGraph

* Fix handling of asynccontextmanager funcs
This commit is contained in:
Nuno Campos
2024-06-21 17:49:52 -07:00
committed by GitHub
parent 571f1bac27
commit 99fa57666e
17 changed files with 223 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/channels/any_value.py
+2 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Generator, Generic, Optional, Sequence, Type
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -32,7 +33,7 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[Value] = None
self, checkpoint: Optional[Value], config: RunnableConfig
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ)
if checkpoint is not None:
libs/langgraph/langgraph/channels/base.py
+4 -3
View File
@@ -10,6 +10,7 @@ from typing import (
TypeVar,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.errors import EmptyChannelError, InvalidUpdateError
@@ -41,18 +42,18 @@ class BaseChannel(Generic[Value, Update, C], ABC):
@contextmanager
@abstractmethod
def from_checkpoint(
self, checkpoint: Optional[C] = None
self, checkpoint: Optional[C], config: RunnableConfig
) -> Generator[Self, None, None]:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
@asynccontextmanager
async def afrom_checkpoint(
self, checkpoint: Optional[C] = None
self, checkpoint: Optional[C], config: RunnableConfig
) -> AsyncGenerator[Self, None]:
"""Return a new identical channel, optionally initialized from a checkpoint.
If the checkpoint contains complex data structures, they should be copied."""
with self.from_checkpoint(checkpoint) as value:
with self.from_checkpoint(checkpoint, config) as value:
yield value
# state methods
libs/langgraph/langgraph/channels/binop.py
+2 -1
View File
@@ -9,6 +9,7 @@ from typing import (
Type,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import NotRequired, Required, Self
from langgraph.channels.base import BaseChannel, Value
@@ -72,7 +73,7 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[Value] = None
self, checkpoint: Optional[Value], config: RunnableConfig
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.operator)
if checkpoint is not None:
libs/langgraph/langgraph/channels/context.py
+40 -32
View File
@@ -1,17 +1,19 @@
from contextlib import asynccontextmanager, contextmanager
from inspect import signature
from typing import (
Any,
AsyncContextManager,
AsyncGenerator,
Callable,
ContextManager,
Generator,
Generic,
Optional,
Sequence,
Type,
Union,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -35,69 +37,75 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
def __init__(
self,
ctx: Optional[Callable[[], ContextManager[Value]]] = None,
actx: Optional[Callable[[], AsyncContextManager[Value]]] = None,
typ: Optional[Type[Value]] = None,
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.typ = typ
self.ctx = ctx
self.actx = actx
@property
def ValueType(self) -> Any:
"""The type of the value stored in the channel."""
return (
self.typ
or (self.ctx if hasattr(self.ctx, "__enter__") else None)
or (self.actx if hasattr(self.actx, "__aenter__") else None)
or None
)
return None
@property
def UpdateType(self) -> Type[None]:
"""The type of the update received by the channel."""
raise InvalidUpdateError()
return None
def checkpoint(self) -> None:
raise EmptyChannelError()
@contextmanager
def from_checkpoint(self, checkpoint: None = None) -> Generator[Self, None, None]:
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, typ=self.typ)
# ContextManager doesn't have a checkpoint
ctx = self.ctx()
empty.value = ctx.__enter__()
try:
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
finally:
ctx.__exit__(None, None, None)
@asynccontextmanager
async def afrom_checkpoint(
self, checkpoint: Optional[str] = None
self, checkpoint: None, config: RunnableConfig
) -> AsyncGenerator[Self, None]:
empty = self.__class__(ctx=self.ctx, actx=self.actx)
if self.actx is not None:
empty = self.__class__(ctx=self.ctx, actx=self.actx, typ=self.typ)
# ContextManager doesn't have a checkpoint
actx = self.actx()
empty.value = await actx.__aenter__()
try:
yield empty
finally:
await actx.__aexit__(None, None, None)
ctx = (
self.actx(config)
if signature(self.actx).parameters.get("config")
else self.actx()
)
else:
with self.from_checkpoint() as empty:
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()
raise InvalidUpdateError("Context channel does not accept writes.")
return False
def get(self) -> Value:
libs/langgraph/langgraph/channels/dynamic_barrier_value.py
+4 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Generator, Generic, NamedTuple, Optional, Sequence, Type, Union
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -46,7 +47,9 @@ class DynamicBarrierValue(
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[tuple[Optional[set[Value]], set[Value]]] = None
self,
checkpoint: Optional[tuple[Optional[set[Value]], set[Value]]],
config: RunnableConfig,
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ)
if checkpoint is not None:
libs/langgraph/langgraph/channels/ephemeral_value.py
+2 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Generator, Generic, Optional, Sequence, Type
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -32,7 +33,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[Value] = None
self, checkpoint: Optional[Value], config: RunnableConfig
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.guard)
if checkpoint is not None:
libs/langgraph/langgraph/channels/last_value.py
+2 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Generator, Generic, Optional, Sequence, Type
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -31,7 +32,7 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[Value] = None
self, checkpoint: Optional[Value], config: RunnableConfig
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ)
if checkpoint is not None:
libs/langgraph/langgraph/channels/manager.py
+19 -20
View File
@@ -1,7 +1,9 @@
from contextlib import asynccontextmanager, contextmanager
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from datetime import datetime, timezone
from typing import Any, AsyncGenerator, Generator, Mapping
from langchain_core.runnables import RunnableConfig
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
from langgraph.checkpoint.id import uuid6
@@ -12,35 +14,32 @@ from langgraph.errors import EmptyChannelError
def ChannelsManager(
channels: Mapping[str, BaseChannel],
checkpoint: Checkpoint,
config: RunnableConfig,
) -> Generator[Mapping[str, BaseChannel], None, None]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
# TODO use https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
empty = {
k: v.from_checkpoint(checkpoint["channel_values"].get(k))
for k, v in channels.items()
}
try:
yield {k: v.__enter__() for k, v in empty.items()}
finally:
for v in empty.values():
v.__exit__(None, None, None)
with ExitStack() as stack:
yield {
k: stack.enter_context(
v.from_checkpoint(checkpoint["channel_values"].get(k), config)
)
for k, v in channels.items()
}
@asynccontextmanager
async def AsyncChannelsManager(
channels: Mapping[str, BaseChannel],
checkpoint: Checkpoint,
config: RunnableConfig,
) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
empty = {
k: v.afrom_checkpoint(checkpoint["channel_values"].get(k))
for k, v in channels.items()
}
try:
yield {k: await v.__aenter__() for k, v in empty.items()}
finally:
for v in empty.values():
await v.__aexit__(None, None, None)
async with AsyncExitStack() as stack:
yield {
k: await stack.enter_async_context(
v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
)
for k, v in channels.items()
}
def create_checkpoint(
libs/langgraph/langgraph/channels/named_barrier_value.py
+2 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Generator, Generic, Optional, Sequence, Type
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -30,7 +31,7 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[set[Value]] = None
self, checkpoint: Optional[set[Value]], config: RunnableConfig
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.names)
if checkpoint is not None:
libs/langgraph/langgraph/channels/topic.py
+4 -1
View File
@@ -1,6 +1,7 @@
from contextlib import contextmanager
from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, Union
from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel, Value
@@ -55,7 +56,9 @@ class Topic(
@contextmanager
def from_checkpoint(
self, checkpoint: Optional[tuple[set[Value], list[Value]]] = None
self,
checkpoint: Optional[tuple[set[Value], list[Value]]],
config: RunnableConfig,
) -> Generator[Self, None, None]:
empty = self.__class__(self.typ, self.unique, self.accumulate)
if checkpoint is not None:
libs/langgraph/langgraph/graph/state.py
+24 -13
View File
@@ -20,6 +20,7 @@ from langchain_core.runnables.base import RunnableLike
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
@@ -240,7 +241,16 @@ class StateGraph(Graph):
# prepare output channels
state_keys = list(self.channels)
output_channels = state_keys[0] if state_keys == ["__root__"] else state_keys
output_channels = (
state_keys[0]
if state_keys == ["__root__"]
else [
key
for key in state_keys
if not isinstance(self.channels[key], Context)
and not is_managed_value(self.channels[key])
]
)
compiled = CompiledStateGraph(
builder=self,
@@ -281,16 +291,7 @@ class CompiledStateGraph(CompiledGraph):
def get_input_schema(
self, config: Optional[RunnableConfig] = None
) -> type[BaseModel]:
if isinstance(self.builder.schema, BaseModel):
return self.builder.schema
return super().get_input_schema(config)
def get_output_schema(self, config: Optional[RunnableConfig] = None) -> BaseModel:
if isinstance(self.builder.schema, BaseModel):
return self.builder.schema
return super().get_output_schema(config)
return self.get_output_schema(config)
def attach_node(self, key: str, node: Optional[Runnable]) -> None:
state_keys = list(self.builder.channels)
@@ -477,11 +478,21 @@ def _get_channel(
return manager
else:
raise ValueError(f"This {annotation} not allowed in this position")
elif channel := _is_field_channel(annotation):
return channel
elif channel := _is_field_binop(annotation):
return channel
return LastValue(annotation)
def _is_field_channel(typ: Type[Any]) -> Optional[BaseChannel]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) >= 1 and isinstance(meta[-1], BaseChannel):
return meta[-1]
return None
def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
@@ -502,8 +513,8 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) == 1:
decoration = get_origin(meta[0]) or meta[0]
if len(meta) >= 1:
decoration = get_origin(meta[-1]) or meta[-1]
if is_managed_value(decoration):
return decoration
libs/langgraph/langgraph/managed/few_shot.py
+4 -2
View File
@@ -69,7 +69,9 @@ class FewShotExamples(ManagedValue[Sequence[V]], Generic[V]):
for example in self.graph.checkpointer.list(
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
):
with ChannelsManager(self.graph.channels, example.checkpoint) as channels:
with ChannelsManager(
self.graph.channels, example.checkpoint, self.config
) as channels:
yield read_channels(channels, self.graph.output_channels)
async def aiter(self, score: int = 1) -> AsyncIterator[V]:
@@ -77,7 +79,7 @@ class FewShotExamples(ManagedValue[Sequence[V]], Generic[V]):
None, filter={"score": score, **self.metadata_filter_dict}, limit=self.k
):
async with AsyncChannelsManager(
self.graph.channels, example.checkpoint
self.graph.channels, example.checkpoint, self.config
) as channels:
yield read_channels(channels, self.graph.output_channels)
libs/langgraph/langgraph/pregel/__init__.py
+31 -17
View File
@@ -55,6 +55,7 @@ from langgraph.channels.base import (
BaseChannel,
EmptyChannelError,
)
from langgraph.channels.context import Context
from langgraph.channels.manager import (
AsyncChannelsManager,
ChannelsManager,
@@ -335,7 +336,9 @@ class Pregel(
@property
def stream_channels_asis(self) -> Union[str, Sequence[str]]:
return self.stream_channels or [k for k in self.channels]
return self.stream_channels or [
k for k in self.channels if not isinstance(self.channels[k], Context)
]
@property
def managed_values_dict(self) -> dict[str, ManagedValueSpec]:
@@ -356,7 +359,7 @@ class Pregel(
checkpoint = saved.checkpoint if saved else empty_checkpoint()
config = saved.config if saved else config
with ChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
@@ -388,7 +391,7 @@ class Pregel(
config = saved.config if saved else config
async with AsyncChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
@@ -430,7 +433,7 @@ class Pregel(
config, before=before, limit=limit, filter=filter
):
with ChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
@@ -475,7 +478,7 @@ class Pregel(
parent_config,
) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
async with AsyncChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
@@ -537,7 +540,7 @@ class Pregel(
if as_node is None:
raise InvalidUpdateError("Ambiguous update, specify as_node")
# update channels
with ChannelsManager(self.channels, checkpoint) as channels:
with ChannelsManager(self.channels, checkpoint, config) as channels:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
@@ -560,7 +563,7 @@ class Pregel(
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, task.writes
_local_read, checkpoint, channels, task.writes, config
),
},
),
@@ -625,7 +628,7 @@ class Pregel(
if as_node is None:
raise InvalidUpdateError("Ambiguous update, specify as_node")
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint) as channels:
async with AsyncChannelsManager(self.channels, checkpoint, config) as channels:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
@@ -648,7 +651,7 @@ class Pregel(
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, task.writes
_local_read, checkpoint, channels, task.writes, config
),
},
),
@@ -790,7 +793,7 @@ class Pregel(
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
with ChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, get_executor_for_config(
config
) as executor, ManagedValuesManager(
@@ -957,7 +960,7 @@ class Pregel(
task = futures.pop(fut)
if fut.exception() is not None:
# we got an exception, break out of while loop
# exception will be handle in panic_or_proceed
# exception will be handled in panic_or_proceed
futures.clear()
else:
# yield updates output for the finished task
@@ -1143,7 +1146,7 @@ class Pregel(
start = saved.metadata.get("step", -2) + 1 if saved else -1
# create channels from checkpoint
async with AsyncChannelsManager(
self.channels, checkpoint
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, config, self
) as managed:
@@ -1579,14 +1582,21 @@ def _local_read(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
writes: Sequence[tuple[str, Any]],
config: RunnableConfig,
select: Union[list[str], str],
fresh: bool = False,
) -> Union[dict[str, Any], Any]:
if fresh:
checkpoint = create_checkpoint(checkpoint, channels, -1)
with ChannelsManager(channels, checkpoint) as channels:
_apply_writes(copy_checkpoint(checkpoint), channels, writes, None)
return read_channels(channels, select)
context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
with ChannelsManager(
{k: v for k, v in channels.items() if k not in context_channels},
checkpoint,
config,
) as channels:
all_channels = {**channels, **context_channels}
_apply_writes(copy_checkpoint(checkpoint), all_channels, writes, None)
return read_channels(all_channels, select)
else:
return read_channels(channels, select)
@@ -1739,7 +1749,7 @@ def _prepare_next_tasks(
_local_write, writes.extend, processes, channels
),
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, tasks
_local_read, checkpoint, channels, tasks, config
),
},
),
@@ -1819,7 +1829,11 @@ def _prepare_next_tasks(
_local_write, writes.extend, processes, channels
),
CONFIG_KEY_READ: partial(
_local_read, checkpoint, channels, writes
_local_read,
checkpoint,
channels,
writes,
config,
),
},
),
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+6 -6
View File
File diff suppressed because one or more lines are too long
libs/langgraph/tests/test_channels.py
+42 -38
View File
@@ -4,6 +4,7 @@ from typing import AsyncGenerator, Generator, Sequence, Union
import httpx
import pytest
from langchain_core.runnables import RunnableConfig
from pytest_mock import MockerFixture
from langgraph.channels.binop import BinaryOperatorAggregate
@@ -14,7 +15,7 @@ from langgraph.errors import EmptyChannelError, InvalidUpdateError
def test_last_value() -> None:
with LastValue(int).from_checkpoint() as channel:
with LastValue(int).from_checkpoint(None, {}) as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -28,12 +29,12 @@ def test_last_value() -> None:
channel.update([4])
assert channel.get() == 4
checkpoint = channel.checkpoint()
with LastValue(int).from_checkpoint(checkpoint) as channel:
with LastValue(int).from_checkpoint(checkpoint, {}) as channel:
assert channel.get() == 4
async def test_last_value_async() -> None:
async with LastValue(int).afrom_checkpoint() as channel:
async with LastValue(int).afrom_checkpoint(None, {}) as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -47,12 +48,12 @@ async def test_last_value_async() -> None:
channel.update([4])
assert channel.get() == 4
checkpoint = channel.checkpoint()
async with LastValue(int).afrom_checkpoint(checkpoint) as channel:
async with LastValue(int).afrom_checkpoint(checkpoint, {}) as channel:
assert channel.get() == 4
def test_topic() -> None:
with Topic(str).from_checkpoint() as channel:
with Topic(str).from_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -67,16 +68,16 @@ def test_topic() -> None:
assert channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
with Topic(str).from_checkpoint(checkpoint) as channel:
with Topic(str).from_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["e"]
with Topic(str).from_checkpoint(checkpoint) as channel_copy:
with Topic(str).from_checkpoint(checkpoint, {}) as channel_copy:
channel_copy.update(["f"])
assert channel_copy.get() == ["f"]
assert channel.get() == ["e"]
async def test_topic_async() -> None:
async with Topic(str).afrom_checkpoint() as channel:
async with Topic(str).afrom_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -91,12 +92,12 @@ async def test_topic_async() -> None:
assert channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
async with Topic(str).afrom_checkpoint(checkpoint) as channel:
async with Topic(str).afrom_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["e"]
def test_topic_unique() -> None:
with Topic(str, unique=True).from_checkpoint() as channel:
with Topic(str, unique=True).from_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -111,14 +112,14 @@ def test_topic_unique() -> None:
assert channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
with Topic(str, unique=True).from_checkpoint(checkpoint) as channel:
with Topic(str, unique=True).from_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["e"]
assert channel.update(["d", "f"])
assert channel.get() == ["f"], "de-dupes from checkpoint"
async def test_topic_unique_async() -> None:
async with Topic(str, unique=True).afrom_checkpoint() as channel:
async with Topic(str, unique=True).afrom_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -133,14 +134,14 @@ async def test_topic_unique_async() -> None:
assert channel.update(["e"])
assert channel.get() == ["e"]
checkpoint = channel.checkpoint()
async with Topic(str, unique=True).afrom_checkpoint(checkpoint) as channel:
async with Topic(str, unique=True).afrom_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["e"]
assert channel.update(["d", "f"])
assert channel.get() == ["f"], "de-dupes from checkpoint"
def test_topic_accumulate() -> None:
with Topic(str, accumulate=True).from_checkpoint() as channel:
with Topic(str, accumulate=True).from_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -151,14 +152,14 @@ def test_topic_accumulate() -> None:
assert not channel.update([])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
checkpoint = channel.checkpoint()
with Topic(str, accumulate=True).from_checkpoint(checkpoint) as channel:
with Topic(str, accumulate=True).from_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
assert channel.update(["e"])
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
async def test_topic_accumulate_async() -> None:
async with Topic(str, accumulate=True).afrom_checkpoint() as channel:
async with Topic(str, accumulate=True).afrom_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -169,14 +170,14 @@ async def test_topic_accumulate_async() -> None:
assert not channel.update([])
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
checkpoint = channel.checkpoint()
async with Topic(str, accumulate=True).afrom_checkpoint(checkpoint) as channel:
async with Topic(str, accumulate=True).afrom_checkpoint(checkpoint, {}) as channel:
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
assert channel.update(["e"])
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
def test_topic_unique_accumulate() -> None:
with Topic(str, unique=True, accumulate=True).from_checkpoint() as channel:
with Topic(str, unique=True, accumulate=True).from_checkpoint(None, {}) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -189,7 +190,7 @@ def test_topic_unique_accumulate() -> None:
assert channel.get() == ["a", "b", "c", "d"]
checkpoint = channel.checkpoint()
with Topic(str, unique=True, accumulate=True).from_checkpoint(
checkpoint
checkpoint, {}
) as channel:
assert channel.get() == ["a", "b", "c", "d"]
assert channel.update(["d", "e"])
@@ -197,7 +198,9 @@ def test_topic_unique_accumulate() -> None:
async def test_topic_unique_accumulate_async() -> None:
async with Topic(str, unique=True, accumulate=True).afrom_checkpoint() as channel:
async with Topic(str, unique=True, accumulate=True).afrom_checkpoint(
None, {}
) as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, list[str]]
@@ -209,7 +212,7 @@ async def test_topic_unique_accumulate_async() -> None:
assert channel.get() == ["a", "b", "c", "d"]
checkpoint = channel.checkpoint()
async with Topic(str, unique=True, accumulate=True).afrom_checkpoint(
checkpoint
checkpoint, {}
) as channel:
assert channel.get() == ["a", "b", "c", "d"]
channel.update(["d", "e"])
@@ -217,7 +220,9 @@ async def test_topic_unique_accumulate_async() -> None:
def test_binop() -> None:
with BinaryOperatorAggregate(int, operator.add).from_checkpoint() as channel:
with BinaryOperatorAggregate(int, operator.add).from_checkpoint(
None, {}
) as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -229,13 +234,15 @@ def test_binop() -> None:
assert channel.get() == 10
checkpoint = channel.checkpoint()
with BinaryOperatorAggregate(int, operator.add).from_checkpoint(
checkpoint
checkpoint, {}
) as channel:
assert channel.get() == 10
async def test_binop_async() -> None:
async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint() as channel:
async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint(
None, {}
) as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -247,7 +254,7 @@ async def test_binop_async() -> None:
assert channel.get() == 10
checkpoint = channel.checkpoint()
async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint(
checkpoint
checkpoint, {}
) as channel:
assert channel.get() == 10
@@ -264,13 +271,12 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
finally:
cleanup()
with Context(an_int, None, int).from_checkpoint() as channel:
with Context(an_int, None).from_checkpoint(None, {}) as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
assert channel.ValueType is int
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
assert channel.ValueType is None
assert channel.UpdateType is None
assert channel.get() == 5
@@ -282,10 +288,9 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
with Context(httpx.Client).from_checkpoint() as channel:
assert channel.ValueType is httpx.Client
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is 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)
@@ -301,7 +306,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
cleanup = mocker.Mock()
@contextmanager
def an_int_sync() -> Generator[int, None, None]:
def an_int_sync(config: RunnableConfig) -> Generator[int, None, None]:
try:
yield 5
finally:
@@ -315,13 +320,12 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
finally:
cleanup()
async with Context(an_int_sync, an_int, int).afrom_checkpoint() as channel:
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 int
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
assert channel.ValueType is None
assert channel.UpdateType is None
assert channel.get() == 5
libs/langgraph/tests/test_pregel.py
+8 -1
View File
@@ -18,6 +18,7 @@ from typing import (
Union,
)
import httpx
import pytest
from langchain_core.runnables import (
RunnableConfig,
@@ -1245,7 +1246,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"ctx": Context(an_int, typ=int),
"ctx": Context(an_int),
"output": LastValue(int),
"input": LastValue(int),
},
@@ -2107,6 +2108,7 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
input: str
agent_outcome: Optional[Union[AgentAction, AgentFinish]]
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
session: Annotated[httpx.Client, Context(httpx.Client)]
# Assemble the tools
@tool()
@@ -2147,6 +2149,9 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
# Define tool execution logic
def execute_tools(data: AgentState) -> dict:
# check session in data
assert isinstance(data["session"], httpx.Client)
# execute the tool
agent_action: AgentAction = data.pop("agent_outcome")
observation = {t.name: t for t in tools}[agent_action.tool].invoke(
agent_action.tool_input
@@ -2155,6 +2160,8 @@ def test_conditional_state_graph(snapshot: SnapshotAssertion) -> None:
# 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"
libs/langgraph/tests/test_pregel_async.py
+27 -1
View File
@@ -18,6 +18,7 @@ from typing import (
)
from uuid import UUID
import httpx
import pytest
from langchain_core.runnables import (
RunnableConfig,
@@ -26,6 +27,7 @@ from langchain_core.runnables import (
RunnablePick,
)
from langchain_core.utils.aiter import aclosing
from pydantic import BaseModel
from pytest_mock import MockerFixture
from syrupy import SnapshotAssertion
@@ -1342,7 +1344,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
"input": LastValue(int),
"output": LastValue(int),
"inbox": Topic(int),
"ctx": Context(an_int, an_int_async, typ=int),
"ctx": Context(an_int, an_int_async),
},
input_channels="input",
output_channels=["inbox", "output"],
@@ -2208,10 +2210,29 @@ async def test_conditional_graph_state() -> None:
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool
class MyPydanticContextModel(BaseModel):
class Config:
arbitrary_types_allowed = True
session: httpx.AsyncClient
something_else: str
@asynccontextmanager
async def make_context(
config: RunnableConfig,
) -> AsyncIterator[MyPydanticContextModel]:
assert isinstance(config, dict)
session = httpx.AsyncClient()
try:
yield MyPydanticContextModel(session=session, something_else="hello")
finally:
await session.aclose()
class AgentState(TypedDict):
input: str
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()
@@ -2252,6 +2273,9 @@ async def test_conditional_graph_state() -> None:
# 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(
agent_action.tool_input
@@ -2260,6 +2284,8 @@ async def test_conditional_graph_state() -> None:
# 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"