From 02326d74bf1d62dcfcac383a7adb55954555dcf3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 08:11:46 -0700 Subject: [PATCH 01/11] WIP --- tests/__snapshots__/test_pregel_async.ambr | 100 --------------------- tests/test_pregel_async.py | 20 ++++- 2 files changed, 16 insertions(+), 104 deletions(-) delete mode 100644 tests/__snapshots__/test_pregel_async.ambr diff --git a/tests/__snapshots__/test_pregel_async.ambr b/tests/__snapshots__/test_pregel_async.ambr deleted file mode 100644 index 809ff8731..000000000 --- a/tests/__snapshots__/test_pregel_async.ambr +++ /dev/null @@ -1,100 +0,0 @@ -# serializer version: 1 -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** ... - * . - ** ... - +--------------+ . - | analyzer_one | . - +--------------+ . - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +---------------+ - | rewrite_query | - +---------------+ - *** ... - * . - ** ... - +--------------+ . - | analyzer_one | . - +--------------+ . - * . - * . - * . - +---------------+ +---------------+ - | retriever_one | | retriever_two | - +---------------+ +---------------+ - *** *** - * * - ** ** - +----+ - | qa | - +----+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- -# name: test_nested_graph - ''' - +-----------+ - | __start__ | - +-----------+ - * - * - * - +-------+ - | inner | - +-------+ - * - * - * - +------+ - | side | - +------+ - * - * - * - +---------+ - | __end__ | - +---------+ - ''' -# --- diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 34e4a783c..204523765 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -10,6 +10,7 @@ from typing import ( AsyncIterator, Dict, Generator, + NamedTuple, Optional, Sequence, TypedDict, @@ -2496,6 +2497,7 @@ async def test_state_graph_few_shot() -> None: class BaseState(TypedDict): messages: Annotated[list[AnyMessage], add_messages] + tool_results: Annotated[list[str], operator.add] class AgentState(BaseState): examples: Annotated[ @@ -2547,19 +2549,29 @@ Some examples of past conversations: response = await model.ainvoke(formatted) return {"messages": response} + class Do(NamedTuple): + node: str + kwargs: dict[str, Any] + # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit - if not data["messages"][-1].tool_calls: - return "exit" + if tool_calls := data["messages"][-1].tool_calls: + return [ + Do(node="tools", kwargs={"tool_call": tool_call}) + for tool_call in tool_calls + ] else: - return "continue" + return "exit" + + def tools_node(data: AgentState, *, tool_call: ToolCall) -> AgentState: + return {"tool_results": ...} # Define a new graph workflow = StateGraph(AgentState) workflow.add_node("agent", agent) - workflow.add_node("tools", ToolNode(tools)) + workflow.add_node("tools", tools_node) workflow.set_entry_point("agent") workflow.add_conditional_edges( "agent", should_continue, {"continue": "tools", "exit": END} From 7a2de5e369e8c941e1521730900a82b14c281d48 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 10:42:05 -0700 Subject: [PATCH 02/11] WIP --- tests/test_pregel_async.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 204523765..211e4ab6d 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2485,7 +2485,13 @@ async def test_state_graph_few_shot() -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) - from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage + from langchain_core.messages import ( + AIMessage, + AnyMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) from langchain_core.prompts import ChatPromptTemplate from langchain_core.tools import tool @@ -2549,18 +2555,15 @@ Some examples of past conversations: response = await model.ainvoke(formatted) return {"messages": response} - class Do(NamedTuple): - node: str - kwargs: dict[str, Any] + class GoTo: + def __init__(self, /, __node__: str, **kwargs: Any) -> None: + pass # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [ - Do(node="tools", kwargs={"tool_call": tool_call}) - for tool_call in tool_calls - ] + return [GoTo("tools", tool_call=tool_call) for tool_call in tool_calls] else: return "exit" From 7a1feb466b04d04b301e89e61af487df7e4d997b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 15:36:32 -0700 Subject: [PATCH 03/11] Implement map-reduce api for StateGraph/Pregel - current Pregel primitive is pull-based, ie. nodes write to channels, and it's up to other nodes to subscribe to those channels to "pull" updates - this adds a "push" primitive where a node can directly schedule a node (more than once if desired) for execution in the next step, with additional kwargs to be passed in. Nodes scheduled in this way get called with both the current state and any kwargs passed to Packet --- langgraph/channels/base.py | 1 + langgraph/checkpoint/base.py | 7 + langgraph/constants.py | 17 ++ langgraph/graph/graph.py | 39 +-- langgraph/graph/state.py | 21 +- langgraph/pregel/__init__.py | 199 +++++++++++---- langgraph/pregel/debug.py | 21 +- langgraph/pregel/io.py | 58 +++-- langgraph/pregel/read.py | 19 +- langgraph/pregel/types.py | 1 + langgraph/pregel/validate.py | 10 +- langgraph/pregel/write.py | 61 +++-- langgraph/serde/jsonplus.py | 5 + langgraph/utils.py | 30 +-- tests/test_pregel_async.py | 473 ++++++++++++++++++++++++++++++++++- 15 files changed, 813 insertions(+), 149 deletions(-) diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index a2a4c2d0d..9dc5fcbbc 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -128,6 +128,7 @@ def create_checkpoint( channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], + pending_packets=checkpoint["pending_packets"], ) diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 2bdb52055..673a83278 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -5,6 +5,7 @@ from typing import ( Any, AsyncIterator, Iterator, + List, Literal, NamedTuple, Optional, @@ -14,6 +15,7 @@ from typing import ( from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig from langgraph.checkpoint.id import uuid6 +from langgraph.constants import Packet from langgraph.serde.base import SerializerProtocol from langgraph.serde.jsonplus import JsonPlusSerializer @@ -72,6 +74,9 @@ class Checkpoint(TypedDict): Used to determine which nodes to execute next. """ + pending_packets: List[Packet] + """List of packets sent to nodes but not yet processed. + Cleared by the next checkpoint.""" def _seen_dict(): @@ -86,6 +91,7 @@ def empty_checkpoint() -> Checkpoint: channel_values={}, channel_versions=defaultdict(int), versions_seen=defaultdict(_seen_dict), + pending_packets=[], ) @@ -100,6 +106,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: _seen_dict, {k: defaultdict(int, v) for k, v in checkpoint["versions_seen"].items()}, ), + pending_packets=checkpoint["pending_packets"].copy(), ) diff --git a/langgraph/constants.py b/langgraph/constants.py index e75e97452..d880aa733 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,6 +1,23 @@ +from typing import Any + +from langgraph.errors import InvalidUpdateError + CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" INTERRUPT = "__interrupt__" +TASKS = "__pregel_tasks" + +RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ} TAG_HIDDEN = "langsmith:hidden" + + +class Packet: + def __init__(self, /, __node__: str, **kwargs: Any) -> None: + if not kwargs: + raise InvalidUpdateError( + "Packet must have at least one keyword argument to pass to node" + ) + self.node = __node__ + self.kwargs = kwargs diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 5659b5031..4e00d148a 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -26,7 +26,8 @@ from langchain_core.runnables.graph import ( from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import TAG_HIDDEN, Packet +from langgraph.errors import InvalidUpdateError from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All @@ -69,13 +70,7 @@ class Branch(NamedTuple): writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: result = self.path.invoke(reader(config) if reader else input, config) - if not isinstance(result, list): - result = [result] - if self.ends: - destinations = [self.ends[r] for r in result] - else: - destinations = result - return writer(destinations) or input + return self._finish(writer, input, result) async def _aroute( self, @@ -86,14 +81,25 @@ class Branch(NamedTuple): writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: result = await self.path.ainvoke(reader(config) if reader else input, config) + return self._finish(writer, input, result) + + def _finish( + self, writer: Callable[[list[str]], Optional[Runnable]], input: Any, result: Any + ): if not isinstance(result, list): result = [result] if self.ends: - destinations = [self.ends[r] for r in result] + destinations = [ + r if isinstance(r, Packet) else self.ends[r] for r in result + ] else: destinations = result - if any(dest is None for dest in destinations): + if any(dest is None or dest == START for dest in destinations): raise ValueError("Branch did not return a valid destination") + if any(p.node == END for p in destinations if isinstance(p, Packet)): + raise InvalidUpdateError( + "Cannot send a packet with keyword arguments to the END node" + ) return writer(destinations) or input @@ -384,13 +390,14 @@ class CompiledGraph(Pregel): self.nodes[end].channels.append(start) def attach_branch(self, start: str, name: str, branch: Branch) -> None: - def branch_writer(ends: list[str]) -> Optional[ChannelWrite]: - channels = [ - f"branch:{start}:{name}:{end}" if end != END else END for end in ends + def branch_writer(packets: list[Union[str, Packet]]) -> Optional[ChannelWrite]: + writes = [ + ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END) + if not isinstance(p, Packet) + else p + for p in packets ] - return ChannelWrite( - [ChannelWriteEntry(ch) for ch in channels], tags=[TAG_HIDDEN] - ) + return ChannelWrite(writes, tags=[TAG_HIDDEN]) # add hidden start node if start == START and start not in self.nodes: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 02df2cc3a..01b16aa69 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -27,7 +27,7 @@ from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.constants import TAG_HIDDEN from langgraph.errors import InvalidUpdateError -from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph +from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Packet from langgraph.managed.base import ManagedValue, is_managed_value from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All @@ -325,7 +325,7 @@ class CompiledStateGraph(CompiledGraph): ], ) else: - self.channels[key] = EphemeralValue(Any) + self.channels[key] = EphemeralValue(Any, guard=False) self.nodes[key] = PregelNode( triggers=[], # read state keys and managed values @@ -377,17 +377,24 @@ class CompiledStateGraph(CompiledGraph): ) def attach_branch(self, start: str, name: str, branch: Branch) -> None: - def branch_writer(ends: list[str]) -> Optional[ChannelWrite]: - if filtered_ends := [end for end in ends if end != END]: + def branch_writer(packets: list[Union[str, Packet]]) -> Optional[ChannelWrite]: + if filtered := [p for p in packets if p != END]: writes = [ - ChannelWriteEntry(f"branch:{start}:{name}:{end}", start) - for end in filtered_ends + ChannelWriteEntry(f"branch:{start}:{name}:{p}", start) + if not isinstance(p, Packet) + else p + for p in filtered ] if branch.then and branch.then != END: writes.append( ChannelWriteEntry( f"branch:{start}:{name}:then", - WaitForNames(set(filtered_ends)), + WaitForNames( + { + p.node if isinstance(p, Packet) else p + for p in filtered + } + ), ) ) return ChannelWrite(writes, tags=[TAG_HIDDEN]) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index bd4fea0c1..1b50065ea 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -66,6 +66,8 @@ from langgraph.constants import ( CONFIG_KEY_SEND, INTERRUPT, TAG_HIDDEN, + TASKS, + Packet, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -877,7 +879,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _ in next_tasks: + for _, _, _, writes, _, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1180,7 +1182,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _ in next_tasks: + for _, _, _, writes, _, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1444,10 +1446,10 @@ def _should_interrupt( checkpoint["channel_versions"][chan] > seen[chan] for chan in snapshot_channels ) - # and any channel written to is in interrupt_nodes list + # and any triggered node is in interrupt_nodes list and any( node - for node, _, _, _, config, _ in tasks + for node, _, _, _, config, _, _ in tasks if ( (not config or TAG_HIDDEN not in config.get("tags")) if interrupt_nodes == "*" @@ -1473,15 +1475,40 @@ def _local_read( return read_channels(channels, select) +def _local_write( + commit: Callable[[Sequence[tuple[str, Any]]], None], + processes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + writes: Sequence[tuple[str, Any]], +) -> None: + for chan, value in writes: + if chan == TASKS: + if not isinstance(value, Packet): + raise InvalidUpdateError( + f"Invalid packet type, expected Packet, got {value}" + ) + if value.node not in processes: + raise InvalidUpdateError(f"Invalid node name {value.node} in packet") + elif chan not in channels: + logger.warning(f"Skipping write for channel '{chan}' which has no readers") + commit(writes) + + def _apply_writes( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], pending_writes: Sequence[tuple[str, Any]], ) -> None: + if checkpoint["pending_packets"]: + raise RuntimeError("Cannot apply writes when there are pending packets") + pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: - pending_writes_by_channel[chan].append(val) + if chan == TASKS: + checkpoint["pending_packets"].append(val) + else: + pending_writes_by_channel[chan].append(val) # Find the highest version of all channels if checkpoint["channel_versions"]: @@ -1497,12 +1524,10 @@ def _apply_writes( channels[chan].update(vals) except InvalidUpdateError as e: raise InvalidUpdateError( - f"Invalid update for channel {chan}: {e}" + f"Invalid update for channel {chan} with values {vals}" ) from e checkpoint["channel_versions"][chan] = max_version + 1 updated_channels.add(chan) - else: - logger.warning(f"Skipping write for channel '{chan}' which has no readers") # Channels that weren't updated in this step are notified of a new step for chan in channels: if chan not in updated_channels: @@ -1550,6 +1575,64 @@ def _prepare_next_tasks( ) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: checkpoint = copy_checkpoint(checkpoint) tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] + # Consume pending packets + for packet in checkpoint["pending_packets"]: + try: + val = next( + _proc_input( + step, + packet.node, + processes[packet.node], + managed, + channels, + catch=True, + ) + ) + except StopIteration: + logger.warn('No input for node "%s" in packet, skipping', packet.node) + continue + if for_execution: + if node := processes[packet.node].get_node(packet.kwargs): + writes = deque() + tasks.append( + PregelExecutableTask( + packet.node, + val, + node, + writes, + patch_config( + merge_configs( + config, + processes[packet.node].config, + { + "metadata": { + "langgraph_step": step, + "langgraph_node": packet.node, + "langgraph_triggers": [TASKS], + } + }, + ), + run_name=packet.node, + callbacks=manager.get_child(f"graph:step:{step}") + if manager + else None, + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + _local_write, writes.extend, processes, channels + ), + CONFIG_KEY_READ: partial( + _local_read, checkpoint, channels, tasks + ), + }, + ), + [TASKS], + packet.kwargs, + ) + ) + else: + tasks.append(PregelTaskDescription(packet.node, val)) + checkpoint["pending_packets"].clear() # Check if any processes should be run in next step # If so, prepare the values to be passed to them for name, proc in processes.items(): @@ -1563,43 +1646,10 @@ def _prepare_next_tasks( ) and checkpoint["channel_versions"][chan] > seen[chan] ]: - # If all trigger channels subscribed by this process are not empty - # then invoke the process with the values of all non-empty channels - if isinstance(proc.channels, dict): - try: - val: dict = { - k: read_channel(channels, chan, catch=chan not in proc.triggers) - for k, chan in proc.channels.items() - if isinstance(chan, str) - } - - managed_values = {} - for key, chan in proc.channels.items(): - if is_managed_value(chan): - managed_values[key] = managed[key]( - step, PregelTaskDescription(name, val) - ) - - val.update(managed_values) - except EmptyChannelError: - continue - elif isinstance(proc.channels, list): - for chan in proc.channels: - try: - val = read_channel(channels, chan, catch=False) - break - except EmptyChannelError: - pass - else: - continue - else: - raise RuntimeError( - "Invalid channels type, expected list or dict, got {proc.channels}" - ) - - # If the process has a mapper, apply it to the value - if proc.mapper is not None: - val = proc.mapper(val) + try: + val = next(_proc_input(step, name, proc, managed, channels)) + except StopIteration: + continue # update seen versions if for_execution: @@ -1613,6 +1663,7 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): writes = deque() + triggers = sorted(triggers) tasks.append( PregelExecutableTask( name, @@ -1627,6 +1678,7 @@ def _prepare_next_tasks( "metadata": { "langgraph_step": step, "langgraph_node": name, + "langgraph_triggers": triggers, } }, ), @@ -1636,13 +1688,15 @@ def _prepare_next_tasks( else None, configurable={ # deque.extend is thread-safe - CONFIG_KEY_SEND: writes.extend, + CONFIG_KEY_SEND: partial( + _local_write, writes.extend, processes, channels + ), CONFIG_KEY_READ: partial( _local_read, checkpoint, channels, writes ), }, ), - sorted(triggers), + triggers, ) ) else: @@ -1650,6 +1704,57 @@ def _prepare_next_tasks( return checkpoint, tasks +def _proc_input( + step: int, + name: str, + proc: PregelNode, + managed: ManagedValueMapping, + channels: Mapping[str, BaseChannel], + catch: bool = False, +) -> Iterator[Any]: + # If all trigger channels subscribed by this process are not empty + # then invoke the process with the values of all non-empty channels + if isinstance(proc.channels, dict): + try: + val: dict = { + k: read_channel( + channels, chan, catch=catch or chan not in proc.triggers + ) + for k, chan in proc.channels.items() + if isinstance(chan, str) + } + + managed_values = {} + for key, chan in proc.channels.items(): + if is_managed_value(chan): + managed_values[key] = managed[key]( + step, PregelTaskDescription(name, val) + ) + + val.update(managed_values) + except EmptyChannelError: + return + elif isinstance(proc.channels, list): + for chan in proc.channels: + try: + val = read_channel(channels, chan, catch=False) + break + except EmptyChannelError: + pass + else: + return + else: + raise RuntimeError( + "Invalid channels type, expected list or dict, got {proc.channels}" + ) + + # If the process has a mapper, apply it to the value + if proc.mapper is not None: + val = proc.mapper(val) + + yield val + + def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]: if on: for chunk in iter: diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 0fd00bdc8..6dc86643f 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -64,20 +64,23 @@ def map_debug_tasks( step: int, tasks: list[PregelExecutableTask] ) -> Iterator[DebugOutputTask]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, input, _, _, config, triggers) in enumerate(tasks): + for idx, (name, input, _, _, config, triggers, kwargs) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue + payload = { + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), + "name": name, + "input": input, + "triggers": triggers, + } + if kwargs is not None: + payload["kwargs"] = kwargs yield { "type": "task", "timestamp": ts, "step": step, - "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), - "name": name, - "input": input, - "triggers": triggers, - }, + "payload": payload, } @@ -87,7 +90,7 @@ def map_debug_task_results( stream_channels_list: Sequence[str], ) -> Iterator[DebugOutputTaskResult]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, _, _, writes, config, _) in enumerate(tasks): + for idx, (name, _, _, writes, config, _, _) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -130,7 +133,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: ) + "\n".join( f"- {get_colored_text(name, 'green')} -> {pformat(val)}" - for name, val, _, _, _, _ in next_tasks + for name, val, _, _, _, _, _ in next_tasks ) ) diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 6228167cc..3e2bf173d 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,9 +1,11 @@ +from collections import defaultdict +from itertools import groupby from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union from langchain_core.runnables.utils import AddableDict from langgraph.channels.base import BaseChannel, EmptyChannelError -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import TAG_HIDDEN, TASKS from langgraph.pregel.log import logger from langgraph.pregel.types import PregelExecutableTask @@ -102,24 +104,44 @@ def map_output_updates( t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags") ] if isinstance(output_channels, str): - if updated := AddableUpdatesDict( - { - node: value - for node, _, _, writes, _, _ in output_tasks - for chan, value in writes - if chan == output_channels - } - ): - yield updated + if updated := [ + (triggers == [TASKS], node, value) + for node, _, _, writes, _, triggers, _ in output_tasks + for chan, value in writes + if chan == output_channels + ]: + grouped = defaultdict(list) + for from_packet, node, value in updated: + if from_packet: + grouped[node].append(value) + for from_packet, node, value in updated: + if not from_packet: + if grouped[node]: + grouped[node].append(value) + else: + grouped[node] = value + yield AddableUpdatesDict(grouped) else: - if updated := AddableUpdatesDict( - { - node: {chan: value for chan, value in writes if chan in output_channels} - for node, _, _, writes, _, _ in output_tasks - if any(chan in output_channels for chan, _ in writes) - } - ): - yield updated + if updated := [ + ( + triggers == [TASKS], + node, + {chan: value for chan, value in writes if chan in output_channels}, + ) + for node, _, _, writes, _, triggers, _ in output_tasks + if any(chan in output_channels for chan, _ in writes) + ]: + grouped = defaultdict(list) + for from_packet, node, value in updated: + if from_packet: + grouped[node].append(value) + for from_packet, node, value in updated: + if not from_packet: + if grouped[node]: + grouped[node].append(value) + else: + grouped[node] = value + yield AddableUpdatesDict(grouped) T = TypeVar("T") diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index eefed2b47..2e90a505b 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( @@ -121,11 +121,17 @@ class PregelNode(RunnableBindingBase): and isinstance(writers[-2], ChannelWrite) ): # we can combine writes if they are consecutive - writers[-2].writes += writers[-1].writes + # careful to not modify the original writers list or ChannelWrite + writers[-2] = ChannelWrite( + writes=writers[-2].writes + writers[-1].writes, + tags=writers[-2].config["tags"] if writers[-2].config else None, + ) writers.pop() return writers - def get_node(self) -> Optional[Runnable[Any, Any]]: + def get_node( + self, kwargs: Optional[Dict[str, Any]] = None + ) -> Optional[Runnable[Any, Any]]: writers = self.get_writers() if self.bound is DEFAULT_BOUND and not writers: return None @@ -134,9 +140,12 @@ class PregelNode(RunnableBindingBase): elif self.bound is DEFAULT_BOUND: return RunnableSequence(*writers) elif writers: - return RunnableSequence(self.bound, *writers) + return RunnableSequence( + self.bound.bind(**kwargs) if kwargs is not None else self.bound, + *writers, + ) else: - return self.bound + return self.bound.bind(**kwargs) if kwargs is not None else self.bound def __init__( self, diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index a3826314a..ea9cb69ef 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -18,6 +18,7 @@ class PregelExecutableTask(NamedTuple): writes: deque[tuple[str, Any]] config: Optional[RunnableConfig] triggers: list[str] + kwargs: Optional[dict[str, Any]] = None class StateSnapshot(NamedTuple): diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index ad4959657..8627642e9 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -1,7 +1,7 @@ from typing import Mapping, Optional, Sequence, Union from langgraph.channels.base import BaseChannel -from langgraph.constants import INTERRUPT +from langgraph.constants import RESERVED from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All @@ -16,13 +16,13 @@ def validate_graph( interrupt_before_nodes: Union[All, Sequence[str]], ) -> None: for chan in channels: - if chan == INTERRUPT: - raise ValueError(f"Channel name {INTERRUPT} is reserved") + if chan in RESERVED: + raise ValueError(f"Channel names {RESERVED} are reserved") subscribed_channels = set[str]() for name, node in nodes.items(): - if name == INTERRUPT: - raise ValueError(f"Node name {INTERRUPT} is reserved") + if name in RESERVED: + raise ValueError(f"Node names {RESERVED} are reserved") if isinstance(node, PregelNode): subscribed_channels.update(node.triggers) else: diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 4b2ad9ea5..f9e9c1233 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,12 +1,14 @@ from __future__ import annotations import asyncio -from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar +from typing import Any, Callable, List, NamedTuple, Optional, Sequence, Tuple, TypeVar from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec +from sympy import Union -from langgraph.constants import CONFIG_KEY_SEND +from langgraph.constants import CONFIG_KEY_SEND, TASKS, Packet +from langgraph.errors import InvalidUpdateError from langgraph.utils import RunnableCallable TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] @@ -25,7 +27,7 @@ class ChannelWriteEntry(NamedTuple): class ChannelWrite(RunnableCallable): - writes: Sequence[ChannelWriteEntry] + writes: Sequence[Union[ChannelWriteEntry, Packet]] """ Sequence of write entries, each of which is a tuple of: - channel name @@ -34,7 +36,10 @@ class ChannelWrite(RunnableCallable): """ def __init__( - self, writes: Sequence[ChannelWriteEntry], *, tags: Optional[list[str]] = None + self, + writes: Sequence[Union[ChannelWriteEntry, Packet]], + *, + tags: Optional[list[str]] = None, ): super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) self.writes = writes @@ -46,7 +51,7 @@ class ChannelWrite(RunnableCallable): self, suffix: Optional[str] = None, *, name: Optional[str] = None ) -> str: if not name: - name = f"ChannelWrite<{','.join(chan for chan, _, _, _ in self.writes)}>" + name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else w.node for w in self.writes)}>" return super().get_name(suffix, name=name) @property @@ -62,47 +67,69 @@ class ChannelWrite(RunnableCallable): ] def _write(self, input: Any, config: RunnableConfig) -> None: + # split packets and entries + writes = [ + (TASKS, packet) for packet in self.writes if isinstance(packet, Packet) + ] + entries = [ + write for write in self.writes if isinstance(write, ChannelWriteEntry) + ] + for entry in entries: + if entry.channel == TASKS: + raise InvalidUpdateError("Cannot write to the reserved channel TASKS") + # process entries into values values = [ - input if write.value is PASSTHROUGH else write.value - for write in self.writes + input if write.value is PASSTHROUGH else write.value for write in entries ] values = [ val if write.mapper is None else write.mapper.invoke(val, config) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) ] values = [ (write.channel, val) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) if not write.skip_none or val is not None ] - self.do_write(config, **dict(values)) + # write packets and values + self.do_write(config, writes + values) return input async def _awrite(self, input: Any, config: RunnableConfig) -> None: + # split packets and entries + writes = [ + (TASKS, packet) for packet in self.writes if isinstance(packet, Packet) + ] + entries = [ + write for write in self.writes if isinstance(write, ChannelWriteEntry) + ] + for entry in entries: + if entry.channel == TASKS: + raise InvalidUpdateError("Cannot write to the reserved channel TASKS") + # process entries into values values = [ - input if write.value is PASSTHROUGH else write.value - for write in self.writes + input if write.value is PASSTHROUGH else write.value for write in entries ] values = await asyncio.gather( *( _mk_future(val) if write.mapper is None else write.mapper.ainvoke(val, config) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) ) ) values = [ (write.channel, val) - for val, write in zip(values, self.writes) + for val, write in zip(values, entries) if not write.skip_none or val is not None ] - self.do_write(config, **dict(values)) + # write packets and values + self.do_write(config, writes + values) return input @staticmethod - def do_write(config: RunnableConfig, **values: Any) -> None: + def do_write(config: RunnableConfig, values: List[Tuple[str, Any]]) -> None: write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE]) + write([(chan, val) for chan, val in values if val is not SKIP_WRITE]) @staticmethod def is_writer(runnable: Runnable) -> bool: diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py index e2e836689..2b89f5af7 100644 --- a/langgraph/serde/jsonplus.py +++ b/langgraph/serde/jsonplus.py @@ -9,6 +9,7 @@ from uuid import UUID from langchain_core.load.load import Reviver from langchain_core.load.serializable import Serializable +from langgraph.constants import Packet from langgraph.serde.base import SerializerProtocol LC_REVIVER = Reviver() @@ -63,6 +64,10 @@ class JsonPlusSerializer(SerializerProtocol): ) elif isinstance(obj, Enum): return self._encode_constructor_args(obj.__class__, args=[obj.value]) + elif isinstance(obj, Packet): + return self._encode_constructor_args( + Packet, args=[obj.node], kwargs=obj.kwargs + ) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" diff --git a/langgraph/utils.py b/langgraph/utils.py index e766a7fb8..7d94b81e9 100644 --- a/langgraph/utils.py +++ b/langgraph/utils.py @@ -59,7 +59,7 @@ class RunnableCallable(Runnable): pass self.func = func self.afunc = afunc - self.config = {"tags": tags} if tags else None + self.config: Optional[RunnableConfig] = {"tags": tags} if tags else None self.kwargs = kwargs self.trace = trace self.recurse = recurse @@ -72,47 +72,47 @@ class RunnableCallable(Runnable): } return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})" - def invoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + def invoke( + self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Any: if self.func is None: raise TypeError( f'No synchronous function provided to "{self.name}".' "\nEither initialize with a synchronous function or invoke" " via the async API (ainvoke, astream, etc.)" ) + kwargs = {**self.kwargs, **kwargs} if self.trace: ret = self._call_with_config( - self.func, input, merge_configs(self.config, config), **self.kwargs + self.func, input, merge_configs(self.config, config), **kwargs ) else: config = merge_configs(self.config, config) context = copy_context() context.run(var_child_runnable_config.set, config) - kwargs = ( - {**self.kwargs, "config": config} - if accepts_config(self.func) - else self.kwargs - ) + if accepts_config(self.func): + kwargs["config"] = config ret = context.run(self.func, input, **kwargs) if isinstance(ret, Runnable) and self.recurse: return ret.invoke(input, config) return ret - async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None) -> Any: + async def ainvoke( + self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Any: if not self.afunc: return self.invoke(input, config) + kwargs = {**self.kwargs, **kwargs} if self.trace: ret = await self._acall_with_config( - self.afunc, input, merge_configs(self.config, config), **self.kwargs + self.afunc, input, merge_configs(self.config, config), **kwargs ) else: config = merge_configs(self.config, config) context = copy_context() context.run(var_child_runnable_config.set, config) - kwargs = ( - {**self.kwargs, "config": config} - if accepts_config(self.afunc) - else self.kwargs - ) + if accepts_config(self.afunc): + kwargs["config"] = config if sys.version_info >= (3, 11): ret = await asyncio.create_task( self.afunc(input, **kwargs), context=context diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 211e4ab6d..7515c0351 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -10,7 +10,6 @@ from typing import ( AsyncIterator, Dict, Generator, - NamedTuple, Optional, Sequence, TypedDict, @@ -19,7 +18,12 @@ from typing import ( from uuid import UUID import pytest -from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough +from langchain_core.runnables import ( + RunnableConfig, + RunnableLambda, + RunnablePassthrough, + RunnablePick, +) from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -28,6 +32,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver +from langgraph.constants import Packet from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START @@ -2503,7 +2508,7 @@ async def test_state_graph_few_shot() -> None: class BaseState(TypedDict): messages: Annotated[list[AnyMessage], add_messages] - tool_results: Annotated[list[str], operator.add] + # tool_results: Annotated[list[str], operator.add] class AgentState(BaseState): examples: Annotated[ @@ -2518,6 +2523,7 @@ async def test_state_graph_few_shot() -> None: return f"result for {query}" tools = [search_api] + tools_by_name = {t.name: t for t in tools} prompt = ChatPromptTemplate.from_messages( [ @@ -2555,20 +2561,23 @@ Some examples of past conversations: response = await model.ainvoke(formatted) return {"messages": response} - class GoTo: - def __init__(self, /, __node__: str, **kwargs: Any) -> None: - pass - # Define decision-making logic def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [GoTo("tools", tool_call=tool_call) for tool_call in tool_calls] + return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] else: return "exit" - def tools_node(data: AgentState, *, tool_call: ToolCall) -> AgentState: - return {"tool_results": ...} + def tools_node( + _: AgentState, config: RunnableConfig, *, tool_call: ToolCall + ) -> AgentState: + output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + return { + "messages": ToolMessage( + content=output, name=tool_call["name"], tool_call_id=tool_call["id"] + ) + } # Define a new graph workflow = StateGraph(AgentState) @@ -3061,6 +3070,450 @@ async def test_prebuilt_chat() -> None: ] +async def test_state_graph_packets() -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) + from langchain_core.tools import tool + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + tools_by_name = {t.name: t for t in tools} + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + else: + return END + + def tools_node( + _: AgentState, config: RunnableConfig, *, tool_call: ToolCall + ) -> AgentState: + output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + return { + "messages": ToolMessage( + content=output, name=tool_call["name"], tool_call_id=tool_call["id"] + ) + } + + # Define a new graph + workflow = StateGraph(AgentState) + + # Define the two nodes we will cycle between + workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) + workflow.add_node("tools", tools_node) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges("agent", should_continue) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert await app.ainvoke( + {"messages": HumanMessage(content="what is weather in sf")} + ) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ), + ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + AIMessage(content="answer", id="ai3"), + ] + } + + assert [ + c + async for c in app.astream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="a1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + }, + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="a2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + } + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ) + }, + { + "messages": ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + }, + ] + }, + {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ) + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ), + ], + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "loop", + "step": 1, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + id="ai1", + ) + }, + }, + ) + + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"' + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + ], + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "update", + "step": 2, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ) + }, + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ) + }, + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + id="ai2", + ) + }, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"another"'} + }, + id="ai2", + ), + ], + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"another"', + } + }, + id="ai2", + ) + }, + }, + ) + + await app_w_interrupt.aupdate_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + content="", + additional_kwargs={ + "function_call": { + "name": "search_api", + "arguments": '"a different query"', + } + }, + id="ai1", + ), + FunctionMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + ), + AIMessage(content="answer", id="ai2"), + ], + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + }, + ) + + async def test_message_graph() -> None: from langchain_core.agents import AgentAction from langchain_core.language_models.fake_chat_models import ( From 75f14cc06df8f1878c88926efdaa9bfb6ab9d488 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 16:08:30 -0700 Subject: [PATCH 04/11] Finish tests --- langgraph/pregel/__init__.py | 6 +- langgraph/pregel/io.py | 1 - tests/__snapshots__/test_pregel_async.ambr | 100 +++++ tests/test_pregel.py | 500 +++++++++++++++++++++ tests/test_pregel_async.py | 342 ++++++++------ 5 files changed, 806 insertions(+), 143 deletions(-) create mode 100644 tests/__snapshots__/test_pregel_async.ambr diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 1b50065ea..aaf6662ac 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1500,7 +1500,7 @@ def _apply_writes( pending_writes: Sequence[tuple[str, Any]], ) -> None: if checkpoint["pending_packets"]: - raise RuntimeError("Cannot apply writes when there are pending packets") + checkpoint["pending_packets"].clear() pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel @@ -1630,8 +1630,8 @@ def _prepare_next_tasks( packet.kwargs, ) ) - else: - tasks.append(PregelTaskDescription(packet.node, val)) + else: + tasks.append(PregelTaskDescription(packet.node, val)) checkpoint["pending_packets"].clear() # Check if any processes should be run in next step # If so, prepare the values to be passed to them diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 3e2bf173d..e21a5f4f0 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,5 +1,4 @@ from collections import defaultdict -from itertools import groupby from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union from langchain_core.runnables.utils import AddableDict diff --git a/tests/__snapshots__/test_pregel_async.ambr b/tests/__snapshots__/test_pregel_async.ambr new file mode 100644 index 000000000..809ff8731 --- /dev/null +++ b/tests/__snapshots__/test_pregel_async.ambr @@ -0,0 +1,100 @@ +# serializer version: 1 +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +---------------+ + | rewrite_query | + +---------------+ + *** ... + * . + ** ... + +--------------+ . + | analyzer_one | . + +--------------+ . + * . + * . + * . + +---------------+ +---------------+ + | retriever_one | | retriever_two | + +---------------+ +---------------+ + *** *** + * * + ** ** + +----+ + | qa | + +----+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- +# name: test_nested_graph + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | inner | + +-------+ + * + * + * + +------+ + | side | + +------+ + * + * + * + +---------+ + | __end__ | + +---------+ + ''' +# --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index c58ff7e7b..3d18fcd6c 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -34,6 +34,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.sqlite import SqliteSaver +from langgraph.constants import Packet from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph from langgraph.graph.graph import START @@ -3427,6 +3428,505 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: ] +def test_state_graph_packets() -> None: + from langchain_core.language_models.fake_chat_models import ( + FakeMessagesListChatModel, + ) + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + ToolCall, + ToolMessage, + ) + from langchain_core.tools import tool + + class AgentState(TypedDict): + messages: Annotated[list[BaseMessage], add_messages] + + @tool() + def search_api(query: str) -> str: + """Searches the API for the query.""" + return f"result for {query}" + + tools = [search_api] + tools_by_name = {t.name: t for t in tools} + + model = FakeMessagesListChatModel( + responses=[ + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(id="ai3", content="answer"), + ] + ) + + # Define decision-making logic + def should_continue(data: AgentState) -> str: + # Logic to decide whether to continue in the loop or exit + if tool_calls := data["messages"][-1].tool_calls: + return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + else: + return END + + def tools_node( + _: AgentState, config: RunnableConfig, *, tool_call: ToolCall + ) -> AgentState: + output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) + return { + "messages": ToolMessage( + content=output, name=tool_call["name"], tool_call_id=tool_call["id"] + ) + } + + # Define a new graph + workflow = StateGraph(AgentState) + + # Define the two nodes we will cycle between + workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) + workflow.add_node("tools", tools_node) + + # Set the entrypoint as `agent` + # This means that this node is the first one called + workflow.set_entry_point("agent") + + # We now add a conditional edge + workflow.add_conditional_edges("agent", should_continue) + + # We now add a normal edge from `tools` to `agent`. + # This means that after `tools` is called, `agent` node is called next. + workflow.add_edge("tools", "agent") + + # Finally, we compile it! + # This compiles it into a LangChain Runnable, + # meaning you can use it as you would any other runnable + app = workflow.compile() + + assert app.invoke({"messages": HumanMessage(content="what is weather in sf")}) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ), + ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + AIMessage(content="answer", id="ai3"), + ] + } + + assert [ + c + for c in app.stream( + {"messages": [HumanMessage(content="what is weather in sf")]} + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + }, + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + } + }, + { + "tools": [ + { + "messages": ToolMessage( + content="result for another", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call234", + ) + }, + { + "messages": ToolMessage( + content="result for a third one", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call567", + ), + }, + ] + }, + {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, + ] + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + next=("tools",), + config=(app_w_interrupt.checkpointer.get_tuple(config)).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + }, + ) + + # modify ai message + last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + app_w_interrupt.update_state(config, {"messages": last_message}) + + # message was replaced instead of appended + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "update", + "step": 2, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } + }, + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": [ + { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ] + }, + next=("tools", "tools"), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "loop", + "step": 4, + "writes": { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, + }, + }, + ) + + app_w_interrupt.update_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + next=(), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "source": "update", + "step": 5, + "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, + }, + ) + + def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 7515c0351..4913e1f05 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2508,7 +2508,6 @@ async def test_state_graph_few_shot() -> None: class BaseState(TypedDict): messages: Annotated[list[AnyMessage], add_messages] - # tool_results: Annotated[list[str], operator.add] class AgentState(BaseState): examples: Annotated[ @@ -3097,7 +3096,7 @@ async def test_state_graph_packets() -> None: model = FakeMessagesListChatModel( responses=[ AIMessage( - id="a1", + id="ai1", content="", tool_calls=[ { @@ -3108,7 +3107,7 @@ async def test_state_graph_packets() -> None: ], ), AIMessage( - id="a2", + id="ai2", content="", tool_calls=[ { @@ -3174,7 +3173,7 @@ async def test_state_graph_packets() -> None: "messages": [ HumanMessage(content="what is weather in sf", id=AnyStr()), AIMessage( - id="a1", + id="ai1", content="", tool_calls=[ { @@ -3191,7 +3190,7 @@ async def test_state_graph_packets() -> None: tool_call_id="tool_call123", ), AIMessage( - id="a2", + id="ai2", content="", tool_calls=[ { @@ -3231,7 +3230,7 @@ async def test_state_graph_packets() -> None: { "agent": { "messages": AIMessage( - id="a1", + id="ai1", content="", tool_calls=[ { @@ -3258,7 +3257,7 @@ async def test_state_graph_packets() -> None: { "agent": { "messages": AIMessage( - id="a2", + id="ai2", content="", tool_calls=[ { @@ -3307,34 +3306,46 @@ async def test_state_graph_packets() -> None: assert [ c async for c in app_w_interrupt.astream( - HumanMessage(content="what is weather in sf"), config + {"messages": HumanMessage(content="what is weather in sf")}, config ) ] == [ { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ) + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } }, ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - HumanMessage( - content="what is weather in sf", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ), - ], + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3344,40 +3355,49 @@ async def test_state_graph_packets() -> None: "source": "loop", "step": 1, "writes": { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ) + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } }, }, ) # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values[-1] - last_message.additional_kwargs["function_call"]["arguments"] = '"a different query"' - await app_w_interrupt.aupdate_state(config, last_message) + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) # message was replaced instead of appended assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - HumanMessage( - content="what is weather in sf", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - ], + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3387,69 +3407,101 @@ async def test_state_graph_packets() -> None: "source": "update", "step": 2, "writes": { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ) + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ) + } }, }, ) assert [c async for c in app_w_interrupt.astream(None, config)] == [ { - "tools": FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ) + "tools": [ + { + "messages": ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ) + } + ] }, { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ) + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, }, ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - HumanMessage( - content="what is weather in sf", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"another"'} - }, - id="ai2", - ), - ], - next=("tools",), + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + ] + }, + next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ "ts" @@ -3458,49 +3510,61 @@ async def test_state_graph_packets() -> None: "source": "loop", "step": 4, "writes": { - "agent": AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"another"', - } - }, - id="ai2", - ) + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ) + }, }, }, ) await app_w_interrupt.aupdate_state( config, - AIMessage(content="answer", id="ai2"), + {"messages": AIMessage(content="answer", id="ai2")}, ) # replaces message even if object identity is different, as long as id is the same assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - HumanMessage( - content="what is weather in sf", - id=AnyStr(), - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id=AnyStr(), - ), - AIMessage(content="answer", id="ai2"), - ], + values={ + "messages": [ + HumanMessage( + content="what is weather in sf", + id=AnyStr(), + ), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ToolMessage( + content="result for a different query", + name="search_api", + id=AnyStr(), + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3509,7 +3573,7 @@ async def test_state_graph_packets() -> None: metadata={ "source": "update", "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, }, ) From fb9fdcd3455e1bbf60ecab327af20c793b3a2f66 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 16:09:29 -0700 Subject: [PATCH 05/11] Lint --- langgraph/pregel/write.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index f9e9c1233..5dd2ceca4 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,11 +1,20 @@ from __future__ import annotations import asyncio -from typing import Any, Callable, List, NamedTuple, Optional, Sequence, Tuple, TypeVar +from typing import ( + Any, + Callable, + List, + NamedTuple, + Optional, + Sequence, + Tuple, + TypeVar, + Union, +) from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec -from sympy import Union from langgraph.constants import CONFIG_KEY_SEND, TASKS, Packet from langgraph.errors import InvalidUpdateError From b19c426a33f0305adf57febe243cef101eecc13a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 16:41:14 -0700 Subject: [PATCH 06/11] In state graph validate that all nodes either return None or write to one of the state keys --- langgraph/graph/state.py | 7 ++++++- langgraph/pregel/read.py | 1 + langgraph/pregel/write.py | 32 ++++++++++++++++++++++++++++---- tests/test_pregel.py | 15 +++++++++++++++ 4 files changed, 50 insertions(+), 5 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 01b16aa69..49bd34e37 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -321,7 +321,11 @@ class CompiledStateGraph(CompiledGraph): triggers=[START], channels=[START], writers=[ - ChannelWrite(state_write_entries, tags=[TAG_HIDDEN]), + ChannelWrite( + state_write_entries, + tags=[TAG_HIDDEN], + require_at_least_one_of=state_keys, + ), ], ) else: @@ -345,6 +349,7 @@ class CompiledStateGraph(CompiledGraph): ChannelWrite( [ChannelWriteEntry(key, key)] + state_write_entries, tags=[TAG_HIDDEN], + require_at_least_one_of=state_keys, ), ], ).pipe(node) diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 2e90a505b..a6041d5da 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -125,6 +125,7 @@ class PregelNode(RunnableBindingBase): writers[-2] = ChannelWrite( writes=writers[-2].writes + writers[-1].writes, tags=writers[-2].config["tags"] if writers[-2].config else None, + require_at_least_one_of=writers[-2].require_at_least_one_of, ) writers.pop() return writers diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 5dd2ceca4..82610c322 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -43,15 +43,21 @@ class ChannelWrite(RunnableCallable): - runnable to map input, or None to use the input, or any other value to use instead - whether to skip writing if the mapped value is None """ + require_at_least_one_of: Optional[Sequence[str]] + """ + If defined, at least one of these channels must be written to. + """ def __init__( self, writes: Sequence[Union[ChannelWriteEntry, Packet]], *, tags: Optional[list[str]] = None, + require_at_least_one_of: Optional[Sequence[str]] = None, ): super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags) self.writes = writes + self.require_at_least_one_of = require_at_least_one_of def __repr_args__(self) -> Any: return [("writes", self.writes)] @@ -100,7 +106,11 @@ class ChannelWrite(RunnableCallable): if not write.skip_none or val is not None ] # write packets and values - self.do_write(config, writes + values) + self.do_write( + config, + writes + values, + self.require_at_least_one_of if input is not None else None, + ) return input async def _awrite(self, input: Any, config: RunnableConfig) -> None: @@ -132,13 +142,27 @@ class ChannelWrite(RunnableCallable): if not write.skip_none or val is not None ] # write packets and values - self.do_write(config, writes + values) + self.do_write( + config, + writes + values, + self.require_at_least_one_of if input is not None else None, + ) return input @staticmethod - def do_write(config: RunnableConfig, values: List[Tuple[str, Any]]) -> None: + def do_write( + config: RunnableConfig, + values: List[Tuple[str, Any]], + require_at_least_one_of: Optional[Sequence[str]] = None, + ) -> None: + filtered = [(chan, val) for chan, val in values if val is not SKIP_WRITE] + if require_at_least_one_of is not None: + if not {chan for chan, _ in filtered} & set(require_at_least_one_of): + raise InvalidUpdateError( + f"Must write to at least one of {require_at_least_one_of}" + ) write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - write([(chan, val) for chan, val in values if val is not SKIP_WRITE]) + write(filtered) @staticmethod def is_writer(runnable: Runnable) -> bool: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 3d18fcd6c..a70533a9f 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -141,6 +141,21 @@ def test_graph_validation() -> None: with pytest.raises(ValueError): # extra is dead-end workflow.compile() + class State(TypedDict): + hello: str + + def node_a(state: State) -> State: + # typo + return {"hel": "world"} + + builder = StateGraph(State) + builder.add_node("a", node_a) + builder.set_entry_point("a") + builder.set_finish_point("a") + graph = builder.compile() + with pytest.raises(InvalidUpdateError): + assert graph.invoke({"hello": "there"}) == {"hello": "world"} + def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) From dc9c7254e71865ac00687f22d4adf07c654c8734 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 16:59:08 -0700 Subject: [PATCH 07/11] Passthrough additional keys from node to cond edge - this can be used to eg inform what gets sent in packets, without needing to write them to state first --- langgraph/graph/graph.py | 20 ++++++++++++++++++-- tests/test_pregel.py | 30 +++++++++++++++++++++++++----- 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 4e00d148a..b4fb5511e 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -69,7 +69,15 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: - result = self.path.invoke(reader(config) if reader else input, config) + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if isinstance(value, dict) and isinstance(input, dict): + value = {**input, **value} + else: + value = input + result = self.path.invoke(value, config) return self._finish(writer, input, result) async def _aroute( @@ -80,7 +88,15 @@ class Branch(NamedTuple): reader: Optional[Callable[[], Any]], writer: Callable[[list[str]], Optional[Runnable]], ) -> Runnable: - result = await self.path.ainvoke(reader(config) if reader else input, config) + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if isinstance(value, dict) and isinstance(input, dict): + value = {**input, **value} + else: + value = input + result = await self.path.ainvoke(value, config) return self._finish(writer, input, result) def _finish( diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a70533a9f..379386e40 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3500,8 +3500,17 @@ def test_state_graph_packets() -> None: ] ) + def agent(data: AgentState) -> AgentState: + return { + "messages": model.invoke(data["messages"]), + "something_extra": "hi there", + } + # Define decision-making logic def should_continue(data: AgentState) -> str: + assert ( + data["something_extra"] == "hi there" + ), "nodes can pass extra data to their cond edges, which isn't saved in state" # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] @@ -3522,7 +3531,7 @@ def test_state_graph_packets() -> None: workflow = StateGraph(AgentState) # Define the two nodes we will cycle between - workflow.add_node("agent", {"messages": RunnablePick("messages") | model}) + workflow.add_node("agent", agent) workflow.add_node("tools", tools_node) # Set the entrypoint as `agent` @@ -3745,7 +3754,9 @@ def test_state_graph_packets() -> None: # modify ai message last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] last_message.tool_calls[0]["args"]["query"] = "a different query" - app_w_interrupt.update_state(config, {"messages": last_message}) + app_w_interrupt.update_state( + config, {"messages": last_message, "something_extra": "hi there"} + ) # message was replaced instead of appended assert app_w_interrupt.get_state(config) == StateSnapshot( @@ -3786,7 +3797,8 @@ def test_state_graph_packets() -> None: "args": {"query": "a different query"}, }, ], - ) + ), + "something_extra": "hi there", } }, }, @@ -3900,7 +3912,10 @@ def test_state_graph_packets() -> None: app_w_interrupt.update_state( config, - {"messages": AIMessage(content="answer", id="ai2")}, + { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + }, ) # replaces message even if object identity is different, as long as id is the same @@ -3937,7 +3952,12 @@ def test_state_graph_packets() -> None: metadata={ "source": "update", "step": 5, - "writes": {"agent": {"messages": AIMessage(content="answer", id="ai2")}}, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + } + }, }, ) From 683fac83dc8632e8c4462e2d1d1f2e9e161db222 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 17:00:35 -0700 Subject: [PATCH 08/11] don't clear pending packets when just viewing --- langgraph/pregel/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index aaf6662ac..9fd2d3b75 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1632,7 +1632,8 @@ def _prepare_next_tasks( ) else: tasks.append(PregelTaskDescription(packet.node, val)) - checkpoint["pending_packets"].clear() + if for_execution: + checkpoint["pending_packets"].clear() # Check if any processes should be run in next step # If so, prepare the values to be passed to them for name, proc in processes.items(): From ae9c5639f07f7bbde1d761260dce48bc2c89dca7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 30 May 2024 17:01:21 -0700 Subject: [PATCH 09/11] Lint --- tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 379386e40..dfa3bfd43 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -146,7 +146,7 @@ def test_graph_validation() -> None: def node_a(state: State) -> State: # typo - return {"hel": "world"} + return {"hell": "world"} builder = StateGraph(State) builder.add_node("a", node_a) @@ -154,7 +154,7 @@ def test_graph_validation() -> None: builder.set_finish_point("a") graph = builder.compile() with pytest.raises(InvalidUpdateError): - assert graph.invoke({"hello": "there"}) == {"hello": "world"} + graph.invoke({"hello": "there"}) def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: From 9fd459c7fcc9c01b6502c4a9cf12049e30a09a3f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 31 May 2024 13:19:29 -0700 Subject: [PATCH 10/11] Update to use a single arg like other nodes --- langgraph/constants.py | 15 ++++----------- langgraph/pregel/__init__.py | 27 ++++++--------------------- langgraph/pregel/debug.py | 8 +++----- langgraph/pregel/io.py | 4 ++-- langgraph/pregel/read.py | 11 +++-------- langgraph/pregel/types.py | 1 - langgraph/serde/jsonplus.py | 8 +++----- tests/test_pregel.py | 6 ++---- tests/test_pregel_async.py | 12 ++++-------- 9 files changed, 27 insertions(+), 65 deletions(-) diff --git a/langgraph/constants.py b/langgraph/constants.py index d880aa733..f4c372e71 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,6 +1,4 @@ -from typing import Any - -from langgraph.errors import InvalidUpdateError +from typing import Any, NamedTuple CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" @@ -13,11 +11,6 @@ RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ} TAG_HIDDEN = "langsmith:hidden" -class Packet: - def __init__(self, /, __node__: str, **kwargs: Any) -> None: - if not kwargs: - raise InvalidUpdateError( - "Packet must have at least one keyword argument to pass to node" - ) - self.node = __node__ - self.kwargs = kwargs +class Packet(NamedTuple): + node: str + arg: Any diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 9fd2d3b75..933392b46 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -879,7 +879,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _, _ in next_tasks: + for _, _, _, writes, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1182,7 +1182,7 @@ class Pregel( # combine pending writes from all tasks pending_writes = deque[tuple[str, Any]]() - for _, _, _, writes, _, _, _ in next_tasks: + for _, _, _, writes, _, _ in next_tasks: pending_writes.extend(writes) if debug: @@ -1449,7 +1449,7 @@ def _should_interrupt( # and any triggered node is in interrupt_nodes list and any( node - for node, _, _, _, config, _, _ in tasks + for node, _, _, _, config, _ in tasks if ( (not config or TAG_HIDDEN not in config.get("tags")) if interrupt_nodes == "*" @@ -1577,27 +1577,13 @@ def _prepare_next_tasks( tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] # Consume pending packets for packet in checkpoint["pending_packets"]: - try: - val = next( - _proc_input( - step, - packet.node, - processes[packet.node], - managed, - channels, - catch=True, - ) - ) - except StopIteration: - logger.warn('No input for node "%s" in packet, skipping', packet.node) - continue if for_execution: - if node := processes[packet.node].get_node(packet.kwargs): + if node := processes[packet.node].get_node(): writes = deque() tasks.append( PregelExecutableTask( packet.node, - val, + packet.arg, node, writes, patch_config( @@ -1627,11 +1613,10 @@ def _prepare_next_tasks( }, ), [TASKS], - packet.kwargs, ) ) else: - tasks.append(PregelTaskDescription(packet.node, val)) + tasks.append(PregelTaskDescription(packet.node, packet.arg)) if for_execution: checkpoint["pending_packets"].clear() # Check if any processes should be run in next step diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index 6dc86643f..f9cf7e59d 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -64,7 +64,7 @@ def map_debug_tasks( step: int, tasks: list[PregelExecutableTask] ) -> Iterator[DebugOutputTask]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, input, _, _, config, triggers, kwargs) in enumerate(tasks): + for idx, (name, input, _, _, config, triggers) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -74,8 +74,6 @@ def map_debug_tasks( "input": input, "triggers": triggers, } - if kwargs is not None: - payload["kwargs"] = kwargs yield { "type": "task", "timestamp": ts, @@ -90,7 +88,7 @@ def map_debug_task_results( stream_channels_list: Sequence[str], ) -> Iterator[DebugOutputTaskResult]: ts = datetime.now(timezone.utc).isoformat() - for idx, (name, _, _, writes, config, _, _) in enumerate(tasks): + for idx, (name, _, _, writes, config, _) in enumerate(tasks): if config is not None and TAG_HIDDEN in config.get("tags", []): continue @@ -133,7 +131,7 @@ def print_step_tasks(step: int, next_tasks: list[PregelExecutableTask]) -> None: ) + "\n".join( f"- {get_colored_text(name, 'green')} -> {pformat(val)}" - for name, val, _, _, _, _, _ in next_tasks + for name, val, _, _, _, _ in next_tasks ) ) diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index e21a5f4f0..0b6979606 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -105,7 +105,7 @@ def map_output_updates( if isinstance(output_channels, str): if updated := [ (triggers == [TASKS], node, value) - for node, _, _, writes, _, triggers, _ in output_tasks + for node, _, _, writes, _, triggers in output_tasks for chan, value in writes if chan == output_channels ]: @@ -127,7 +127,7 @@ def map_output_updates( node, {chan: value for chan, value in writes if chan in output_channels}, ) - for node, _, _, writes, _, triggers, _ in output_tasks + for node, _, _, writes, _, triggers in output_tasks if any(chan in output_channels for chan, _ in writes) ]: grouped = defaultdict(list) diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index a6041d5da..902d96cf4 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -130,9 +130,7 @@ class PregelNode(RunnableBindingBase): writers.pop() return writers - def get_node( - self, kwargs: Optional[Dict[str, Any]] = None - ) -> Optional[Runnable[Any, Any]]: + def get_node(self) -> Optional[Runnable[Any, Any]]: writers = self.get_writers() if self.bound is DEFAULT_BOUND and not writers: return None @@ -141,12 +139,9 @@ class PregelNode(RunnableBindingBase): elif self.bound is DEFAULT_BOUND: return RunnableSequence(*writers) elif writers: - return RunnableSequence( - self.bound.bind(**kwargs) if kwargs is not None else self.bound, - *writers, - ) + return RunnableSequence(self.bound, *writers) else: - return self.bound.bind(**kwargs) if kwargs is not None else self.bound + return self.bound def __init__( self, diff --git a/langgraph/pregel/types.py b/langgraph/pregel/types.py index ea9cb69ef..a3826314a 100644 --- a/langgraph/pregel/types.py +++ b/langgraph/pregel/types.py @@ -18,7 +18,6 @@ class PregelExecutableTask(NamedTuple): writes: deque[tuple[str, Any]] config: Optional[RunnableConfig] triggers: list[str] - kwargs: Optional[dict[str, Any]] = None class StateSnapshot(NamedTuple): diff --git a/langgraph/serde/jsonplus.py b/langgraph/serde/jsonplus.py index 2b89f5af7..f4f49a7a4 100644 --- a/langgraph/serde/jsonplus.py +++ b/langgraph/serde/jsonplus.py @@ -3,7 +3,7 @@ import importlib import json from datetime import datetime, timedelta, timezone from enum import Enum -from typing import Any, Optional +from typing import Any, NamedTuple, Optional from uuid import UUID from langchain_core.load.load import Reviver @@ -64,10 +64,8 @@ class JsonPlusSerializer(SerializerProtocol): ) elif isinstance(obj, Enum): return self._encode_constructor_args(obj.__class__, args=[obj.value]) - elif isinstance(obj, Packet): - return self._encode_constructor_args( - Packet, args=[obj.node], kwargs=obj.kwargs - ) + elif isinstance(obj, NamedTuple): + return self._encode_constructor_args(Packet, args=[*obj]) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" diff --git a/tests/test_pregel.py b/tests/test_pregel.py index dfa3bfd43..bfabb59b2 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3513,13 +3513,11 @@ def test_state_graph_packets() -> None: ), "nodes can pass extra data to their cond edges, which isn't saved in state" # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + return [Packet("tools", tool_call) for tool_call in tool_calls] else: return END - def tools_node( - _: AgentState, config: RunnableConfig, *, tool_call: ToolCall - ) -> AgentState: + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) return { "messages": ToolMessage( diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4913e1f05..0f9868b65 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -2564,13 +2564,11 @@ Some examples of past conversations: def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + return [Packet("tools", tool_call) for tool_call in tool_calls] else: return "exit" - def tools_node( - _: AgentState, config: RunnableConfig, *, tool_call: ToolCall - ) -> AgentState: + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) return { "messages": ToolMessage( @@ -3130,13 +3128,11 @@ async def test_state_graph_packets() -> None: def should_continue(data: AgentState) -> str: # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: - return [Packet("tools", tool_call=tool_call) for tool_call in tool_calls] + return [Packet("tools", tool_call) for tool_call in tool_calls] else: return END - def tools_node( - _: AgentState, config: RunnableConfig, *, tool_call: ToolCall - ) -> AgentState: + def tools_node(tool_call: ToolCall, config: RunnableConfig) -> AgentState: output = tools_by_name[tool_call["name"]].invoke(tool_call["args"], config) return { "messages": ToolMessage( From 9c513ee7e6d28ea88f6d0704a999fa40d6d953f7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 31 May 2024 13:23:04 -0700 Subject: [PATCH 11/11] Lint --- langgraph/graph/graph.py | 4 +--- langgraph/pregel/debug.py | 13 ++++++------- langgraph/pregel/read.py | 2 +- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index b4fb5511e..8acc8f9b4 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -113,9 +113,7 @@ class Branch(NamedTuple): if any(dest is None or dest == START for dest in destinations): raise ValueError("Branch did not return a valid destination") if any(p.node == END for p in destinations if isinstance(p, Packet)): - raise InvalidUpdateError( - "Cannot send a packet with keyword arguments to the END node" - ) + raise InvalidUpdateError("Cannot send a packet to the END node") return writer(destinations) or input diff --git a/langgraph/pregel/debug.py b/langgraph/pregel/debug.py index f9cf7e59d..0fd00bdc8 100644 --- a/langgraph/pregel/debug.py +++ b/langgraph/pregel/debug.py @@ -68,17 +68,16 @@ def map_debug_tasks( if config is not None and TAG_HIDDEN in config.get("tags", []): continue - payload = { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), - "name": name, - "input": input, - "triggers": triggers, - } yield { "type": "task", "timestamp": ts, "step": step, - "payload": payload, + "payload": { + "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, idx)))), + "name": name, + "input": input, + "triggers": triggers, + }, } diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 902d96cf4..a14c704e4 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union +from typing import Any, Callable, Mapping, Optional, Sequence, Union from langchain_core.pydantic_v1 import Field from langchain_core.runnables import (