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

Remove ChannelsManager, ManagedValues are now static classes and never instantiated (#4812)

Browse Source
This commit is contained in:
Nuno Campos
2025-05-28 14:02:43 -07:00
committed by GitHub
parent 2e4645ce61 9c79d887a0
commit dc95f95181
13 changed files with 901 additions and 992 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/managed/base.py
+5 -37
View File
@@ -1,6 +1,4 @@
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from inspect import isclass
from typing import (
Any,
@@ -8,48 +6,18 @@ from typing import (
TypeVar,
)
from typing_extensions import Self, TypeGuard
from typing_extensions import TypeGuard
from langgraph.types import LoopProtocol
from langgraph.types import PregelScratchpad
V = TypeVar("V")
U = TypeVar("U")
class ManagedValue(ABC, Generic[V]):
def __init__(self, loop: LoopProtocol) -> None:
self.loop = loop
@classmethod
@contextmanager
def enter(cls, loop: LoopProtocol, **kwargs: Any) -> Iterator[Self]:
try:
value = cls(loop, **kwargs)
yield value
finally:
# because managed value and Pregel have reference to each other
# let's make sure to break the reference on exit
try:
del value
except UnboundLocalError:
pass
@classmethod
@asynccontextmanager
async def aenter(cls, loop: LoopProtocol, **kwargs: Any) -> AsyncIterator[Self]:
try:
value = cls(loop, **kwargs)
yield value
finally:
# because managed value and Pregel have reference to each other
# let's make sure to break the reference on exit
try:
del value
except UnboundLocalError:
pass
@staticmethod
@abstractmethod
def __call__(self) -> V: ...
def get(scratchpad: PregelScratchpad) -> V: ...
ManagedValueSpec = type[ManagedValue]
@@ -59,4 +27,4 @@ def is_managed_value(value: Any) -> TypeGuard[ManagedValueSpec]:
return isclass(value) and issubclass(value, ManagedValue)
ManagedValueMapping = dict[str, ManagedValue]
ManagedValueMapping = dict[str, ManagedValueSpec]
libs/langgraph/langgraph/managed/is_last_step.py
+7 -4
View File
@@ -1,19 +1,22 @@
from typing import Annotated
from langgraph.managed.base import ManagedValue
from langgraph.types import PregelScratchpad
class IsLastStepManager(ManagedValue[bool]):
def __call__(self) -> bool:
return self.loop.step == self.loop.stop - 1
@staticmethod
def get(scratchpad: PregelScratchpad) -> bool:
return scratchpad.step == scratchpad.stop - 1
IsLastStep = Annotated[bool, IsLastStepManager]
class RemainingStepsManager(ManagedValue[int]):
def __call__(self) -> int:
return self.loop.stop - self.loop.step
@staticmethod
def get(scratchpad: PregelScratchpad) -> int:
return scratchpad.stop - scratchpad.step
RemainingSteps = Annotated[int, RemainingStepsManager]
libs/langgraph/langgraph/pregel/__init__.py
+558 -584
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/langgraph/pregel/algo.py
+47 -26
View File
@@ -364,6 +364,7 @@ def prepare_next_tasks(
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
stop: int,
*,
for_execution: Literal[False],
store: Literal[None] = None,
@@ -385,6 +386,7 @@ def prepare_next_tasks(
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
stop: int,
*,
for_execution: Literal[True],
store: Optional[BaseStore],
@@ -405,6 +407,7 @@ def prepare_next_tasks(
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
stop: int,
*,
for_execution: bool,
store: Optional[BaseStore] = None,
@@ -459,6 +462,7 @@ def prepare_next_tasks(
managed=managed,
config=config,
step=step,
stop=stop,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
@@ -504,6 +508,7 @@ def prepare_next_tasks(
managed=managed,
config=config,
step=step,
stop=stop,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
@@ -532,6 +537,7 @@ def prepare_single_task(
managed: ManagedValueMapping,
config: RunnableConfig,
step: int,
stop: int,
for_execution: bool,
store: Optional[BaseStore] = None,
checkpointer: Optional[BaseCheckpointSaver] = None,
@@ -632,6 +638,8 @@ def prepare_single_task(
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
),
},
),
@@ -752,6 +760,8 @@ def prepare_single_task(
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
),
CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
PREVIOUS, None
@@ -785,23 +795,6 @@ def prepare_single_task(
proc,
):
triggers = tuple(sorted(proc.triggers))
try:
val = _proc_input(
proc,
managed,
channels,
for_execution=for_execution,
input_cache=input_cache,
)
if val is MISSING:
return
except Exception as exc:
if SUPPORTS_EXC_NOTES:
exc.add_note(
f"Before task with name '{name}' and path '{task_path[:3]}'"
)
raise
# create task id
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = task_id_func(
@@ -813,6 +806,35 @@ def prepare_single_task(
*triggers,
)
task_checkpoint_ns = f"{checkpoint_ns}{NS_END}{task_id}"
# create scratchpad
scratchpad = _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
step,
stop,
)
# create task input
try:
val = _proc_input(
proc,
managed,
channels,
for_execution=for_execution,
input_cache=input_cache,
scratchpad=scratchpad,
)
if val is MISSING:
return
except Exception as exc:
if SUPPORTS_EXC_NOTES:
exc.add_note(
f"Before task with name '{name}' and path '{task_path[:3]}'"
)
raise
metadata = {
"langgraph_step": step,
"langgraph_node": name,
@@ -888,13 +910,7 @@ def prepare_single_task(
},
CONFIG_KEY_CHECKPOINT_ID: None,
CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns,
CONFIG_KEY_SCRATCHPAD: _scratchpad(
config[CONF].get(CONFIG_KEY_SCRATCHPAD),
pending_writes,
task_id,
xxh3_128_hexdigest(task_checkpoint_ns.encode()),
config[CONF].get(CONFIG_KEY_RESUME_MAP),
),
CONFIG_KEY_SCRATCHPAD: scratchpad,
CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get(
PREVIOUS, None
),
@@ -947,6 +963,8 @@ def _scratchpad(
task_id: str,
namespace_hash: str,
resume_map: Optional[dict[str, Any]],
step: int,
stop: int,
) -> PregelScratchpad:
if len(pending_writes) > 0:
# find global resume value
@@ -994,6 +1012,8 @@ def _scratchpad(
# using itertools.count as an atomic counter (+= 1 is not thread-safe)
return PregelScratchpad(
step=step,
stop=stop,
# call
call_counter=LazyAtomicCounter(),
# interrupt
@@ -1011,6 +1031,7 @@ def _proc_input(
channels: Mapping[str, BaseChannel],
*,
for_execution: bool,
scratchpad: PregelScratchpad,
input_cache: Optional[dict[INPUT_CACHE_KEY_TYPE, Any]],
) -> Any:
"""Prepare input for a PULL task, based on the process's channels and triggers."""
@@ -1026,7 +1047,7 @@ def _proc_input(
if channels[chan].is_available():
val[k] = channels[chan].get()
else:
val[k] = managed[k]()
val[k] = managed[k].get(scratchpad)
elif isinstance(proc.channels, list):
for chan in proc.channels:
if chan in channels:
@@ -1034,7 +1055,7 @@ def _proc_input(
val = channels[chan].get()
break
else:
val = managed[chan]()
val = managed[chan].get(scratchpad)
break
else:
return MISSING
libs/langgraph/langgraph/pregel/checkpoint.py
+23 -1
View File
@@ -1,11 +1,12 @@
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Optional
from typing import Optional, Union
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
from langgraph.checkpoint.base.id import uuid6
from langgraph.constants import MISSING
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
LATEST_VERSION = 3
@@ -50,3 +51,24 @@ def create_checkpoint(
versions_seen=checkpoint["versions_seen"],
pending_sends=checkpoint.get("pending_sends", []),
)
def channels_from_checkpoint(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
) -> tuple[Mapping[str, BaseChannel], ManagedValueMapping]:
"""Get channels from a checkpoint."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
if isinstance(v, BaseChannel):
channel_specs[k] = v
else:
managed_specs[k] = v
return (
{
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
for k, v in channel_specs.items()
},
managed_specs,
)
libs/langgraph/langgraph/pregel/draw.py
+165 -165
View File
@@ -15,12 +15,11 @@ from langgraph.pregel.algo import (
increment,
prepare_next_tasks,
)
from langgraph.pregel.checkpoint import empty_checkpoint
from langgraph.pregel.checkpoint import channels_from_checkpoint, empty_checkpoint
from langgraph.pregel.io import map_input
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite
from langgraph.types import All, Checkpointer, LoopProtocol
from langgraph.types import All, Checkpointer
def draw_graph(
@@ -56,31 +55,102 @@ def draw_graph(
if isinstance(checkpointer, BaseCheckpointSaver)
else increment
)
with ChannelsManager(
channels, managed = channels_from_checkpoint(
specs,
checkpoint,
LoopProtocol(step=step, stop=-1, config=config),
) as (channels, managed):
static_seen: set[Any] = set()
sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
# remove node mappers
nodes = {
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
for k, v in nodes.items()
)
static_seen: set[Any] = set()
sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
# remove node mappers
nodes = {
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
for k, v in nodes.items()
}
# apply input writes
input_writes = list(map_input(input_channels, {}))
updated_channels = apply_writes(
checkpoint,
channels,
[
PregelTaskWrites((), INPUT, input_writes, []),
],
get_next_version,
trigger_to_nodes,
)
# prepare first tasks
tasks = prepare_next_tasks(
checkpoint,
[],
nodes,
channels,
managed,
config,
step,
-1,
for_execution=True,
store=None,
checkpointer=None,
manager=None,
trigger_to_nodes=trigger_to_nodes,
updated_channels=updated_channels,
)
start_tasks = tasks
# run the pregel loop
for step in range(step, limit):
if not tasks:
break
conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
# run task writers
for task in tasks.values():
for w in task.writers:
# apply regular writes
if isinstance(w, ChannelWrite):
empty_input = (
cast(BaseChannel, specs["__root__"]).ValueType()
if "__root__" in specs
else None
)
w.invoke(empty_input, task.config)
# apply conditional writes declared for static analysis, only once
if w not in static_seen:
static_seen.add(w)
# apply static writes
if writes := ChannelWrite.get_static_writes(w):
# END writes are not written, but become edges directly
for t in writes:
if t[0] == END:
edges.add((task.name, t[0], True, t[2]))
writes = [t for t in writes if t[0] != END]
conditionals.update(
{(task.name, t[0], t[1] or None): t[2] for t in writes}
)
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
# collect sources
step_sources = {
task.name: {
(
w[0],
(task.name, w[0], w[1] or None) in conditionals,
conditionals.get((task.name, w[0], w[1] or None)),
)
for w in task.writes
}
for task in tasks.values()
}
# apply input writes
input_writes = list(map_input(input_channels, {}))
updated_channels = apply_writes(
checkpoint,
channels,
[
PregelTaskWrites((), INPUT, input_writes, []),
],
get_next_version,
trigger_to_nodes,
sources.update(step_sources)
# invert triggers
trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
defaultdict(set)
)
# prepare first tasks
for src, triggers in sources.items():
for trigger, cond, label in triggers:
trigger_to_sources[trigger].add((src, cond, label))
# apply writes
updated_channels = apply_writes(
checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes
)
# prepare next tasks
tasks = prepare_next_tasks(
checkpoint,
[],
@@ -89,6 +159,7 @@ def draw_graph(
managed,
config,
step,
limit,
for_execution=True,
store=None,
checkpointer=None,
@@ -96,149 +167,78 @@ def draw_graph(
trigger_to_nodes=trigger_to_nodes,
updated_channels=updated_channels,
)
start_tasks = tasks
# run the pregel loop
for step in range(step, limit):
if not tasks:
break
conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
# run task writers
for task in tasks.values():
for w in task.writers:
# apply regular writes
if isinstance(w, ChannelWrite):
empty_input = (
cast(BaseChannel, specs["__root__"]).ValueType()
if "__root__" in specs
else None
)
w.invoke(empty_input, task.config)
# apply conditional writes declared for static analysis, only once
if w not in static_seen:
static_seen.add(w)
# apply static writes
if writes := ChannelWrite.get_static_writes(w):
# END writes are not written, but become edges directly
for t in writes:
if t[0] == END:
edges.add((task.name, t[0], True, t[2]))
writes = [t for t in writes if t[0] != END]
conditionals.update(
{(task.name, t[0], t[1] or None): t[2] for t in writes}
)
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
# collect sources
step_sources = {
task.name: {
(
w[0],
(task.name, w[0], w[1] or None) in conditionals,
conditionals.get((task.name, w[0], w[1] or None)),
)
for w in task.writes
}
for task in tasks.values()
}
sources.update(step_sources)
# invert triggers
trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
defaultdict(set)
)
for src, triggers in sources.items():
for trigger, cond, label in triggers:
trigger_to_sources[trigger].add((src, cond, label))
# apply writes
updated_channels = apply_writes(
checkpoint, channels, tasks.values(), get_next_version, trigger_to_nodes
)
# prepare next tasks
tasks = prepare_next_tasks(
checkpoint,
[],
nodes,
channels,
managed,
config,
step,
for_execution=True,
store=None,
checkpointer=None,
manager=None,
trigger_to_nodes=trigger_to_nodes,
updated_channels=updated_channels,
)
# collect edges
for task in tasks.values():
added = False
for trigger in task.triggers:
for src, cond, label in sorted(trigger_to_sources[trigger]):
edges.add((src, task.name, cond, label))
# if the edge is from this step, skip adding the implicit edges
if (trigger, cond, label) in step_sources.get(src, set()):
added = True
else:
sources[src].discard((trigger, cond, label))
# if no edges from this step, add implicit edges from all previous tasks
if not added:
for src in step_sources:
edges.add((src, task.name, True, None))
# assemble the graph
graph = Graph()
# add nodes
for name, node in nodes.items():
metadata = dict(node.metadata or {})
if name in interrupt_before_nodes and name in interrupt_after_nodes:
metadata["__interrupt"] = "before,after"
elif name in interrupt_before_nodes:
metadata["__interrupt"] = "before"
elif name in interrupt_after_nodes:
metadata["__interrupt"] = "after"
graph.add_node(node.bound, name, metadata=metadata or None)
# add start node
if START not in nodes:
graph.add_node(None, START)
for task in start_tasks.values():
add_edge(graph, START, task.name)
# add discovered edges
for src, dest, is_conditional, label in sorted(edges):
add_edge(
graph,
src,
dest,
data=label if label != dest else None,
conditional=is_conditional,
)
# add end edges
termini = {d for _, d, _, _ in edges if d != END}.difference(
s for s, _, _, _ in edges
# collect edges
for task in tasks.values():
added = False
for trigger in task.triggers:
for src, cond, label in sorted(trigger_to_sources[trigger]):
edges.add((src, task.name, cond, label))
# if the edge is from this step, skip adding the implicit edges
if (trigger, cond, label) in step_sources.get(src, set()):
added = True
else:
sources[src].discard((trigger, cond, label))
# if no edges from this step, add implicit edges from all previous tasks
if not added:
for src in step_sources:
edges.add((src, task.name, True, None))
# assemble the graph
graph = Graph()
# add nodes
for name, node in nodes.items():
metadata = dict(node.metadata or {})
if name in interrupt_before_nodes and name in interrupt_after_nodes:
metadata["__interrupt"] = "before,after"
elif name in interrupt_before_nodes:
metadata["__interrupt"] = "before"
elif name in interrupt_after_nodes:
metadata["__interrupt"] = "after"
graph.add_node(node.bound, name, metadata=metadata or None)
# add start node
if START not in nodes:
graph.add_node(None, START)
for task in start_tasks.values():
add_edge(graph, START, task.name)
# add discovered edges
for src, dest, is_conditional, label in sorted(edges):
add_edge(
graph,
src,
dest,
data=label if label != dest else None,
conditional=is_conditional,
)
if termini:
for src in sorted(termini):
add_edge(graph, src, END)
elif len(step_sources) == 1:
for src in sorted(step_sources):
add_edge(graph, src, END, conditional=True)
# replace subgraphs
for name, subgraph in subgraphs.items():
if (
len(subgraph.nodes) > 1
and name in graph.nodes
and subgraph.first_node()
and subgraph.last_node()
):
subgraph.trim_first_node()
subgraph.trim_last_node()
# replace the node with the subgraph
graph.nodes.pop(name)
first, last = graph.extend(subgraph, prefix=name)
for idx, edge in enumerate(graph.edges):
if edge.source == name:
edge = edge.copy(source=cast(Node, last).id)
if edge.target == name:
edge = edge.copy(target=cast(Node, first).id)
graph.edges[idx] = edge
# add end edges
termini = {d for _, d, _, _ in edges if d != END}.difference(
s for s, _, _, _ in edges
)
if termini:
for src in sorted(termini):
add_edge(graph, src, END)
elif len(step_sources) == 1:
for src in sorted(step_sources):
add_edge(graph, src, END, conditional=True)
# replace subgraphs
for name, subgraph in subgraphs.items():
if (
len(subgraph.nodes) > 1
and name in graph.nodes
and subgraph.first_node()
and subgraph.last_node()
):
subgraph.trim_first_node()
subgraph.trim_last_node()
# replace the node with the subgraph
graph.nodes.pop(name)
first, last = graph.extend(subgraph, prefix=name)
for idx, edge in enumerate(graph.edges):
if edge.source == name:
edge = edge.copy(source=cast(Node, last).id)
if edge.target == name:
edge = edge.copy(target=cast(Node, first).id)
graph.edges[idx] = edge
return graph
return graph
def add_edge(
libs/langgraph/langgraph/pregel/loop.py
+24 -15
View File
@@ -89,7 +89,11 @@ from langgraph.pregel.algo import (
should_interrupt,
task_path_str,
)
from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint
from langgraph.pregel.checkpoint import (
channels_from_checkpoint,
create_checkpoint,
empty_checkpoint,
)
from langgraph.pregel.debug import (
map_debug_checkpoint,
map_debug_task_results,
@@ -111,7 +115,6 @@ from langgraph.pregel.io import (
read_channels,
single,
)
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.utils import get_new_channel_versions, is_xxh3_128_hexdigest
from langgraph.store.base import BaseStore
@@ -119,7 +122,6 @@ from langgraph.types import (
All,
CachePolicy,
Command,
LoopProtocol,
PregelExecutableTask,
PregelScratchpad,
RetryPolicy,
@@ -146,7 +148,13 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
class PregelLoop(LoopProtocol):
class PregelLoop:
config: RunnableConfig
store: Optional["BaseStore"]
stream: Optional[StreamProtocol]
step: int
stop: int
input: Optional[Any]
input_model: Optional[type[BaseModel]]
cache: Optional[BaseCache[WritesT]]
@@ -226,13 +234,11 @@ class PregelLoop(LoopProtocol):
cache_policy: Optional[CachePolicy] = None,
checkpoint_during: bool = True,
) -> None:
super().__init__(
step=0,
stop=0,
config=config,
stream=stream,
store=store,
)
self.stream = stream
self.config = config
self.store = store
self.step = 0
self.stop = 0
self.input = input
self.input_model = input_model
self.checkpointer = checkpointer
@@ -423,6 +429,7 @@ class PregelLoop(LoopProtocol):
managed=self.managed,
config=task.config,
step=self.step,
stop=self.stop,
for_execution=True,
store=self.store,
checkpointer=self.checkpointer,
@@ -554,6 +561,7 @@ class PregelLoop(LoopProtocol):
self.managed,
self.config,
self.step,
self.stop,
for_execution=True,
manager=self.manager,
store=self.store,
@@ -738,6 +746,7 @@ class PregelLoop(LoopProtocol):
self.managed,
self.config,
self.step,
self.stop,
for_execution=True,
store=None,
checkpointer=None,
@@ -1146,8 +1155,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
)
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
self.channels, self.managed = self.stack.enter_context(
ChannelsManager(self.specs, self.checkpoint, self)
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
@@ -1341,8 +1350,8 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
self.submit = await self.stack.enter_async_context(
AsyncBackgroundExecutor(self.config)
)
self.channels, self.managed = await self.stack.enter_async_context(
AsyncChannelsManager(self.specs, self.checkpoint, self)
self.channels, self.managed = channels_from_checkpoint(
self.specs, self.checkpoint
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
libs/langgraph/langgraph/pregel/manager.py
-76
View File
@@ -1,76 +0,0 @@
import asyncio
from collections.abc import AsyncIterator, Iterator, Mapping
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from typing import Union
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
from langgraph.constants import MISSING
from langgraph.managed.base import (
ManagedValueMapping,
ManagedValueSpec,
)
from langgraph.types import LoopProtocol
@contextmanager
def ChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
loop: LoopProtocol,
) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
if isinstance(v, BaseChannel):
channel_specs[k] = v
else:
managed_specs[k] = v
with ExitStack() as stack:
yield (
{
k: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
for k, v in channel_specs.items()
},
ManagedValueMapping(
{
key: stack.enter_context(value.enter(loop))
for key, value in managed_specs.items()
}
),
)
@asynccontextmanager
async def AsyncChannelsManager(
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
checkpoint: Checkpoint,
loop: LoopProtocol,
) -> AsyncIterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
channel_specs: dict[str, BaseChannel] = {}
managed_specs: dict[str, ManagedValueSpec] = {}
for k, v in specs.items():
if 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.aenter(loop))): 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: v.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
for k, v in channel_specs.items()
},
# managed: build mapping from spec to result
ManagedValueMapping({tasks[task]: task.result() for task in done}),
)
libs/langgraph/langgraph/types.py
+2 -24
View File
@@ -27,7 +27,6 @@ from langgraph.utils.fields import get_update_as_tuples
if TYPE_CHECKING:
from langgraph.pregel.protocol import PregelProtocol
from langgraph.store.base import BaseStore
try:
@@ -380,31 +379,10 @@ class StreamProtocol:
self.modes = modes
class LoopProtocol:
config: RunnableConfig
store: Optional["BaseStore"]
stream: Optional[StreamProtocol]
step: int
stop: int
def __init__(
self,
*,
step: int,
stop: int,
config: RunnableConfig,
store: Optional["BaseStore"] = None,
stream: Optional[StreamProtocol] = None,
) -> None:
self.stream = stream
self.config = config
self.store = store
self.step = step
self.stop = stop
@dataclasses.dataclass(**_DC_KWARGS)
class PregelScratchpad:
step: int
stop: int
# call
call_counter: Callable[[], int]
# interrupt
libs/langgraph/tests/test_algo.py
+32 -30
View File
@@ -1,46 +1,48 @@
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.constants import PULL, PUSH
from langgraph.pregel.algo import prepare_next_tasks, task_path_str
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.checkpoint import channels_from_checkpoint
def test_prepare_next_tasks() -> None:
config = {}
processes = {}
checkpoint = empty_checkpoint()
channels, managed = channels_from_checkpoint({}, checkpoint)
with ChannelsManager({}, checkpoint, config) as (channels, managed):
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
config,
0,
for_execution=False,
)
== {}
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
config,
0,
-1,
for_execution=False,
)
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
config,
0,
for_execution=True,
checkpointer=None,
store=None,
manager=None,
)
== {}
== {}
)
assert (
prepare_next_tasks(
checkpoint,
{},
processes,
channels,
managed,
config,
0,
-1,
for_execution=True,
checkpointer=None,
store=None,
manager=None,
)
== {}
)
# TODO: add more tests
# TODO: add more tests
def test_tuple_str() -> None:
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+12 -28
View File
@@ -22,12 +22,12 @@ from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR
from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound
from langgraph.pregel import Pregel
from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task
from langgraph.pregel.checkpoint import channels_from_checkpoint
from langgraph.pregel.executor import (
AsyncBackgroundExecutor,
BackgroundExecutor,
Submit,
)
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.runner import PregelRunner
from langgraph.scheduler.kafka.retry import aretry, retry
from langgraph.scheduler.kafka.types import (
@@ -41,7 +41,7 @@ from langgraph.scheduler.kafka.types import (
Sendable,
Topics,
)
from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy
from langgraph.types import PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable, recast_checkpoint_ns
@@ -184,19 +184,10 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
raise RuntimeError("Checkpoint not found")
if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
raise CheckpointNotLatest()
async with (
AsyncChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
AsyncBackgroundExecutor(msg["config"]) as submit,
):
async with AsyncBackgroundExecutor(msg["config"]) as submit:
channels, managed = channels_from_checkpoint(
graph.channels, saved.checkpoint
)
if task := await asyncio.to_thread(
prepare_single_task,
msg["task"]["path"],
@@ -208,6 +199,7 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
managed=managed,
config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}),
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
for_execution=True,
checkpointer=self.graph.checkpointer,
store=self.graph.store,
@@ -404,19 +396,10 @@ class KafkaExecutor(AbstractContextManager):
raise RuntimeError("Checkpoint not found")
if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]:
raise CheckpointNotLatest()
with (
ChannelsManager(
graph.channels,
saved.checkpoint,
LoopProtocol(
config=msg["config"],
store=self.graph.store,
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
),
) as (channels, managed),
BackgroundExecutor({}) as submit,
):
with BackgroundExecutor({}) as submit:
channels, managed = channels_from_checkpoint(
graph.channels, saved.checkpoint
)
if task := prepare_single_task(
msg["task"]["path"],
msg["task"]["id"],
@@ -427,6 +410,7 @@ class KafkaExecutor(AbstractContextManager):
managed=managed,
config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}),
step=saved.metadata["step"] + 1,
stop=saved.metadata["step"] + 2,
for_execution=True,
checkpointer=self.graph.checkpointer,
checkpoint_id_bytes=binascii.unhexlify(
libs/scheduler-kafka/tests/test_subgraph.py
+13 -1
View File
@@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph
from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyInt
from tests.drain import drain_topics_async
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -199,6 +199,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -272,6 +274,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -375,6 +379,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -488,6 +494,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -556,6 +564,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -680,6 +690,8 @@ async def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
libs/scheduler-kafka/tests/test_subgraph_sync.py
+13 -1
View File
@@ -15,7 +15,7 @@ from langgraph.pregel import Pregel
from langgraph.scheduler.kafka import serde
from langgraph.scheduler.kafka.default_sync import DefaultProducer
from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics
from tests.any import AnyDict
from tests.any import AnyDict, AnyInt
from tests.drain import drain_topics
from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage
@@ -198,6 +198,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -271,6 +273,8 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -374,6 +378,8 @@ def test_subgraph_w_interrupt(
"__pregel_task_id": history[0].tasks[0].id,
"__pregel_previous": None,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -486,6 +492,8 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -554,6 +562,8 @@ def test_subgraph_w_interrupt(
"__pregel_previous": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,
@@ -678,6 +688,8 @@ def test_subgraph_w_interrupt(
"__pregel_store": None,
"__pregel_task_id": history[1].tasks[0].id,
"__pregel_scratchpad": {
"step": AnyInt(),
"stop": AnyInt(),
"subgraph_counter": None,
"call_counter": None,
"interrupt_counter": None,