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(