From d90478b50c547286d5310ca9e7a72fef7e5c294f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 16:53:09 -0800 Subject: [PATCH] Remove unnecessary runs from StateGraph/MessageGraph run tree --- langgraph/graph/state.py | 53 ++++++++++++++++++++++++------------ langgraph/pregel/__init__.py | 16 ++++++++--- langgraph/pregel/read.py | 7 +++++ langgraph/pregel/write.py | 37 +++++++++++++++++-------- 4 files changed, 81 insertions(+), 32 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 75d5f3d1f..c448b7628 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -3,7 +3,7 @@ from functools import partial from inspect import signature from typing import Any, Optional, Sequence, Type -from langchain_core.runnables import RunnableLambda, RunnablePassthrough +from langchain_core.runnables import RunnableLambda from langchain_core.runnables.base import RunnableLike from langgraph.channels.any_value import AnyValue @@ -14,8 +14,8 @@ from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, START, CompiledGraph, Graph from langgraph.pregel import Channel -from langgraph.pregel.read import ChannelRead -from langgraph.pregel.write import SKIP_WRITE, ChannelWrite +from langgraph.pregel.read import ChannelInvoke +from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry class StateGraph(Graph): @@ -46,18 +46,25 @@ class StateGraph(Graph): state_keys = list(self.channels) state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys + state_channels = ( + {chan: chan for chan in state_keys} + if isinstance(state_keys_read, list) + else {None: state_keys_read} + ) update_channels = ( - [("__root__", None, True)] + [ChannelWriteEntry("__root__", None, True)] if not isinstance(state_keys_read, list) else [ - (key, RunnableLambda(partial(_dict_getter, state_keys, key)), False) + ChannelWriteEntry( + key, RunnableLambda(partial(_dict_getter, state_keys, key)), False + ) for key in state_keys_read ] ) coerce_state = ( partial(_coerce_state, self.schema) if isinstance(state_keys_read, list) - else RunnablePassthrough() + else None ) outgoing_edges = defaultdict(list) @@ -66,10 +73,15 @@ class StateGraph(Graph): nodes = { key: ( - Channel.subscribe_to(f"{key}:inbox") - | coerce_state # coerce/validate using schema + ChannelInvoke( + triggers=[f"{key}:inbox"], + channels=state_channels, + mapper=coerce_state, + ) | node - | ChannelWrite(channels=[(key, None, False)] + update_channels) + | ChannelWrite( + channels=[ChannelWriteEntry(key, None, False)] + update_channels + ) ) for key, node in self.nodes.items() } @@ -89,11 +101,16 @@ class StateGraph(Graph): outgoing = outgoing_edges[key] edges_key = f"{key}:edges" if outgoing or key in self.branches: - nodes[edges_key] = Channel.subscribe_to( - key, tags=["langsmith:hidden"] - ) | ChannelRead(state_keys_read) + nodes[edges_key] = ChannelInvoke( + triggers=[key], tags=["langsmith:hidden"], channels=state_channels + ) if outgoing: - nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing]) + nodes[edges_key] |= ChannelWrite( + channels=[ + ChannelWriteEntry(dest, None if dest == END else key, True) + for dest in outgoing + ] + ) if key in self.branches: for branch in self.branches[key]: nodes[edges_key] |= RunnableLambda( @@ -102,10 +119,12 @@ class StateGraph(Graph): nodes[START] = Channel.subscribe_to( f"{START}:inbox", tags=["langsmith:hidden"] - ) | ChannelWrite(channels=[(START, None, False)] + update_channels) - nodes[f"{START}:edges"] = Channel.subscribe_to( - START, tags=["langsmith:hidden"] - ) | ChannelRead(state_keys_read) + ) | ChannelWrite( + channels=[ChannelWriteEntry(START, None, False)] + update_channels + ) + nodes[f"{START}:edges"] = ChannelInvoke( + triggers=[START], tags=["langsmith:hidden"], channels=state_channels + ) if self.entry_point: nodes[f"{START}:edges"] |= Channel.write_to(f"{self.entry_point}:inbox") elif self.entry_point_branch: diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d7fadd9e5..d512bd3dd 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -63,7 +63,7 @@ from langgraph.pregel.log import logger from langgraph.pregel.read import ChannelBatch, ChannelInvoke from langgraph.pregel.reserved import ReservedChannels from langgraph.pregel.validate import validate_graph, validate_keys -from langgraph.pregel.write import ChannelWrite +from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry WriteValue = Union[ Runnable[Input, Output], @@ -149,8 +149,11 @@ class Channel: """Writes to channels the result of the lambda, or None to skip writing.""" return ChannelWrite( channels=( - [(c, None, False) for c in channels] - + [(k, _coerce_write_value(v), True) for k, v in kwargs.items()] + [ChannelWriteEntry(c, None, False) for c in channels] + + [ + ChannelWriteEntry(k, _coerce_write_value(v), True) + for k, v in kwargs.items() + ] ) ) @@ -789,7 +792,8 @@ def _prepare_next_tasks( checkpoint["channel_versions"][chan] > seen[chan] for chan in proc.triggers ): - # If all channels subscribed by this process are not empty + # If all trigger channels subscribed by this process are not empty + # then invoke the process with the values of all non-empty channels try: val: Any = { k: _read_channel( @@ -800,6 +804,10 @@ def _prepare_next_tasks( except EmptyChannelError: continue + # If the process has a mapper, apply it to the value + if proc.mapper is not None: + val = proc.mapper(val) + # Processes that subscribe to a single keyless channel get # the value directly, instead of a dict if list(proc.channels.keys()) == [None]: diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 08903082e..9006b639a 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -79,6 +79,8 @@ class ChannelInvoke(RunnableBindingBase): triggers: list[str] = Field(default_factory=list) + mapper: Optional[Callable[[Any], Any]] = None + when: Optional[Callable[[Any], bool]] = None bound: Runnable[Any, Any] = Field(default=default_bound) @@ -89,6 +91,7 @@ class ChannelInvoke(RunnableBindingBase): self, channels: Mapping[None, str] | Mapping[str, str], triggers: Sequence[str], + mapper: Optional[Callable[[Any], Any]] = None, when: Optional[Callable[[Any], bool]] = None, tags: Optional[list[str]] = None, *, @@ -100,6 +103,7 @@ class ChannelInvoke(RunnableBindingBase): super().__init__( channels=channels, triggers=triggers, + mapper=mapper, when=when, bound=bound or default_bound, kwargs=kwargs or {}, @@ -120,6 +124,7 @@ class ChannelInvoke(RunnableBindingBase): **{chan: chan for chan in channels}, }, triggers=self.triggers, + mapper=self.mapper, when=self.when, bound=self.bound, kwargs=self.kwargs, @@ -138,6 +143,7 @@ class ChannelInvoke(RunnableBindingBase): return ChannelInvoke( channels=self.channels, triggers=self.triggers, + mapper=self.mapper, when=self.when, bound=coerce_to_runnable(other), kwargs=self.kwargs, @@ -147,6 +153,7 @@ class ChannelInvoke(RunnableBindingBase): return ChannelInvoke( channels=self.channels, triggers=self.triggers, + mapper=self.mapper, when=self.when, # delegate to __or__ in self.bound bound=self.bound | other, diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 3a050925f..40c7bb33f 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,7 +1,7 @@ from __future__ import annotations import asyncio -from typing import Any, Callable, Optional, Sequence +from typing import Any, Callable, NamedTuple, Optional, Sequence, Union from langchain_core.runnables import ( Runnable, @@ -18,21 +18,25 @@ TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] SKIP_WRITE = object() +class ChannelWriteEntry(NamedTuple): + channel: str + value: Optional[Union[Any, Runnable]] + skip_none: bool + + class ChannelWrite(RunnablePassthrough): - channels: Sequence[tuple[str, Optional[Runnable], bool]] + channels: Sequence[ChannelWriteEntry] """ - Mapping of write channels to Runnables that return the value to be written, - or None to skip writing. + Sequence of write entries, each of which is a tuple of: + - channel name + - 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 """ class Config: arbitrary_types_allowed = True - def __init__( - self, - *, - channels: Sequence[tuple[str, Optional[Runnable], bool]], - ): + def __init__(self, *, channels: Sequence[ChannelWriteEntry]): super().__init__(func=self._write, afunc=self._awrite, channels=channels) self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>" @@ -53,7 +57,14 @@ class ChannelWrite(RunnablePassthrough): def _write(self, input: Any, config: RunnableConfig) -> None: values = [ - (chan, r.invoke(input, config) if r else input) + ( + chan, + r.invoke(input, config) + if isinstance(r, Runnable) + else r + if r is not None + else input, + ) for chan, r, _ in self.channels ] values = [ @@ -67,7 +78,11 @@ class ChannelWrite(RunnablePassthrough): async def _awrite(self, input: Any, config: RunnableConfig) -> None: values = await asyncio.gather( *( - r.ainvoke(input, config) if r else _mk_future(input) + r.ainvoke(input, config) + if isinstance(r, Runnable) + else _mk_future(r) + if r is not None + else _mk_future(input) for _, r, _ in self.channels ) )