mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 19:29:43 +02:00
Implement AsyncLoop
This commit is contained in:
Nuno Campos
@@ -61,4 +61,6 @@ def create_checkpoint(
|
||||
channel_versions=checkpoint["channel_versions"],
|
||||
versions_seen=checkpoint["versions_seen"],
|
||||
pending_sends=checkpoint.get("pending_sends", []),
|
||||
# checkpoints are saved only at the end of a step, ie. no current tasks
|
||||
current_tasks={},
|
||||
)
|
||||
|
||||
@@ -54,6 +54,10 @@ class CheckpointMetadata(TypedDict, total=False):
|
||||
"""
|
||||
|
||||
|
||||
class TaskInfo(TypedDict):
|
||||
status: Literal["scheduled", "success", "error"]
|
||||
|
||||
|
||||
class Checkpoint(TypedDict):
|
||||
"""State snapshot at a given point in time."""
|
||||
|
||||
@@ -85,6 +89,8 @@ class Checkpoint(TypedDict):
|
||||
pending_sends: List[Send]
|
||||
"""List of packets sent to nodes but not yet processed.
|
||||
Cleared by the next checkpoint."""
|
||||
current_tasks: Dict[str, TaskInfo]
|
||||
"""Map from task ID to task info."""
|
||||
|
||||
|
||||
def empty_checkpoint() -> Checkpoint:
|
||||
@@ -96,6 +102,7 @@ def empty_checkpoint() -> Checkpoint:
|
||||
channel_versions={},
|
||||
versions_seen=defaultdict(dict),
|
||||
pending_sends=[],
|
||||
current_tasks={},
|
||||
)
|
||||
|
||||
|
||||
@@ -111,6 +118,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
|
||||
{k: v.copy() for k, v in checkpoint["versions_seen"].items()},
|
||||
),
|
||||
pending_sends=checkpoint.get("pending_sends", []).copy(),
|
||||
current_tasks=checkpoint.get("current_tasks", {}).copy(),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -59,7 +59,6 @@ from langgraph.channels.manager import (
|
||||
)
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
CheckpointMetadata,
|
||||
copy_checkpoint,
|
||||
empty_checkpoint,
|
||||
)
|
||||
@@ -77,28 +76,20 @@ from langgraph.managed.base import (
|
||||
)
|
||||
from langgraph.pregel.algo import (
|
||||
apply_writes,
|
||||
increment,
|
||||
local_read,
|
||||
prepare_next_tasks,
|
||||
should_interrupt,
|
||||
)
|
||||
from langgraph.pregel.debug import (
|
||||
map_debug_checkpoint,
|
||||
map_debug_task_results,
|
||||
map_debug_tasks,
|
||||
print_step_checkpoint,
|
||||
print_step_tasks,
|
||||
print_step_writes,
|
||||
)
|
||||
from langgraph.pregel.executor import AsyncBackgroundExecutor
|
||||
from langgraph.pregel.io import (
|
||||
map_input,
|
||||
map_output_updates,
|
||||
map_output_values,
|
||||
read_channels,
|
||||
single,
|
||||
)
|
||||
from langgraph.pregel.loop import PregelLoop
|
||||
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry
|
||||
from langgraph.pregel.types import (
|
||||
@@ -839,8 +830,7 @@ class Pregel(
|
||||
debug=debug,
|
||||
)
|
||||
|
||||
# create channels from checkpoint
|
||||
with PregelLoop(
|
||||
with SyncPregelLoop(
|
||||
input, config=config, checkpointer=self.checkpointer, graph=self
|
||||
) as loop:
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
@@ -852,6 +842,7 @@ class Pregel(
|
||||
output_keys=output_keys,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
manager=run_manager,
|
||||
):
|
||||
# debug flag
|
||||
if self.debug:
|
||||
@@ -871,19 +862,16 @@ class Pregel(
|
||||
# debug flag
|
||||
if debug:
|
||||
print_step_tasks(loop.step, loop.tasks)
|
||||
# TODO move to tick() ?
|
||||
if "debug" in stream_modes:
|
||||
yield from _with_mode(
|
||||
"debug",
|
||||
isinstance(stream_mode, list),
|
||||
map_debug_tasks(loop.step, loop.tasks),
|
||||
)
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
# yield updates/debug output as each task finishes
|
||||
futures = {
|
||||
loop.submit(run_with_retry, task, self.retry_policy): task
|
||||
loop.submit(
|
||||
run_with_retry,
|
||||
task,
|
||||
self.retry_policy,
|
||||
): task
|
||||
for task in loop.tasks
|
||||
if not task.writes
|
||||
}
|
||||
@@ -1063,7 +1051,6 @@ class Pregel(
|
||||
None,
|
||||
)
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
if config["recursion_limit"] < 1:
|
||||
raise ValueError("recursion_limit must be at least 1")
|
||||
if self.checkpointer and not config.get("configurable"):
|
||||
@@ -1085,211 +1072,59 @@ class Pregel(
|
||||
interrupt_after=interrupt_after,
|
||||
debug=debug,
|
||||
)
|
||||
# copy nodes to ignore mutations during execution
|
||||
processes = {**self.nodes}
|
||||
# get checkpoint from saver, or create an empty one
|
||||
saved = (
|
||||
await self.checkpointer.aget_tuple(config)
|
||||
if self.checkpointer
|
||||
else None
|
||||
)
|
||||
checkpoint = saved.checkpoint if saved else empty_checkpoint()
|
||||
|
||||
# merge configurable fields with previous checkpoint config
|
||||
checkpoint_config = config
|
||||
if saved:
|
||||
checkpoint_config = {
|
||||
**config,
|
||||
**saved.config,
|
||||
"configurable": {
|
||||
**config.get("configurable", {}),
|
||||
**saved.config["configurable"],
|
||||
},
|
||||
}
|
||||
|
||||
start = saved.metadata.get("step", -2) + 1 if saved else -1
|
||||
# create channels from checkpoint
|
||||
async with AsyncBackgroundExecutor() as submit, AsyncChannelsManager(
|
||||
self.channels, checkpoint, config
|
||||
) as channels, AsyncManagedValuesManager(
|
||||
self.managed_values_dict, config, self
|
||||
) as managed:
|
||||
|
||||
def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
if self.checkpointer is not None:
|
||||
submit(
|
||||
self.checkpointer.aput_writes,
|
||||
{
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]:
|
||||
print(metadata)
|
||||
nonlocal checkpoint, checkpoint_config, channels
|
||||
|
||||
if self.checkpointer is None:
|
||||
return
|
||||
if debug:
|
||||
print_step_checkpoint(
|
||||
metadata["step"], channels, self.stream_channels_list
|
||||
)
|
||||
|
||||
# create new checkpoint
|
||||
checkpoint = create_checkpoint(
|
||||
checkpoint, channels, metadata["step"]
|
||||
)
|
||||
# save it, without blocking
|
||||
submit(
|
||||
self.checkpointer.aput,
|
||||
checkpoint_config,
|
||||
copy_checkpoint(checkpoint),
|
||||
metadata,
|
||||
)
|
||||
|
||||
# update checkpoint config
|
||||
checkpoint_config = {
|
||||
**checkpoint_config,
|
||||
"configurable": {
|
||||
**checkpoint_config["configurable"],
|
||||
"thread_ts": checkpoint["id"],
|
||||
},
|
||||
}
|
||||
# yield debug checkpoint event
|
||||
if "debug" in stream_modes:
|
||||
yield from _with_mode(
|
||||
"debug",
|
||||
isinstance(stream_mode, list),
|
||||
map_debug_checkpoint(
|
||||
metadata["step"],
|
||||
checkpoint_config,
|
||||
channels,
|
||||
self.stream_channels_asis,
|
||||
metadata,
|
||||
),
|
||||
)
|
||||
|
||||
# map inputs to channel updates
|
||||
if input_writes := deque(map_input(self.input_channels, input)):
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
checkpoint, _ = prepare_next_tasks(
|
||||
checkpoint,
|
||||
processes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
-1,
|
||||
for_execution=True,
|
||||
get_next_version=(
|
||||
self.checkpointer.get_next_version
|
||||
if self.checkpointer
|
||||
else increment
|
||||
),
|
||||
)
|
||||
# apply input writes
|
||||
apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
input_writes,
|
||||
(
|
||||
self.checkpointer.get_next_version
|
||||
if self.checkpointer
|
||||
else increment
|
||||
),
|
||||
)
|
||||
# save input checkpoint
|
||||
for chunk in put_checkpoint(
|
||||
{"source": "input", "step": start, "writes": input}
|
||||
):
|
||||
yield chunk
|
||||
# increment start to 0
|
||||
start += 1
|
||||
else:
|
||||
# no input is taken as signal to proceed past previous interrupt
|
||||
checkpoint = copy_checkpoint(checkpoint)
|
||||
for k in channels:
|
||||
if k in checkpoint["channel_versions"]:
|
||||
version = checkpoint["channel_versions"][k]
|
||||
checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
|
||||
async with AsyncPregelLoop(
|
||||
input, config=config, checkpointer=self.checkpointer, graph=self
|
||||
) as loop:
|
||||
aioloop = asyncio.get_event_loop()
|
||||
# Similarly to Bulk Synchronous Parallel / Pregel model
|
||||
# computation proceeds in steps, while there are channel updates
|
||||
# channel updates from step N are only visible in step N+1,
|
||||
# channel updates from step N are only visible in step N+1
|
||||
# channels are guaranteed to be immutable for the duration of the step,
|
||||
# channel updates being applied only at the transition between steps
|
||||
stop = start + config["recursion_limit"] + 1
|
||||
for step in range(start, stop):
|
||||
next_checkpoint, next_tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
processes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
for_execution=True,
|
||||
manager=run_manager,
|
||||
get_next_version=(
|
||||
self.checkpointer.get_next_version
|
||||
if self.checkpointer
|
||||
else increment
|
||||
),
|
||||
)
|
||||
|
||||
# assign pending writes to tasks
|
||||
if saved and saved.pending_writes:
|
||||
for task in next_tasks:
|
||||
task.writes.extend(
|
||||
(c, v)
|
||||
for tid, c, v in saved.pending_writes
|
||||
if tid == task.id
|
||||
)
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not next_tasks:
|
||||
if step == start:
|
||||
raise ValueError("No tasks to run in graph.")
|
||||
else:
|
||||
break
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if should_interrupt(checkpoint, interrupt_before, next_tasks):
|
||||
break
|
||||
else:
|
||||
checkpoint = next_checkpoint
|
||||
|
||||
# with channel updates applied only at the transition between steps
|
||||
while loop.tick(
|
||||
output_keys=output_keys,
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
manager=run_manager,
|
||||
):
|
||||
# debug flag
|
||||
if self.debug:
|
||||
print_step_checkpoint(
|
||||
loop.checkpoint_metadata,
|
||||
loop.channels,
|
||||
self.stream_channels_list,
|
||||
)
|
||||
# emit output
|
||||
while loop.stream:
|
||||
mode, payload = loop.stream.popleft()
|
||||
if mode in stream_modes:
|
||||
if isinstance(stream_mode, list):
|
||||
yield (mode, payload)
|
||||
else:
|
||||
yield payload
|
||||
# debug flag
|
||||
if debug:
|
||||
print_step_tasks(step, next_tasks)
|
||||
if "debug" in stream_modes:
|
||||
for chunk in _with_mode(
|
||||
"debug",
|
||||
isinstance(stream_mode, list),
|
||||
map_debug_tasks(step, next_tasks),
|
||||
):
|
||||
yield chunk
|
||||
print_step_tasks(loop.step, loop.tasks)
|
||||
|
||||
# execute tasks, and wait for one to fail or all to finish.
|
||||
# each task is independent from all other concurrent tasks
|
||||
# yield updates/debug output as each task finishes
|
||||
futures = {
|
||||
submit(
|
||||
loop.submit(
|
||||
arun_with_retry,
|
||||
task,
|
||||
self.retry_policy,
|
||||
do_stream,
|
||||
stream=do_stream,
|
||||
__name__=task.name,
|
||||
__cancel_on_exit__=True,
|
||||
): task
|
||||
for task in next_tasks
|
||||
for task in loop.tasks
|
||||
if not task.writes
|
||||
}
|
||||
end_time = (
|
||||
self.step_timeout + loop.time() if self.step_timeout else None
|
||||
self.step_timeout + aioloop.time()
|
||||
if self.step_timeout
|
||||
else None
|
||||
)
|
||||
if not futures:
|
||||
done, inflight = set(), set()
|
||||
@@ -1298,7 +1133,7 @@ class Pregel(
|
||||
futures,
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
timeout=(
|
||||
max(0, end_time - loop.time()) if end_time else None
|
||||
max(0, end_time - aioloop.time()) if end_time else None
|
||||
),
|
||||
)
|
||||
if not done:
|
||||
@@ -1307,13 +1142,12 @@ class Pregel(
|
||||
task = futures.pop(fut)
|
||||
if fut.exception() is not None:
|
||||
# we got an exception, break out of while loop
|
||||
# exception will be handle in panic_or_proceed
|
||||
# exception will be handled in panic_or_proceed
|
||||
futures.clear()
|
||||
else:
|
||||
# save task writes to checkpointer, unless this
|
||||
# is the single or last task in this step
|
||||
if futures:
|
||||
put_writes(task.id, task.writes)
|
||||
print(loop.step, task.name, stream_modes)
|
||||
# save task writes to checkpointer
|
||||
loop.put_writes(task.id, task.writes)
|
||||
# yield updates output for the finished task
|
||||
if "updates" in stream_modes:
|
||||
for chunk in _with_mode(
|
||||
@@ -1327,7 +1161,9 @@ class Pregel(
|
||||
"debug",
|
||||
isinstance(stream_mode, list),
|
||||
map_debug_task_results(
|
||||
step, [task], self.stream_channels_list
|
||||
loop.step,
|
||||
[task],
|
||||
self.stream_channels_list,
|
||||
),
|
||||
):
|
||||
yield chunk
|
||||
@@ -1336,71 +1172,36 @@ class Pregel(
|
||||
del fut, task
|
||||
|
||||
# panic on failure or timeout
|
||||
_panic_or_proceed(done, inflight, step, asyncio.TimeoutError)
|
||||
_panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError)
|
||||
# don't keep futures around in memory longer than needed
|
||||
del done, inflight, futures
|
||||
|
||||
# combine pending writes from all tasks
|
||||
pending_writes = deque[tuple[str, Any]]()
|
||||
for task in next_tasks:
|
||||
pending_writes.extend(task.writes)
|
||||
|
||||
# debug flag
|
||||
if debug:
|
||||
print_step_writes(
|
||||
step, pending_writes, self.stream_channels_list
|
||||
loop.step,
|
||||
[w for t in loop.tasks for w in t.writes],
|
||||
self.stream_channels_list,
|
||||
)
|
||||
|
||||
# apply writes to channels
|
||||
apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
pending_writes,
|
||||
(
|
||||
self.checkpointer.get_next_version
|
||||
if self.checkpointer
|
||||
else increment
|
||||
),
|
||||
)
|
||||
|
||||
# yield current values
|
||||
if "values" in stream_modes:
|
||||
for chunk in _with_mode(
|
||||
"values",
|
||||
isinstance(stream_mode, list),
|
||||
map_output_values(output_keys, pending_writes, channels),
|
||||
):
|
||||
yield chunk
|
||||
|
||||
# save end of step checkpoint
|
||||
for chunk in put_checkpoint(
|
||||
{
|
||||
"source": "loop",
|
||||
"step": step,
|
||||
"writes": (
|
||||
single(map_output_updates(output_keys, next_tasks))
|
||||
if self.stream_mode == "updates"
|
||||
else single(
|
||||
map_output_values(
|
||||
output_keys, pending_writes, channels
|
||||
)
|
||||
)
|
||||
),
|
||||
}
|
||||
):
|
||||
yield chunk
|
||||
|
||||
# after execution, check if we should interrupt
|
||||
if should_interrupt(checkpoint, interrupt_after, next_tasks):
|
||||
break
|
||||
else:
|
||||
# emit output
|
||||
while loop.stream:
|
||||
mode, payload = loop.stream.popleft()
|
||||
if mode in stream_modes:
|
||||
if isinstance(stream_mode, list):
|
||||
yield (mode, payload)
|
||||
else:
|
||||
yield payload
|
||||
# handle exit
|
||||
if loop.status == "out_of_steps":
|
||||
raise GraphRecursionError(
|
||||
f"Recursion limit of {config['recursion_limit']} reached"
|
||||
"without hitting a stop condition. You can increase the limit"
|
||||
"by setting the `recursion_limit` config key."
|
||||
f"Recursion limit of {config['recursion_limit']} reached "
|
||||
"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(channels, output_keys))
|
||||
await run_manager.on_chain_end(
|
||||
read_channels(loop.channels, output_keys)
|
||||
)
|
||||
except BaseException as e:
|
||||
await asyncio.shield(run_manager.on_chain_error(e))
|
||||
raise
|
||||
|
||||
@@ -6,6 +6,7 @@ from contextvars import copy_context
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
AsyncContextManager,
|
||||
Awaitable,
|
||||
Callable,
|
||||
Iterator,
|
||||
Optional,
|
||||
@@ -78,7 +79,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
|
||||
|
||||
def submit(
|
||||
self,
|
||||
fn: Callable[P, T],
|
||||
fn: Callable[P, Awaitable[T]],
|
||||
*args: P.args,
|
||||
__name__: Optional[str] = None,
|
||||
__cancel_on_exit__: bool = False,
|
||||
@@ -101,7 +102,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
|
||||
else:
|
||||
self.tasks.pop(task)
|
||||
|
||||
async def __aenter__(self) -> Submit:
|
||||
async def __aenter__(self) -> "submit":
|
||||
return self.submit
|
||||
|
||||
async def exit(self) -> None:
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from contextlib import ExitStack
|
||||
from contextlib import AsyncExitStack, ExitStack
|
||||
from types import TracebackType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncContextManager,
|
||||
Callable,
|
||||
ContextManager,
|
||||
List,
|
||||
@@ -17,11 +19,16 @@ from typing import (
|
||||
Union,
|
||||
)
|
||||
|
||||
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.manager import ChannelsManager, create_checkpoint
|
||||
from langgraph.channels.manager import (
|
||||
AsyncChannelsManager,
|
||||
ChannelsManager,
|
||||
create_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
@@ -32,15 +39,23 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.constants import INTERRUPT
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValuesManager
|
||||
from langgraph.managed.base import (
|
||||
AsyncManagedValuesManager,
|
||||
ManagedValueMapping,
|
||||
ManagedValuesManager,
|
||||
)
|
||||
from langgraph.pregel.algo import (
|
||||
apply_writes,
|
||||
increment,
|
||||
prepare_next_tasks,
|
||||
should_interrupt,
|
||||
)
|
||||
from langgraph.pregel.debug import map_debug_checkpoint
|
||||
from langgraph.pregel.executor import BackgroundExecutor, Submit
|
||||
from langgraph.pregel.debug import map_debug_checkpoint, map_debug_tasks
|
||||
from langgraph.pregel.executor import (
|
||||
AsyncBackgroundExecutor,
|
||||
BackgroundExecutor,
|
||||
Submit,
|
||||
)
|
||||
from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single
|
||||
from langgraph.pregel.types import PregelExecutableTask
|
||||
|
||||
@@ -52,11 +67,17 @@ V = TypeVar("V")
|
||||
INPUT_DONE = object()
|
||||
|
||||
|
||||
class PregelLoop(ContextManager):
|
||||
class PregelLoop:
|
||||
input: Optional[Any]
|
||||
config: RunnableConfig
|
||||
checkpointer: Optional[BaseCheckpointSaver]
|
||||
get_next_version: Callable[[Optional[V]], V]
|
||||
checkpointer_get_next_version: Callable[[Optional[V]], V]
|
||||
checkpointer_put_writes: Optional[
|
||||
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], ...]
|
||||
]
|
||||
checkpointer_put: Optional[
|
||||
Callable[[RunnableConfig, Checkpoint, CheckpointMetadata], ...]
|
||||
]
|
||||
graph: "Pregel"
|
||||
|
||||
submit: Submit
|
||||
@@ -73,6 +94,203 @@ class PregelLoop(ContextManager):
|
||||
tasks: Sequence[PregelExecutableTask]
|
||||
stream: deque[Tuple[str, Any]]
|
||||
|
||||
# public
|
||||
|
||||
def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None:
|
||||
"""Mark tasks as scheduled, to be used by queue-based executors."""
|
||||
raise NotImplementedError
|
||||
|
||||
def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
|
||||
if self.checkpointer_put_writes is not None:
|
||||
self.submit(
|
||||
self.checkpointer_put_writes,
|
||||
{
|
||||
**self.checkpoint_config,
|
||||
"configurable": {
|
||||
**self.checkpoint_config["configurable"],
|
||||
"thread_ts": self.checkpoint["id"],
|
||||
},
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def tick(
|
||||
self,
|
||||
*,
|
||||
output_keys: Union[str, Sequence[str]] = None,
|
||||
interrupt_after: Optional[Sequence[str]] = None,
|
||||
interrupt_before: Optional[Sequence[str]] = None,
|
||||
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
|
||||
) -> bool:
|
||||
"""Execute a single iteration of the Pregel loop.
|
||||
Returns True if more iterations are needed."""
|
||||
|
||||
if self.status != "pending":
|
||||
raise RuntimeError("Cannot tick when status is no longer 'pending'")
|
||||
|
||||
if self.input is not INPUT_DONE:
|
||||
self._first()
|
||||
elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len(
|
||||
self.tasks
|
||||
):
|
||||
# assign writes to tasks, apply them in order
|
||||
grouped: dict[str, list[tuple[str, Any]]] = {}
|
||||
for tid, k, v in self.checkpoint_pending_writes:
|
||||
grouped.setdefault(tid, []).append((k, v))
|
||||
writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])]
|
||||
# all tasks have finished
|
||||
apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
writes,
|
||||
self.checkpointer_get_next_version,
|
||||
)
|
||||
# produce values output
|
||||
self.stream.extend(
|
||||
("values", v)
|
||||
for v in map_output_values(output_keys, writes, self.channels)
|
||||
)
|
||||
# clear pending writes
|
||||
self.checkpoint_pending_writes.clear()
|
||||
# save checkpoint
|
||||
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)
|
||||
),
|
||||
}
|
||||
)
|
||||
# after execution, check if we should interrupt
|
||||
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
|
||||
self.status = "interrupt_after"
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
# check if iteration limit is reached
|
||||
if self.step > self.config["recursion_limit"]:
|
||||
self.status = "out_of_steps"
|
||||
return False
|
||||
|
||||
# prepare next tasks
|
||||
prev_checkpoint = self.checkpoint
|
||||
self.checkpoint, self.tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.graph.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
for_execution=True,
|
||||
get_next_version=self.checkpointer_get_next_version,
|
||||
manager=manager,
|
||||
)
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not self.tasks:
|
||||
self.status = "done"
|
||||
return False
|
||||
|
||||
# if there are pending writes from a previous loop, apply them
|
||||
if self.checkpoint_pending_writes:
|
||||
for tid, k, v in self.checkpoint_pending_writes:
|
||||
if task := next((t for t in self.tasks if t.id == tid), None):
|
||||
task.writes.append((k, v))
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if should_interrupt(prev_checkpoint, interrupt_before, self.tasks):
|
||||
self.status = "interrupt_before"
|
||||
return False
|
||||
|
||||
# produce debug output
|
||||
self.stream.extend(("debug", v) for v in map_debug_tasks(self.step, self.tasks))
|
||||
|
||||
return True
|
||||
|
||||
# private
|
||||
|
||||
def _first(self) -> None:
|
||||
# map inputs to channel updates
|
||||
if input_writes := deque(map_input(self.graph.input_channels, self.input)):
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
self.checkpoint, _ = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.graph.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
for_execution=True,
|
||||
get_next_version=self.checkpointer_get_next_version,
|
||||
)
|
||||
# apply input writes
|
||||
apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
input_writes,
|
||||
self.checkpointer_get_next_version,
|
||||
)
|
||||
# save input checkpoint
|
||||
self._put_checkpoint({"source": "input", "writes": self.input})
|
||||
else:
|
||||
# no input is taken as signal to proceed past previous interrupt
|
||||
self.checkpoint = copy_checkpoint(self.checkpoint)
|
||||
for k in self.channels:
|
||||
if k in self.checkpoint["channel_versions"]:
|
||||
version = self.checkpoint["channel_versions"][k]
|
||||
self.checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
# done with input
|
||||
self.input = INPUT_DONE
|
||||
|
||||
def _put_checkpoint(
|
||||
self,
|
||||
metadata: CheckpointMetadata,
|
||||
) -> None:
|
||||
# assign step
|
||||
metadata["step"] = self.step
|
||||
# bail if no checkpointer
|
||||
if self.checkpointer_put is not None:
|
||||
# create new checkpoint
|
||||
self.checkpoint_metadata = metadata
|
||||
self.checkpoint = create_checkpoint(
|
||||
self.checkpoint, self.channels, self.step
|
||||
)
|
||||
# save it, without blocking
|
||||
self.submit(
|
||||
self.checkpointer_put,
|
||||
self.checkpoint_config,
|
||||
copy_checkpoint(self.checkpoint),
|
||||
self.checkpoint_metadata,
|
||||
)
|
||||
self.checkpoint_config = {
|
||||
**self.checkpoint_config,
|
||||
"configurable": {
|
||||
**self.checkpoint_config["configurable"],
|
||||
"thread_ts": self.checkpoint["id"],
|
||||
},
|
||||
}
|
||||
# produce debug output
|
||||
self.stream.extend(
|
||||
("debug", v)
|
||||
for v in map_debug_checkpoint(
|
||||
self.step,
|
||||
self.checkpoint_config,
|
||||
self.channels,
|
||||
self.graph.stream_channels_asis,
|
||||
self.checkpoint_metadata,
|
||||
)
|
||||
)
|
||||
# increment step
|
||||
self.step += 1
|
||||
|
||||
|
||||
class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
def __init__(
|
||||
self,
|
||||
input: Optional[Any],
|
||||
@@ -86,13 +304,17 @@ class PregelLoop(ContextManager):
|
||||
self.input = input
|
||||
self.config = config
|
||||
self.checkpointer = checkpointer
|
||||
self.get_next_version = (
|
||||
self.checkpointer_get_next_version = (
|
||||
checkpointer.get_next_version if checkpointer else increment
|
||||
)
|
||||
self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None
|
||||
self.checkpointer_put = checkpointer.put if checkpointer else None
|
||||
self.graph = graph
|
||||
# TODO if managed values no longer needs graph we can replace with
|
||||
# managed_specs, channel_specs
|
||||
|
||||
# context manager
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
saved = (
|
||||
self.checkpointer.get_tuple(self.config) if self.checkpointer else None
|
||||
@@ -132,183 +354,73 @@ class PregelLoop(ContextManager):
|
||||
del self.graph
|
||||
return self.stack.__exit__(exc_type, exc_value, traceback)
|
||||
|
||||
def tick(
|
||||
|
||||
class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
def __init__(
|
||||
self,
|
||||
input: Optional[Any],
|
||||
*,
|
||||
output_keys: Union[str, Sequence[str]] = None,
|
||||
interrupt_after: Optional[Sequence[str]] = None,
|
||||
interrupt_before: Optional[Sequence[str]] = None,
|
||||
) -> bool:
|
||||
if self.status != "pending":
|
||||
raise RuntimeError(f"Cannot tick when status is {self.status}")
|
||||
if self.input is not INPUT_DONE:
|
||||
self.first()
|
||||
elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len(
|
||||
self.tasks
|
||||
):
|
||||
# assign writes to tasks, apply them in order
|
||||
grouped: dict[str, list[tuple[str, Any]]] = {}
|
||||
for tid, k, v in self.checkpoint_pending_writes:
|
||||
grouped.setdefault(tid, []).append((k, v))
|
||||
writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])]
|
||||
# all tasks have finished
|
||||
apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
writes,
|
||||
self.get_next_version,
|
||||
)
|
||||
# produce values output
|
||||
self.stream.extend(
|
||||
("values", v)
|
||||
for v in map_output_values(output_keys, writes, self.channels)
|
||||
)
|
||||
# clear pending writes
|
||||
self.checkpoint_pending_writes.clear()
|
||||
# save checkpoint
|
||||
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)
|
||||
),
|
||||
}
|
||||
)
|
||||
# after execution, check if we should interrupt
|
||||
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
|
||||
self.status = "interrupt_after"
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
# check if iteration limit is reached
|
||||
if self.step > self.config["recursion_limit"]:
|
||||
self.status = "out_of_steps"
|
||||
return False
|
||||
|
||||
# prepare next tasks
|
||||
prev_checkpoint = self.checkpoint
|
||||
self.checkpoint, self.tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.graph.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
for_execution=True,
|
||||
get_next_version=self.get_next_version,
|
||||
)
|
||||
|
||||
# if no more tasks, we're done
|
||||
if not self.tasks:
|
||||
self.status = "done"
|
||||
return False
|
||||
|
||||
# TODO how to make this work for both
|
||||
# - online case: we should schedule remaining tasks
|
||||
# - offline case: we should just bail, as other tasks were scheduled before
|
||||
# if there are pending writes from a previous loop, apply them
|
||||
if self.checkpoint_pending_writes:
|
||||
for tid, k, v in self.checkpoint_pending_writes:
|
||||
if task := next((t for t in self.tasks if t.id == tid), None):
|
||||
task.writes.append((k, v))
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if should_interrupt(prev_checkpoint, interrupt_before, self.tasks):
|
||||
self.status = "interrupt_before"
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def first(self) -> None:
|
||||
# map inputs to channel updates
|
||||
if input_writes := deque(map_input(self.graph.input_channels, self.input)):
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
self.checkpoint, _ = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
self.graph.nodes,
|
||||
self.channels,
|
||||
self.managed,
|
||||
self.config,
|
||||
self.step,
|
||||
for_execution=True,
|
||||
get_next_version=self.get_next_version,
|
||||
# TODO missing run_manager
|
||||
)
|
||||
# apply input writes
|
||||
apply_writes(
|
||||
self.checkpoint,
|
||||
self.channels,
|
||||
input_writes,
|
||||
self.get_next_version,
|
||||
)
|
||||
# save input checkpoint
|
||||
self.put_checkpoint({"source": "input", "writes": self.input})
|
||||
else:
|
||||
# no input is taken as signal to proceed past previous interrupt
|
||||
self.checkpoint = copy_checkpoint(self.checkpoint)
|
||||
for k in self.channels:
|
||||
if k in self.checkpoint["channel_versions"]:
|
||||
version = self.checkpoint["channel_versions"][k]
|
||||
self.checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
# done with input
|
||||
self.input = INPUT_DONE
|
||||
|
||||
def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
|
||||
if self.checkpointer is not None:
|
||||
self.submit(
|
||||
self.checkpointer.put_writes,
|
||||
{
|
||||
**self.checkpoint_config,
|
||||
"configurable": {
|
||||
**self.checkpoint_config["configurable"],
|
||||
"thread_ts": self.checkpoint["id"],
|
||||
},
|
||||
},
|
||||
writes,
|
||||
task_id,
|
||||
)
|
||||
|
||||
def put_checkpoint(
|
||||
self,
|
||||
metadata: CheckpointMetadata,
|
||||
config: RunnableConfig,
|
||||
checkpointer: Optional[BaseCheckpointSaver],
|
||||
graph: "Pregel",
|
||||
) -> None:
|
||||
# assign step
|
||||
metadata["step"] = self.step
|
||||
# bail if no checkpointer
|
||||
if self.checkpointer is not None:
|
||||
# create new checkpoint
|
||||
self.checkpoint_metadata = metadata
|
||||
self.checkpoint = create_checkpoint(
|
||||
self.checkpoint, self.channels, self.step
|
||||
self.stream = deque()
|
||||
self.stack = AsyncExitStack()
|
||||
self.input = input
|
||||
self.config = config
|
||||
self.checkpointer = checkpointer
|
||||
self.checkpointer_get_next_version = (
|
||||
checkpointer.get_next_version if checkpointer else increment
|
||||
)
|
||||
self.checkpointer_put_writes = (
|
||||
checkpointer.aput_writes if checkpointer else None
|
||||
)
|
||||
self.checkpointer_put = checkpointer.aput if checkpointer else None
|
||||
self.graph = graph
|
||||
# TODO if managed values no longer needs graph we can replace with
|
||||
# managed_specs, channel_specs
|
||||
|
||||
# context manager
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
saved = (
|
||||
await self.checkpointer.aget_tuple(self.config)
|
||||
if self.checkpointer
|
||||
else None
|
||||
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
|
||||
self.checkpoint_config = {
|
||||
**self.config,
|
||||
**saved.config,
|
||||
"configurable": {
|
||||
**self.config.get("configurable", {}),
|
||||
**saved.config.get("configurable", {}),
|
||||
},
|
||||
}
|
||||
self.checkpoint = saved.checkpoint
|
||||
self.checkpoint_metadata = saved.metadata
|
||||
self.checkpoint_pending_writes = saved.pending_writes
|
||||
|
||||
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.graph
|
||||
)
|
||||
# save it, without blocking
|
||||
self.submit(
|
||||
self.checkpointer.put,
|
||||
self.checkpoint_config,
|
||||
copy_checkpoint(self.checkpoint),
|
||||
self.checkpoint_metadata,
|
||||
)
|
||||
self.checkpoint_config = {
|
||||
**self.checkpoint_config,
|
||||
"configurable": {
|
||||
**self.checkpoint_config["configurable"],
|
||||
"thread_ts": self.checkpoint["id"],
|
||||
},
|
||||
}
|
||||
# produce debug output
|
||||
self.stream.extend(
|
||||
("debug", v)
|
||||
for v in map_debug_checkpoint(
|
||||
self.step,
|
||||
self.checkpoint_config,
|
||||
self.channels,
|
||||
self.graph.stream_channels_asis,
|
||||
self.checkpoint_metadata,
|
||||
)
|
||||
)
|
||||
# increment step
|
||||
self.step += 1
|
||||
)
|
||||
self.status = "pending"
|
||||
self.step = self.checkpoint_metadata["step"] + 1
|
||||
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: Optional[Type[BaseException]],
|
||||
exc_value: Optional[BaseException],
|
||||
traceback: Optional[TracebackType],
|
||||
) -> Optional[bool]:
|
||||
del self.graph
|
||||
return await asyncio.shield(
|
||||
self.stack.__aexit__(exc_type, exc_value, traceback)
|
||||
)
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
from collections import defaultdict
|
||||
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
|
||||
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
import pytest
|
||||
from langchain_core.callbacks import (
|
||||
CallbackManagerForLLMRun,
|
||||
@@ -25,6 +23,7 @@ from langchain_core.tools import BaseTool
|
||||
from langchain_core.tools import tool as dec_tool
|
||||
from pydantic import BaseModel as BaseModelV2
|
||||
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.prebuilt import (
|
||||
ToolNode,
|
||||
ValidationNode,
|
||||
@@ -110,6 +109,7 @@ def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
|
||||
},
|
||||
),
|
||||
"pending_sends": [],
|
||||
"current_tasks": {},
|
||||
}
|
||||
assert saved.metadata == {
|
||||
"source": "loop",
|
||||
@@ -168,6 +168,7 @@ async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
|
||||
},
|
||||
),
|
||||
"pending_sends": [],
|
||||
"current_tasks": {},
|
||||
}
|
||||
assert saved.metadata == {
|
||||
"source": "loop",
|
||||
|
||||
Reference in New Issue
Block a user