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

Merge branch 'main' into vb/update-get-state

Browse Source
This commit is contained in:
vbarda
2024-08-21 18:21:28 -04:00
parent e7bc74e918 9147d05cc4
commit 578ec48031
24 changed files with 1007 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/cli/langgraph_cli/docker.py
+15 -2
View File
@@ -12,6 +12,15 @@ DEFAULT_POSTGRES_URI = (
"postgres://postgres:postgres@langgraph-postgres:5432/postgres?sslmode=disable"
)
REDIS = """
langgraph-redis:
image: redis:6
healthcheck:
test: redis-cli ping
interval: 5s
timeout: 1s
retries: 5
"""
DB = """
langgraph-postgres:
@@ -166,18 +175,22 @@ def compose(
compose_str = f"""{volumes}services:
{db}
{REDIS}
{debugger_compose(port=debugger_port, base_url=debugger_base_url)}
langgraph-api:
ports:
- "{port}:8000\""""
- "{port}:8000\"
depends_on:
langgraph-redis:
condition: service_healthy"""
if include_db:
compose_str += """
depends_on:
langgraph-postgres:
condition: service_healthy"""
compose_str += f"""
environment:
POSTGRES_URI: {postgres_uri}
REDIS_URI: redis://langgraph-redis:6379
"""
if capabilities.healthcheck_start_interval:
compose_str += """ healthcheck:
libs/langgraph/langgraph/channels/manager.py
-39
View File
@@ -1,39 +0,0 @@
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from typing import AsyncGenerator, Generator, Mapping
from langchain_core.runnables import RunnableConfig
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
@contextmanager
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)."""
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)."""
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()
}
libs/langgraph/langgraph/constants.py
+2
View File
@@ -6,6 +6,7 @@ INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
CONFIG_KEY_STORE = "__pregel_store"
CONFIG_KEY_RESUMING = "__pregel_resuming"
CONFIG_KEY_TASK_ID = "__pregel_task_id"
INTERRUPT = "__interrupt__"
@@ -18,6 +19,7 @@ RESERVED = {
CONFIG_KEY_SEND,
CONFIG_KEY_READ,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_STORE,
CONFIG_KEY_RESUMING,
CONFIG_KEY_TASK_ID,
INPUT,
libs/langgraph/langgraph/graph/state.py
+39 -17
View File
@@ -44,10 +44,18 @@ from langgraph.graph.graph import (
Graph,
Send,
)
from langgraph.managed.base import ManagedValue, is_managed_value
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.types import All, RetryPolicy
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.utils import RunnableCallable, coerce_to_runnable
logger = logging.getLogger(__name__)
@@ -125,8 +133,8 @@ class StateGraph(Graph):
nodes: dict[str, StateNodeSpec]
channels: dict[str, BaseChannel]
managed: dict[str, Type[ManagedValue]]
schemas: dict[Type[Any], dict[str, Union[BaseChannel, Type[ManagedValue]]]]
managed: dict[str, ManagedValueSpec]
schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
def __init__(
self,
@@ -382,6 +390,8 @@ class StateGraph(Graph):
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
*,
store: Optional[BaseStore] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
@@ -440,7 +450,11 @@ class StateGraph(Graph):
builder=self,
config_type=self.config_schema,
nodes={},
channels={**self.channels, START: EphemeralValue(self.input)},
channels={
**self.channels,
**self.managed,
START: EphemeralValue(self.input),
},
input_channels=START,
stream_mode="updates",
output_channels=output_channels,
@@ -450,6 +464,7 @@ class StateGraph(Graph):
interrupt_after_nodes=interrupt_after,
auto_validate=False,
debug=debug,
store=store,
)
compiled.attach_node(START, None)
@@ -494,7 +509,7 @@ class CompiledStateGraph(CompiledGraph):
**{
k: (self.channels[k].UpdateType, None)
for k in self.builder.schemas[self.builder.input]
if k in self.channels
if isinstance(self.channels[k], BaseChannel)
and not isinstance(self.channels[k], Context)
},
)
@@ -519,7 +534,11 @@ class CompiledStateGraph(CompiledGraph):
if not isinstance(v, Context) and not is_managed_value(v)
]
else:
output_keys = list(self.builder.channels)
output_keys = list(self.builder.channels) + [
k
for k, v in self.builder.managed.items()
if is_writable_managed_value(v)
]
def _get_state_key(
input: Union[None, dict, Any], config: RunnableConfig, *, key: str
@@ -565,10 +584,7 @@ class CompiledStateGraph(CompiledGraph):
)
else:
input_schema = node.input if node else self.builder.schema
input_values = {
k: v if is_managed_value(v) else k
for k, v in self.builder.schemas[input_schema].items()
}
input_values = {k: k for k in self.builder.schemas[input_schema]}
is_single_input = len(input_values) == 1 and "__root__" in input_values
self.channels[key] = EphemeralValue(Any, guard=False)
@@ -687,12 +703,12 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
def _get_channels(
schema: Type[dict],
) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]:
) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec]]:
if not hasattr(schema, "__annotations__"):
return {"__root__": _get_channel(schema, allow_managed=False)}, {}
return {"__root__": _get_channel("__root__", schema, allow_managed=False)}, {}
all_keys = {
name: _get_channel(typ)
name: _get_channel(name, typ)
for name, typ in get_type_hints(schema, include_extras=True).items()
if name != "__slots__"
}
@@ -703,9 +719,9 @@ def _get_channels(
def _get_channel(
annotation: Any, *, allow_managed: bool = True
) -> Union[BaseChannel, Type[ManagedValue]]:
if manager := _is_field_managed_value(annotation):
name: str, annotation: Any, *, allow_managed: bool = True
) -> Union[BaseChannel, ManagedValueSpec]:
if manager := _is_field_managed_value(name, annotation):
if allow_managed:
return manager
else:
@@ -744,12 +760,18 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
return None
def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]:
def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueSpec]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
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
+35 -51
View File
@@ -1,13 +1,13 @@
import asyncio
from abc import ABC, abstractmethod
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from contextlib import asynccontextmanager, contextmanager
from inspect import isclass
from typing import (
Any,
AsyncGenerator,
Generator,
AsyncIterator,
Generic,
Iterator,
NamedTuple,
Sequence,
Type,
TypeVar,
Union,
@@ -17,6 +17,7 @@ from langchain_core.runnables import RunnableConfig
from typing_extensions import Self, TypeGuard
V = TypeVar("V")
U = TypeVar("U")
class ManagedValue(ABC, Generic[V]):
@@ -25,9 +26,7 @@ class ManagedValue(ABC, Generic[V]):
@classmethod
@contextmanager
def enter(
cls, config: RunnableConfig, **kwargs: Any
) -> Generator[Self, None, None]:
def enter(cls, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
try:
value = cls(config, **kwargs)
yield value
@@ -41,9 +40,7 @@ class ManagedValue(ABC, Generic[V]):
@classmethod
@asynccontextmanager
async def aenter(
cls, config: RunnableConfig, **kwargs: Any
) -> AsyncGenerator[Self, None]:
async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
try:
value = cls(config, **kwargs)
yield value
@@ -60,6 +57,16 @@ class ManagedValue(ABC, Generic[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]
@@ -76,46 +83,23 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
)
@contextmanager
def ManagedValuesManager(
values: dict[str, ManagedValueSpec],
config: RunnableConfig,
) -> Generator[ManagedValueMapping, None, None]:
if values:
with ExitStack() as stack:
yield {
key: stack.enter_context(
value.cls.enter(config, **value.kwargs)
if isinstance(value, ConfiguredManagedValue)
else value.enter(config)
)
for key, value in values.items()
}
else:
yield {}
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)
)
@asynccontextmanager
async def AsyncManagedValuesManager(
values: dict[str, ManagedValueSpec],
config: RunnableConfig,
) -> AsyncGenerator[ManagedValueMapping, None]:
if values:
async with AsyncExitStack() as stack:
# create enter tasks with reference to spec
tasks = {
asyncio.create_task(
stack.enter_async_context(
value.cls.aenter(config, **value.kwargs)
if isinstance(value, ConfiguredManagedValue)
else value.aenter(config)
)
): key
for key, value in values.items()
}
# wait for all enter tasks
done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
# build mapping from spec to result
yield {tasks[task]: task.result() for task in done}
else:
yield {}
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()
libs/langgraph/langgraph/managed/shared_value.py
+126
View File
@@ -0,0 +1,126 @@
import collections.abc
from contextlib import asynccontextmanager, contextmanager
from typing import (
Any,
AsyncIterator,
Iterator,
Optional,
Sequence,
Type,
)
from langchain_core.runnables import RunnableConfig
from typing_extensions import NotRequired, Required, Self
from langgraph.constants import CONFIG_KEY_STORE
from langgraph.errors import InvalidUpdateError
from langgraph.managed.base import (
ChannelKeyPlaceholder,
ChannelTypePlaceholder,
ConfiguredManagedValue,
WritableManagedValue,
)
from langgraph.store.base import BaseStore
V = dict[str, Any]
Value = dict[str, V]
Update = dict[str, Optional[V]]
# Adapted from typing_extensions
def _strip_extras(t):
"""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, config: RunnableConfig, **kwargs: Any) -> Iterator[Self]:
with super().enter(config, **kwargs) as value:
if value.store is not None:
saved = value.store.list([value.ns])
value.value = saved[value.ns] or {}
yield value
@classmethod
@asynccontextmanager
async def aenter(cls, config: RunnableConfig, **kwargs: Any) -> AsyncIterator[Self]:
async with super().aenter(config, **kwargs) as value:
if value.store is not None:
saved = await value.store.alist([value.ns])
value.value = saved[value.ns] or {}
yield value
def __init__(
self, config: RunnableConfig, *, typ: Type[Any], scope: str, key: str
) -> None:
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.config = config
self.value: Value = {}
self.store: BaseStore = config["configurable"].get(CONFIG_KEY_STORE)
if self.store is None:
self.ns: Optional[str] = None
elif scope_value := config["configurable"].get(self.scope):
self.ns = f"scoped:{scope}:{key}:{scope_value}"
else:
raise ValueError(
f"Scope {scope} for shared state key not in config.configurable"
)
def __call__(self, step: int) -> Value:
return self.value.copy()
def _process_update(
self, values: Sequence[Update]
) -> list[tuple[str, str, Optional[dict[str, Any]]]]:
writes = []
for vv in values:
for k, v in vv.items():
if v is None:
if k in self.value:
self.value[k] = None
writes.append((self.ns, k, None))
elif not isinstance(v, dict):
raise InvalidUpdateError("Received a non-dict value")
else:
self.value[k] = v
writes.append((self.ns, k, v))
return writes
def update(self, values: Sequence[Update]) -> None:
if self.store is None:
self._process_update(values)
else:
return self.store.put(self._process_update(values))
async def aupdate(self, writes: Sequence[Update]) -> None:
if self.store is None:
self._process_update(writes)
else:
return await self.store.aput(self._process_update(writes))
libs/langgraph/langgraph/pregel/__init__.py
+55 -68
View File
@@ -52,14 +52,8 @@ from langgraph.channels.base import (
BaseChannel,
)
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.manager import (
AsyncChannelsManager,
ChannelsManager,
)
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
CheckpointTuple,
copy_checkpoint,
create_checkpoint,
empty_checkpoint,
@@ -76,12 +70,7 @@ from langgraph.constants import (
Interrupt,
)
from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValuesManager,
ManagedValueSpec,
is_managed_value,
)
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.algo import (
apply_writes,
local_read,
@@ -100,6 +89,7 @@ from langgraph.pregel.io import (
read_channels,
)
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry
from langgraph.pregel.types import (
@@ -111,6 +101,7 @@ from langgraph.pregel.types import (
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
WriteValue = Union[
Runnable[Input, Output],
@@ -273,7 +264,9 @@ class Pregel(
):
nodes: Mapping[str, PregelNode]
channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field(
default_factory=dict
)
auto_validate: bool = True
@@ -300,6 +293,9 @@ class Pregel(
checkpointer: Optional[BaseCheckpointSaver] = None
"""Checkpointer used to save and load graph state. Defaults to None."""
store: Optional[BaseStore] = None
"""Memory store to use for SharedValues. Defaults to None."""
retry_policy: Optional[RetryPolicy] = None
"""Retry policy to use when running tasks. Set to None to disable."""
@@ -420,19 +416,12 @@ class Pregel(
@property
def stream_channels_asis(self) -> Union[str, Sequence[str]]:
return self.stream_channels or [
k for k in self.channels if not isinstance(self.channels[k], Context)
k
for k in self.channels
if isinstance(self.channels[k], BaseChannel)
and not isinstance(self.channels[k], Context)
]
@property
def managed_values_dict(self) -> dict[str, ManagedValueSpec]:
return {
k: v
for node in self.nodes.values()
if isinstance(node.channels, dict)
for k, v in node.channels.items()
if is_managed_value(v)
}
@property
def subgraphs(self) -> Iterator[Pregel]:
for node in self.nodes.values():
@@ -478,15 +467,11 @@ class Pregel(
graph = checkpoint_ns_to_graph[saved_checkpoint_ns]
with ChannelsManager(
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in graph.channels.items()
},
saved.checkpoint,
saved.config,
) as channels, ManagedValuesManager(
graph.managed_values_dict, ensure_config(saved.config)
) as managed:
graph.channels, saved.checkpoint, saved.config, skip_context=True
) as (
channels,
managed,
):
next_tasks = prepare_next_tasks(
saved.checkpoint,
graph.nodes,
@@ -547,6 +532,7 @@ class Pregel(
existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get(
saved_checkpoint_ns
)
# keep only most recent checkpoint_id
if (
existing_checkpoint_id is None
@@ -559,15 +545,8 @@ class Pregel(
graph = checkpoint_ns_to_graph[saved_checkpoint_ns]
async with AsyncChannelsManager(
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in graph.channels.items()
},
saved.checkpoint,
saved.config,
) as channels, AsyncManagedValuesManager(
graph.managed_values_dict, ensure_config(saved.config)
) as managed:
graph.channels, saved.checkpoint, saved.config, skip_context=True
) as (channels, managed):
next_tasks = prepare_next_tasks(
saved.checkpoint,
graph.nodes,
@@ -742,11 +721,10 @@ class Pregel(
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels
with ChannelsManager(
self.channels, checkpoint, config
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
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()
if not writers:
@@ -777,9 +755,9 @@ class Pregel(
),
)
# apply to checkpoint and save
apply_writes(
assert not apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
)
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
@@ -886,11 +864,10 @@ class Pregel(
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels, acting as the chosen node
async with AsyncChannelsManager(
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
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()
if not writers:
@@ -921,9 +898,9 @@ class Pregel(
),
)
# apply to checkpoint and save
apply_writes(
assert not apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
)
), "Can't write to SharedValues from update_state"
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
@@ -1127,7 +1104,13 @@ class Pregel(
)
with SyncPregelLoop(
input, config=config, checkpointer=checkpointer, graph=self
input,
config=config,
store=self.store,
checkpointer=checkpointer,
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
) as loop:
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
@@ -1135,7 +1118,8 @@ class Pregel(
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
output_keys=output_keys,
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1255,8 +1239,8 @@ class Pregel(
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
run_manager.on_chain_end(read_channels(loop.channels, output_keys))
# set final channel values as run output
run_manager.on_chain_end(loop.output)
except BaseException as e:
run_manager.on_chain_error(e)
raise
@@ -1379,7 +1363,13 @@ class Pregel(
debug=debug,
)
async with AsyncPregelLoop(
input, config=config, checkpointer=checkpointer, graph=self
input,
config=config,
store=self.store,
checkpointer=checkpointer,
nodes=self.nodes,
specs=self.channels,
output_keys=output_keys,
) as loop:
aioloop = asyncio.get_event_loop()
# Similarly to Bulk Synchronous Parallel / Pregel model
@@ -1388,7 +1378,8 @@ class Pregel(
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
output_keys=output_keys,
input_keys=self.input_channels,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
manager=run_manager,
@@ -1511,13 +1502,9 @@ class Pregel(
"without hitting a stop condition. You can increase the "
"limit by setting the `recursion_limit` config key."
)
# set final channel values as run output
await run_manager.on_chain_end(
read_channels(loop.channels, output_keys)
)
# set final channel values as run output
await run_manager.on_chain_end(loop.output)
except BaseException as e:
# TODO use on_chain_end if exc is GraphInterrupt
await asyncio.shield(run_manager.on_chain_error(e))
raise
libs/langgraph/langgraph/pregel/algo.py
+29 -20
View File
@@ -25,7 +25,6 @@ from langchain_core.runnables.config import (
from langgraph.channels.base import BaseChannel
from langgraph.channels.context import Context
from langgraph.channels.manager import ChannelsManager
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
@@ -46,9 +45,10 @@ from langgraph.constants import (
Send,
)
from langgraph.errors import EmptyChannelError, InvalidUpdateError
from langgraph.managed.base import ManagedValueMapping, is_managed_value
from langgraph.managed.base import ManagedValueMapping
from langgraph.pregel.io import read_channel, read_channels
from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
@@ -105,11 +105,10 @@ def local_read(
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(
{k: v for k, v in channels.items() if k not in context_channels},
new_checkpoint,
config,
) as channels:
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)
@@ -121,6 +120,7 @@ def local_write(
commit: Callable[[Sequence[tuple[str, Any]]], None],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
managed: ManagedValueMapping,
writes: Sequence[tuple[str, Any]],
) -> None:
for chan, value in writes:
@@ -131,7 +131,7 @@ def local_write(
)
if value.node not in processes:
raise InvalidUpdateError(f"Invalid node name {value.node} in packet")
elif chan not in channels:
elif chan not in channels and chan not in managed:
logger.warning(f"Skipping write for channel '{chan}' which has no readers")
commit(writes)
@@ -145,7 +145,7 @@ def apply_writes(
channels: Mapping[str, BaseChannel],
tasks: Sequence[WritesProtocol],
get_next_version: Optional[Callable[[int, BaseChannel], int]],
) -> None:
) -> dict[str, list[Any]]:
# update seen versions
for task in tasks:
checkpoint["versions_seen"].setdefault(task.name, {}).update(
@@ -161,6 +161,7 @@ def apply_writes(
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Consume all channels that were read
for chan in {
chan for task in tasks for chan in task.triggers if chan not in RESERVED
@@ -177,12 +178,15 @@ 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 == TASKS:
checkpoint["pending_sends"].append(val)
else:
elif chan in channels:
pending_writes_by_channel[chan].append(val)
else:
pending_writes_by_managed[chan].append(val)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
@@ -214,6 +218,9 @@ def apply_writes(
max_version, channels[chan]
)
# Return managed values writes to be applied externally
return pending_writes_by_managed
@overload
def prepare_next_tasks(
@@ -329,7 +336,11 @@ def prepare_next_tasks(
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
local_write,
writes.extend,
processes,
channels,
managed,
),
CONFIG_KEY_READ: partial(
local_read,
@@ -422,7 +433,11 @@ def prepare_next_tasks(
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
local_write,
writes.extend,
processes,
channels,
managed,
),
CONFIG_KEY_READ: partial(
local_read,
@@ -471,16 +486,10 @@ def _proc_input(
chan,
catch=chan not in proc.triggers,
)
if chan in channels
else managed[k](step)
for k, chan in proc.channels.items()
if isinstance(chan, str)
}
managed_values = {}
for key, chan in proc.channels.items():
if is_managed_value(chan):
managed_values[key] = managed[key](step)
val.update(managed_values)
except EmptyChannelError:
return
elif isinstance(proc.channels, list):
libs/langgraph/langgraph/pregel/loop.py
+96 -54
View File
@@ -4,7 +4,6 @@ from collections import deque
from contextlib import AsyncExitStack, ExitStack
from types import TracebackType
from typing import (
TYPE_CHECKING,
Any,
AsyncContextManager,
Callable,
@@ -18,6 +17,7 @@ from typing import (
Type,
TypeVar,
Union,
cast,
)
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
@@ -25,10 +25,6 @@ from langchain_core.runnables import RunnableConfig
from typing_extensions import Self
from langgraph.channels.base import BaseChannel
from langgraph.channels.manager import (
AsyncChannelsManager,
ChannelsManager,
)
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
@@ -49,9 +45,9 @@ from langgraph.constants import (
)
from langgraph.errors import EmptyInputError, GraphInterrupt
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValueMapping,
ManagedValuesManager,
ManagedValueSpec,
WritableManagedValue,
)
from langgraph.pregel.algo import (
PregelTaskWrites,
@@ -66,13 +62,19 @@ from langgraph.pregel.executor import (
BackgroundExecutor,
Submit,
)
from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
from langgraph.pregel.io import (
map_input,
map_output_updates,
map_output_values,
read_channels,
single,
)
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import PregelExecutableTask
from langgraph.pregel.utils import get_new_channel_versions
if TYPE_CHECKING:
from langgraph.pregel import Pregel
from langgraph.store.base import BaseStore
from langgraph.store.batch import AsyncBatchedStore
V = TypeVar("V")
INPUT_DONE = object()
@@ -83,7 +85,13 @@ EMPTY_SEQ = ()
class PregelLoop:
input: Optional[Any]
config: RunnableConfig
store: Optional[BaseStore]
checkpointer: Optional[BaseCheckpointSaver]
nodes: Mapping[str, PregelNode]
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]]
output_keys: Union[str, Sequence[str]]
is_nested: bool
checkpointer_get_next_version: Callable[[Optional[V]], V]
checkpointer_put_writes: Optional[
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
@@ -99,8 +107,6 @@ class PregelLoop:
Any,
]
]
graph: "Pregel"
submit: Submit
channels: Mapping[str, BaseChannel]
managed: ManagedValueMapping
@@ -118,7 +124,7 @@ class PregelLoop:
]
tasks: Sequence[PregelExecutableTask]
stream: deque[Tuple[str, Any]]
is_nested: bool
output: Union[None, dict[str, Any], Any] = None
# public
@@ -127,16 +133,20 @@ class PregelLoop:
input: Optional[Any],
*,
config: RunnableConfig,
store: Optional[BaseStore],
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]],
) -> None:
self.stream = deque()
self.input = input
self.config = config
self.store = store
self.checkpointer = checkpointer
self.graph = graph
# TODO if managed values no longer needs graph we can replace with
# managed_specs, channel_specs
self.nodes = nodes
self.specs = specs
self.output_keys = output_keys
self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
@@ -166,7 +176,8 @@ class PregelLoop:
def tick(
self,
*,
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
input_keys: Union[str, Sequence[str]],
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
@@ -178,20 +189,23 @@ class PregelLoop:
raise RuntimeError("Cannot tick when status is no longer 'pending'")
if self.input not in (INPUT_DONE, INPUT_RESUMING):
self._first()
self._first(input_keys=input_keys)
elif all(task.writes for task in self.tasks):
writes = [w for t in self.tasks for w in t.writes]
# all tasks have finished
apply_writes(
mv_writes = apply_writes(
self.checkpoint,
self.channels,
self.tasks,
self.checkpointer_get_next_version,
)
# apply writes to managed values
for key, values in mv_writes.items():
self._update_mv(key, values)
# produce values output
self.stream.extend(
("values", v)
for v in map_output_values(output_keys, writes, self.channels)
for v in map_output_values(self.output_keys, writes, self.channels)
)
# clear pending writes
self.checkpoint_pending_writes.clear()
@@ -199,11 +213,7 @@ class PregelLoop:
self._put_checkpoint(
{
"source": "loop",
"writes": single(
map_output_updates(output_keys, self.tasks)
if self.graph.stream_mode == "updates"
else map_output_values(output_keys, writes, self.channels)
),
"writes": single(map_output_updates(self.output_keys, self.tasks)),
}
)
# after execution, check if we should interrupt
@@ -227,7 +237,7 @@ class PregelLoop:
# prepare next tasks
self.tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.nodes,
self.channels,
self.managed,
self.config,
@@ -246,7 +256,7 @@ class PregelLoop:
self.step - 1, # printing checkpoint for previous step
self.checkpoint_config,
self.channels,
self.graph.stream_channels_asis,
stream_keys,
self.checkpoint_metadata,
self.checkpoint,
self.tasks,
@@ -270,7 +280,8 @@ class PregelLoop:
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks):
return self.tick(
output_keys=output_keys,
input_keys=input_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
@@ -294,7 +305,7 @@ class PregelLoop:
# private
def _first(self) -> None:
def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
# resuming from previous checkpoint requires
# - finding a previous checkpoint
# - receiving None input (outer graph) or RESUMING flag (subgraph)
@@ -311,11 +322,11 @@ class PregelLoop:
version = self.checkpoint["channel_versions"][k]
self.checkpoint["versions_seen"][INTERRUPT][k] = version
# map inputs to channel updates
elif input_writes := deque(map_input(self.graph.input_channels, self.input)):
elif input_writes := deque(map_input(input_keys, self.input)):
# discard any unfinished tasks from previous checkpoint
discard_tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.nodes,
self.channels,
self.managed,
self.config,
@@ -324,16 +335,16 @@ class PregelLoop:
manager=None,
)
# apply input writes
apply_writes(
assert not apply_writes(
self.checkpoint,
self.channels,
discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
self.checkpointer_get_next_version,
)
), "Can't write to SharedValues in graph input"
# save input checkpoint
self._put_checkpoint({"source": "input", "writes": self.input})
else:
raise EmptyInputError(f"Received no input for {self.graph.input_channels}")
raise EmptyInputError(f"Received no input for {input_keys}")
# done with input
self.input = INPUT_RESUMING if is_resuming else INPUT_DONE
@@ -395,13 +406,21 @@ class PregelLoop:
# 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]],
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
if isinstance(exc_value, GraphInterrupt) and not self.is_nested:
suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
if suppress or exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys)
if suppress:
# suppress interrupt
return True
@@ -411,10 +430,21 @@ class SyncPregelLoop(PregelLoop, ContextManager):
input: Optional[Any],
*,
config: RunnableConfig,
store: Optional[BaseStore],
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
super().__init__(
input,
config=config,
checkpointer=checkpointer,
store=store,
nodes=nodes,
specs=specs,
output_keys=output_keys,
)
self.stack = ExitStack()
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -438,6 +468,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
finally:
self.checkpointer.put(config, checkpoint, metadata, new_versions)
def _update_mv(self, key: str, values: Sequence[Any]) -> None:
return self.submit(cast(WritableManagedValue, self.managed[key]).update, values)
# context manager
def __enter__(self) -> Self:
@@ -457,11 +490,8 @@ class SyncPregelLoop(PregelLoop, ContextManager):
self.checkpoint_pending_writes = saved.pending_writes or []
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels = self.stack.enter_context(
ChannelsManager(self.graph.channels, self.checkpoint, self.config)
)
self.managed = self.stack.enter_context(
ManagedValuesManager(self.graph.managed_values_dict, self.config)
self.channels, self.managed = self.stack.enter_context(
ChannelsManager(self.specs, self.checkpoint, self.config, self.store)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
@@ -478,7 +508,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
traceback: Optional[TracebackType],
) -> Optional[bool]:
# unwind stack
del self.graph
return self.stack.__exit__(exc_type, exc_value, traceback)
@@ -488,10 +517,22 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
input: Optional[Any],
*,
config: RunnableConfig,
store: Optional[BaseStore],
checkpointer: Optional[BaseCheckpointSaver],
graph: "Pregel",
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
super().__init__(
input,
config=config,
checkpointer=checkpointer,
store=store,
nodes=nodes,
specs=specs,
output_keys=output_keys,
)
self.store = AsyncBatchedStore(self.store) if self.store else None
self.stack = AsyncExitStack()
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
@@ -515,6 +556,11 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
finally:
await self.checkpointer.aput(config, checkpoint, metadata, new_versions)
def _update_mv(self, key: str, values: Sequence[Any]) -> None:
return self.submit(
cast(WritableManagedValue, self.managed[key]).aupdate, values
)
# context manager
async def __aenter__(self) -> Self:
@@ -536,11 +582,8 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
self.checkpoint_pending_writes = saved.pending_writes or []
self.submit = await self.stack.enter_async_context(AsyncBackgroundExecutor())
self.channels = await self.stack.enter_async_context(
AsyncChannelsManager(self.graph.channels, self.checkpoint, self.config)
)
self.managed = await self.stack.enter_async_context(
AsyncManagedValuesManager(self.graph.managed_values_dict, self.config)
self.channels, self.managed = await self.stack.enter_async_context(
AsyncChannelsManager(self.specs, self.checkpoint, self.config, self.store)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
@@ -558,7 +601,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
traceback: Optional[TracebackType],
) -> Optional[bool]:
# unwind stack
del self.graph
return await asyncio.shield(
self.stack.__aexit__(exc_type, exc_value, traceback)
)
libs/langgraph/langgraph/pregel/manager.py
+104
View File
@@ -0,0 +1,104 @@
import asyncio
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
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 (
ConfiguredManagedValue,
ManagedValueMapping,
ManagedValueSpec,
)
from langgraph.store.base import BaseStore
@contextmanager
def ChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
config: RunnableConfig,
store: Optional[BaseStore] = None,
*,
skip_context: bool = False,
) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
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):
channel_specs[k] = v
else:
managed_specs[k] = v
with ExitStack() as stack:
yield (
{
k: stack.enter_context(
v.from_checkpoint(checkpoint["channel_values"].get(k), config)
)
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()
},
)
@asynccontextmanager
async def AsyncChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
config: RunnableConfig,
store: Optional[BaseStore] = None,
*,
skip_context: bool = False,
) -> AsyncIterator[Mapping[str, BaseChannel]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
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):
channel_specs[k] = v
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(config_for_managed, **value.kwargs)
if isinstance(value, ConfiguredManagedValue)
else value.aenter(config_for_managed)
)
): key
for key, value in managed_specs.items()
}:
done, _ = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
else:
done = set()
yield (
# channels: enter each channel with checkpoint
{
k: await stack.enter_async_context(
v.afrom_checkpoint(checkpoint["channel_values"].get(k), config)
)
for k, v in channel_specs.items()
},
# managed: build mapping from spec to result
{tasks[task]: task.result() for task in done},
)
libs/langgraph/langgraph/pregel/read.py
+1 -2
View File
@@ -15,7 +15,6 @@ from langchain_core.runnables.config import merge_configs
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_READ
from langgraph.managed.base import ManagedValueSpec
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.write import ChannelWrite
from langgraph.utils import RunnableCallable
@@ -101,7 +100,7 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
class PregelNode(RunnableBindingBase):
channels: Union[list[str], Mapping[str, Union[str, ManagedValueSpec]]]
channels: Union[list[str], Mapping[str, str]]
triggers: list[str] = Field(default_factory=list)
libs/langgraph/langgraph/pregel/retry.py
+7
View File
@@ -4,6 +4,7 @@ import random
import time
from typing import Optional
from langgraph.errors import GraphInterrupt
from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
logger = logging.getLogger(__name__)
@@ -25,6 +26,9 @@ def run_with_retry(
task.proc.invoke(task.input, task.config)
# if successful, end
break
except GraphInterrupt:
# if interrupted, end
raise
except Exception as exc:
if retry_policy is None:
raise
@@ -75,6 +79,9 @@ async def arun_with_retry(
await task.proc.ainvoke(task.input, task.config)
# if successful, end
break
except GraphInterrupt:
# if interrupted, end
raise
except Exception as exc:
if retry_policy is None:
raise
libs/langgraph/langgraph/store/__init__.py
View File
libs/langgraph/langgraph/store/base.py
+21
View File
@@ -0,0 +1,21 @@
from typing import Any, List, Optional
V = dict[str, Any]
class BaseStore:
def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
# list[namespace] -> dict[namespace, list[value]]
raise NotImplementedError
def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
# list[(namespace, key, value | none)] -> None
raise NotImplementedError
async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
# list[namespace] -> dict[namespace, list[value]]
raise NotImplementedError
async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
# list[(namespace, key, value | none)] -> None
raise NotImplementedError
libs/langgraph/langgraph/store/batch.py
+65
View File
@@ -0,0 +1,65 @@
import asyncio
from typing import NamedTuple, Optional, Union
from langgraph.store.base import BaseStore, V
class ListOp(NamedTuple):
prefixes: list[str]
class PutOp(NamedTuple):
writes: list[tuple[str, str, Optional[V]]]
class AsyncBatchedStore(BaseStore):
def __init__(self, store: BaseStore) -> None:
self.store = store
self.aqueue: dict[asyncio.Future, Union[ListOp, PutOp]] = {}
self.task = asyncio.create_task(_run(self.aqueue, self.store))
def __del__(self) -> None:
self.task.cancel()
async def alist(self, prefixes: list[str]) -> dict[str, dict[str, V]]:
fut = asyncio.get_running_loop().create_future()
self.aqueue[fut] = ListOp(prefixes)
return await fut
async def aput(self, writes: list[tuple[str, str, Optional[V]]]) -> None:
fut = asyncio.get_running_loop().create_future()
self.aqueue[fut] = PutOp(writes)
return await fut
async def _run(
aqueue: dict[asyncio.Future, Union[ListOp, PutOp]], store: BaseStore
) -> None:
while True:
await asyncio.sleep(0)
if not aqueue:
continue
# this could use a lock, if we want thread safety
taken = aqueue.copy()
aqueue.clear()
# action each operation
lists = {f: o for f, o in taken.items() if isinstance(o, ListOp)}
if lists:
try:
results = await store.alist(
[p for op in lists.values() for p in op.prefixes]
)
for fut, op in lists.items():
fut.set_result({k: results.get(k) for k in op.prefixes})
except Exception as e:
for fut in lists:
fut.set_exception(e)
puts = {f: o for f, o in taken.items() if isinstance(o, PutOp)}
if puts:
try:
await store.aput([w for op in puts.values() for w in op.writes])
for fut in puts:
fut.set_result(None)
except Exception as e:
for fut in puts:
fut.set_exception(e)
libs/langgraph/langgraph/store/memory.py
+25
View File
@@ -0,0 +1,25 @@
from collections import defaultdict
from typing import List, Optional
from langgraph.store.base import BaseStore, V
class MemoryStore(BaseStore):
def __init__(self) -> None:
self.data: dict[str, dict[str, V]] = defaultdict(dict)
def list(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
return {prefix: self.data[prefix] for prefix in prefixes}
async def alist(self, prefixes: List[str]) -> dict[str, dict[str, V]]:
return self.list(prefixes)
def put(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
for namespace, key, value in writes:
if value is None:
self.data[namespace].pop(key, None)
else:
self.data[namespace][key] = value
async def aput(self, writes: List[tuple[str, str, Optional[V]]]) -> None:
return self.put(writes)
libs/langgraph/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.6"
version = "0.2.9"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -65,7 +65,7 @@ omit = ["tests/*"]
[tool.pytest-watcher]
now = true
delay = 0.1
runner_args = ["--ff", "-v", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
runner_args = ["--ff", "-v", "-x", "-n", "auto", "--dist", "worksteal", "--snapshot-update", "--tb", "short"]
patterns = ["*.py"]
[build-system]
libs/langgraph/tests/fake_tracer.py
+91
View File
@@ -0,0 +1,91 @@
from typing import Any, Optional
from uuid import UUID
from langchain_core.messages.base import BaseMessage
from langchain_core.outputs.chat_generation import ChatGeneration
from langchain_core.outputs.llm_result import LLMResult
from langchain_core.tracers import BaseTracer, Run
class FakeTracer(BaseTracer):
"""Fake tracer that records LangChain execution.
It replaces run ids with deterministic UUIDs for snapshotting."""
def __init__(self) -> None:
"""Initialize the tracer."""
super().__init__()
self.runs: list[Run] = []
self.uuids_map: dict[UUID, UUID] = {}
self.uuids_generator = (
UUID(f"00000000-0000-4000-8000-{i:012}", version=4) for i in range(10000)
)
def _replace_uuid(self, uuid: UUID) -> UUID:
if uuid not in self.uuids_map:
self.uuids_map[uuid] = next(self.uuids_generator)
return self.uuids_map[uuid]
def _replace_message_id(self, maybe_message: Any) -> Any:
if isinstance(maybe_message, BaseMessage):
maybe_message.id = str(next(self.uuids_generator))
if isinstance(maybe_message, ChatGeneration):
maybe_message.message.id = str(next(self.uuids_generator))
if isinstance(maybe_message, LLMResult):
for i, gen_list in enumerate(maybe_message.generations):
for j, gen in enumerate(gen_list):
maybe_message.generations[i][j] = self._replace_message_id(gen)
if isinstance(maybe_message, dict):
for k, v in maybe_message.items():
maybe_message[k] = self._replace_message_id(v)
if isinstance(maybe_message, list):
for i, v in enumerate(maybe_message):
maybe_message[i] = self._replace_message_id(v)
return maybe_message
def _copy_run(self, run: Run) -> Run:
if run.dotted_order:
levels = run.dotted_order.split(".")
processed_levels = []
for level in levels:
timestamp, run_id = level.split("Z")
new_run_id = self._replace_uuid(UUID(run_id))
processed_level = f"{timestamp}Z{new_run_id}"
processed_levels.append(processed_level)
new_dotted_order = ".".join(processed_levels)
else:
new_dotted_order = None
return run.copy(
update={
"id": self._replace_uuid(run.id),
"parent_run_id": (
self.uuids_map[run.parent_run_id] if run.parent_run_id else None
),
"child_runs": [self._copy_run(child) for child in run.child_runs],
"trace_id": self._replace_uuid(run.trace_id) if run.trace_id else None,
"dotted_order": new_dotted_order,
"inputs": self._replace_message_id(run.inputs),
"outputs": self._replace_message_id(run.outputs),
}
)
def _persist_run(self, run: Run) -> None:
"""Persist a run."""
self.runs.append(self._copy_run(run))
def flattened_runs(self) -> list[Run]:
q = [] + self.runs
result = []
while q:
parent = q.pop()
result.append(parent)
if parent.child_runs:
q.extend(parent.child_runs)
return result
@property
def run_ids(self) -> list[Optional[UUID]]:
runs = self.flattened_runs()
uuids_map = {v: k for k, v in self.uuids_map.items()}
return [uuids_map.get(r.id) for r in runs]
libs/langgraph/tests/test_algo.py
+2 -5
View File
@@ -1,7 +1,6 @@
from langgraph.channels.manager import ChannelsManager
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.managed.base import ManagedValuesManager
from langgraph.pregel.algo import prepare_next_tasks
from langgraph.pregel.manager import ChannelsManager
def test_prepare_next_tasks() -> None:
@@ -9,9 +8,7 @@ def test_prepare_next_tasks() -> None:
processes = {}
checkpoint = empty_checkpoint()
with ManagedValuesManager({}, config) as managed, ChannelsManager(
{}, checkpoint, config
) as channels:
with ChannelsManager({}, checkpoint, config) as (channels, managed):
assert (
prepare_next_tasks(
checkpoint, processes, channels, managed, config, 0, for_execution=False
libs/langgraph/tests/test_pregel.py
+78 -31
View File
@@ -58,6 +58,7 @@ from langgraph.graph import END, Graph
from langgraph.graph.graph import START
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_tool_calling_executor,
)
@@ -70,7 +71,9 @@ from langgraph.pregel import (
)
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.types import PregelTask
from langgraph.store.memory import MemoryStore
from tests.any_str import AnyStr, ExceptionLike
from tests.fake_tracer import FakeTracer
from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
MemorySaverAssertImmutable,
@@ -660,7 +663,7 @@ def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 6, "writes": 5},
metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
created_at=AnyStr(),
parent_config=history[1].config,
),
@@ -675,7 +678,7 @@ def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 5, "writes": None},
metadata={"source": "loop", "step": 5, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[2].config,
),
@@ -705,7 +708,7 @@ def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 3, "writes": None},
metadata={"source": "loop", "step": 3, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[4].config,
),
@@ -735,7 +738,7 @@ def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 1, "writes": 4},
metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
created_at=AnyStr(),
parent_config=history[6].config,
),
@@ -750,7 +753,7 @@ def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 0, "writes": None},
metadata={"source": "loop", "step": 0, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[7].config,
),
@@ -1998,12 +2001,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
},
},
},
@@ -2209,12 +2214,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
}
},
},
@@ -2414,12 +2421,14 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
}
},
},
@@ -6142,20 +6151,34 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
my_key: Annotated[str, operator.add]
market: str
tool_two_node_count = 0
def tool_two_node(s: State) -> State:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
raise NodeInterrupt("Just because...")
return {"my_key": " all good"}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", tool_two_node)
tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
tracer = FakeTracer()
assert tool_two.invoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
@@ -6207,10 +6230,32 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
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)}
def tool_two_fast(data: State, config: RunnableConfig) -> State:
return {"my_key": " fast", **assert_shared_value(data, config)}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two_slow", lambda s: {"my_key": " slow"})
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
tool_two_graph.add_node("tool_two_slow", tool_two_slow)
tool_two_graph.add_node("tool_two_fast", tool_two_fast)
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
)
@@ -6228,14 +6273,16 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
store=MemoryStore(),
checkpointer=saver,
interrupt_before=["tool_two_fast", "tool_two_slow"],
)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
"my_key": "value ⛰️",
@@ -6287,7 +6334,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
)
thread2 = {"configurable": {"thread_id": "2"}}
thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "US"}, thread2) == {
"my_key": "value",
@@ -6327,7 +6374,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config,
)
thread3 = {"configurable": {"thread_id": "3"}}
thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value", "market": "US"}, thread3) == {
"my_key": "value",
libs/langgraph/tests/test_pregel_async.py
+148 -29
View File
@@ -52,6 +52,7 @@ from langgraph.errors import InvalidUpdateError, NodeInterrupt
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.graph import START
from langgraph.graph.message import MessageGraph, add_messages
from langgraph.managed.shared_value import SharedValue
from langgraph.prebuilt.chat_agent_executor import (
create_tool_calling_executor,
)
@@ -65,7 +66,9 @@ from langgraph.pregel import (
)
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.types import PregelTask
from langgraph.store.memory import MemoryStore
from tests.any_str import AnyStr, ExceptionLike
from tests.fake_tracer import FakeTracer
from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
MemorySaverAssertImmutable,
@@ -209,6 +212,91 @@ async def test_node_cancellation_on_other_node_exception() -> None:
assert inner_task_cancelled
async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
tool_two_node_count = 0
async def tool_two_node(s: State) -> State:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
raise NodeInterrupt("Just because...")
return {"my_key": " all good"}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert await tool_two.ainvoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(checkpointer=saver)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
await tool_two.ainvoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert await tool_two.ainvoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1
) == {
"my_key": "value ⛰️",
"market": "DE",
}
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value ⛰️", "market": "DE"},
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
interrupts=(Interrupt("during", "Just because..."),),
),
),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={"source": "loop", "step": 0, "writes": None},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1, limit=2)
][-1].config,
)
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
@@ -803,7 +891,7 @@ async def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 6, "writes": 5},
metadata={"source": "loop", "step": 6, "writes": {"two": 5}},
created_at=AnyStr(),
parent_config=history[1].config,
),
@@ -818,7 +906,7 @@ async def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 5, "writes": None},
metadata={"source": "loop", "step": 5, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[2].config,
),
@@ -848,7 +936,7 @@ async def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 3, "writes": None},
metadata={"source": "loop", "step": 3, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[4].config,
),
@@ -878,7 +966,7 @@ async def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 1, "writes": 4},
metadata={"source": "loop", "step": 1, "writes": {"two": 4}},
created_at=AnyStr(),
parent_config=history[6].config,
),
@@ -893,7 +981,7 @@ async def test_invoke_two_processes_in_out_interrupt(
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "step": 0, "writes": None},
metadata={"source": "loop", "step": 0, "writes": {"one": None}},
created_at=AnyStr(),
parent_config=history[7].config,
),
@@ -2217,12 +2305,14 @@ async def test_conditional_graph() -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
}
},
},
@@ -2443,12 +2533,14 @@ async def test_conditional_graph() -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
}
},
},
@@ -2675,12 +2767,14 @@ async def test_conditional_graph() -> None:
"step": 0,
"writes": {
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"agent": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
}
}
},
},
@@ -4783,10 +4877,33 @@ async def test_start_branch_then() -> 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)}
def tool_two_fast(data: State, config: RunnableConfig) -> State:
return {"my_key": " fast", **assert_shared_value(data, config)}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two_slow", lambda s, config: {"my_key": " slow"})
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
tool_two_graph.add_node("tool_two_slow", tool_two_slow)
tool_two_graph.add_node("tool_two_fast", tool_two_fast)
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
)
@@ -4803,14 +4920,16 @@ async def test_start_branch_then() -> None:
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(
checkpointer=saver, interrupt_before=["tool_two_fast", "tool_two_slow"]
store=MemoryStore(),
checkpointer=saver,
interrupt_before=["tool_two_fast", "tool_two_slow"],
)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
await tool_two.ainvoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}}
# stop when about to enter node
assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == {
"my_key": "value",
@@ -4870,7 +4989,7 @@ async def test_start_branch_then() -> None:
][-1].config,
)
thread2 = {"configurable": {"thread_id": "2"}}
thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}}
# stop when about to enter node
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == {
"my_key": "value",
@@ -4918,7 +5037,7 @@ async def test_start_branch_then() -> None:
][-1].config,
)
thread3 = {"configurable": {"thread_id": "3"}}
thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}}
# stop when about to enter node
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == {
"my_key": "value",
libs/langgraph/tests/test_store.py
+35
View File
@@ -0,0 +1,35 @@
import asyncio
from typing import Any, Optional
from pytest_mock import MockerFixture
from langgraph.store.base import BaseStore
from langgraph.store.batch import AsyncBatchedStore
async def test_async_batch_store(mocker: MockerFixture) -> None:
aget = mocker.stub()
alist = mocker.stub()
class MockStore(BaseStore):
async def aget(
self, pairs: list[tuple[str, str]]
) -> dict[tuple[str, str], Optional[dict[str, Any]]]:
aget(pairs)
return {pair: 1 for pair in pairs}
async def alist(self, prefixes: list[str]) -> dict[str, dict[str, Any]]:
alist(prefixes)
return {prefix: {prefix: 1} for prefix in prefixes}
store = AsyncBatchedStore(MockStore())
# concurrent calls are batched
results = await asyncio.gather(
store.alist(["a", "b"]),
store.alist(["c", "d"]),
)
assert results == [{"a": {"a": 1}, "b": {"b": 1}}, {"c": {"c": 1}, "d": {"d": 1}}]
assert [c.args for c in alist.call_args_list] == [
(["a", "b", "c", "d"],),
]
libs/sdk-py/langgraph_sdk/client.py
+30 -8
View File
@@ -40,8 +40,14 @@ from langgraph_sdk.schema import (
logger = logging.getLogger(__name__)
RESERVED_HEADERS = ("x-api-key",)
def get_client(
*, url: Optional[str] = None, api_key: Optional[str] = None
*,
url: Optional[str] = None,
api_key: Optional[str] = None,
headers: Optional[dict[str, str]] = None,
) -> LangGraphClient:
"""Get a LangGraphClient instance.
@@ -53,6 +59,7 @@ def get_client(
2. LANGGRAPH_API_KEY
3. LANGSMITH_API_KEY
4. LANGCHAIN_API_KEY
headers: Optional custom headers
"""
transport: Optional[httpx.AsyncBaseTransport] = None
if url is None:
@@ -65,17 +72,12 @@ def get_client(
url = "http://localhost:8123"
if transport is None:
transport = httpx.AsyncHTTPTransport(retries=5)
headers = {
"User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}",
}
api_key = _get_api_key(api_key)
if api_key:
headers["x-api-key"] = api_key
client = httpx.AsyncClient(
base_url=url,
transport=transport,
timeout=httpx.Timeout(connect=5, read=60, write=60, pool=5),
headers=headers,
headers=_get_headers(api_key, headers),
)
return LangGraphClient(client)
@@ -1695,3 +1697,23 @@ def _get_api_key(api_key: Optional[str] = None) -> Optional[str]:
if env := os.getenv(f"{prefix}_API_KEY"):
return env.strip().strip('"').strip("'")
return None # type: ignore
def _get_headers(
api_key: Optional[str], custom_headers: Optional[dict[str, str]]
) -> dict[str, str]:
"""Combine api_key and custom user-provided headers."""
custom_headers = custom_headers or {}
for header in RESERVED_HEADERS:
if header in custom_headers:
raise ValueError(f"Cannot set reserved header '{header}'")
headers = {
"User-Agent": f"langgraph-sdk-py/{langgraph_sdk.__version__}",
**custom_headers,
}
api_key = _get_api_key(api_key)
if api_key:
headers["x-api-key"] = api_key
return headers
libs/sdk-py/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-sdk"
version = "0.1.27"
version = "0.1.28"
description = "SDK for interacting with LangGraph API"
authors = []
license = "MIT"