diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6ac6c737a..6e4d3e002 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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 diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index db165513b..aa256b2f7 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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() diff --git a/libs/langgraph/tests/fake_tracer.py b/libs/langgraph/tests/fake_tracer.py new file mode 100644 index 000000000..28ecc88db --- /dev/null +++ b/libs/langgraph/tests/fake_tracer.py @@ -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] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 197dc6c61..74897efcc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e9987c601..ef78c8bdb 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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",