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

Merge pull request #1418 from langchain-ai/nc/21aug/fix-pending-interrupt-run

Browse Source 
Fix pending run when interrupt exception is used
This commit is contained in:
Nuno Campos
2024-08-21 14:29:59 -07:00
committed by GitHub
parent 06c2481783 47ed3d97e9
commit 46171dd759
5 changed files with 143 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+6 -10
View File
@@ -961,6 +961,7 @@ class Pregel(
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
@@ -969,7 +970,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
@@ -1088,8 +1088,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
@@ -1218,6 +1218,7 @@ class Pregel(
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
@@ -1227,7 +1228,6 @@ class Pregel(
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
output_keys=output_keys,
stream_keys=self.stream_channels_asis,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
@@ -1349,13 +1349,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/loop.py
+24 -6
View File
@@ -62,7 +62,13 @@ 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
@@ -83,6 +89,7 @@ class PregelLoop:
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]
@@ -117,6 +124,7 @@ class PregelLoop:
]
tasks: Sequence[PregelExecutableTask]
stream: deque[Tuple[str, Any]]
output: Union[None, dict[str, Any], Any] = None
# public
@@ -129,6 +137,7 @@ class PregelLoop:
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]],
) -> None:
self.stream = deque()
self.input = input
@@ -137,6 +146,7 @@ class PregelLoop:
self.checkpointer = checkpointer
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:
@@ -167,7 +177,6 @@ class PregelLoop:
self,
*,
input_keys: Union[str, Sequence[str]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
@@ -196,7 +205,7 @@ class PregelLoop:
# 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()
@@ -204,7 +213,7 @@ class PregelLoop:
self._put_checkpoint(
{
"source": "loop",
"writes": single(map_output_updates(output_keys, self.tasks)),
"writes": single(map_output_updates(self.output_keys, self.tasks)),
}
)
# after execution, check if we should interrupt
@@ -272,7 +281,7 @@ class PregelLoop:
if all(task.writes for task in self.tasks):
return self.tick(
input_keys=input_keys,
output_keys=output_keys,
stream_keys=stream_keys,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before,
manager=manager,
@@ -406,7 +415,12 @@ class PregelLoop:
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
@@ -420,6 +434,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -428,6 +443,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
store=store,
nodes=nodes,
specs=specs,
output_keys=output_keys,
)
self.stack = ExitStack()
if checkpointer:
@@ -505,6 +521,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
checkpointer: Optional[BaseCheckpointSaver],
nodes: Mapping[str, PregelNode],
specs: Mapping[str, Union[BaseChannel, ManagedValueSpec]],
output_keys: Union[str, Sequence[str]] = EMPTY_SEQ,
) -> None:
super().__init__(
input,
@@ -513,6 +530,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
store=store,
nodes=nodes,
specs=specs,
output_keys=output_keys,
)
self.store = AsyncBatchedStore(self.store) if self.store else None
self.stack = AsyncExitStack()
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_pregel.py
+11 -1
View File
@@ -68,6 +68,7 @@ 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,
@@ -6159,11 +6160,20 @@ def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
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",
libs/langgraph/tests/test_pregel_async.py
+11 -1
View File
@@ -63,6 +63,7 @@ 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,
@@ -225,11 +226,20 @@ async def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}) == {
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",