diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 20fccde54..88c526e2f 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, Type -from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough +from langchain_core.runnables import RunnableLambda, RunnablePassthrough from langchain_core.runnables.base import RunnableLike from langgraph.channels.base import BaseChannel, InvalidUpdateError @@ -13,7 +13,7 @@ from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, Graph from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import ChannelRead -from langgraph.pregel.write import ChannelWrite +from langgraph.pregel.write import SKIP_WRITE, ChannelWrite START = "__start__" @@ -39,10 +39,13 @@ class StateGraph(Graph): state_keys = list(self.channels) state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys - update_state = ( - partial(_update_state_dict, state_keys) - if isinstance(state_keys_read, list) - else _update_state_root + update_channels = ( + [("__root__", None, True)] + if not isinstance(state_keys_read, list) + else [ + (key, RunnableLambda(partial(_dict_getter, state_keys, key)), False) + for key in state_keys_read + ] ) coerce_state = ( partial(_coerce_state, self.schema) @@ -59,8 +62,7 @@ class StateGraph(Graph): Channel.subscribe_to(f"{key}:inbox") | coerce_state # coerce/validate using schema | node - | partial(update_state, key) - | Channel.write_to(key) + | ChannelWrite(channels=[(key, None, False)] + update_channels) ) for key, node in self.nodes.items() } @@ -80,11 +82,9 @@ class StateGraph(Graph): branch.runnable, name=f"{key}_condition" ) - nodes[START] = ( - Channel.subscribe_to(f"{START}:inbox", tags=["langsmith:hidden"]) - | partial(update_state, START) - | Channel.write_to(START) - ) + 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) @@ -105,25 +105,16 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) -def _update_state_dict( - state_keys: list[str], node_name: str, input: dict[str, Any], config: RunnableConfig -) -> dict[str, Any]: +def _dict_getter(allowed_keys: str, key: str, input: dict) -> Any: if input is not None: - if not isinstance(input, dict) or any(key not in state_keys for key in input): + if not isinstance(input, dict) or any(key not in allowed_keys for key in input): raise InvalidUpdateError( - f"Invalid state update from node {node_name}," - f" expected dict with one or more of {state_keys}, got {input}" + f"Invalid state update," + f" expected dict with one or more of {allowed_keys}, got {input}" ) - ChannelWrite.do_write(config, **input) - return input - - -def _update_state_root( - node_name: str, input: Any, config: RunnableConfig -) -> dict[str, Any]: - if input is not None: - ChannelWrite.do_write(config, __root__=input) - return input + return input.get(key, SKIP_WRITE) + else: + return SKIP_WRITE def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]: diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5ab1e9223..7c630579e 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -149,8 +149,8 @@ class Channel: """Writes to channels the result of the lambda, or None to skip writing.""" return ChannelWrite( channels=( - [(c, None) for c in channels] - + [(k, _coerce_write_value(v)) for k, v in kwargs.items()] + [(c, None, False) for c in channels] + + [(k, _coerce_write_value(v), True) for k, v in kwargs.items()] ) ) diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index a43d8d876..3a050925f 100644 --- a/langgraph/pregel/write.py +++ b/langgraph/pregel/write.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from typing import Any, Callable, Optional, Sequence from langchain_core.runnables import ( @@ -14,8 +15,11 @@ from langgraph.constants import CONFIG_KEY_SEND TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] +SKIP_WRITE = object() + + class ChannelWrite(RunnablePassthrough): - channels: Sequence[tuple[str, Optional[Runnable]]] + channels: Sequence[tuple[str, Optional[Runnable], bool]] """ Mapping of write channels to Runnables that return the value to be written, or None to skip writing. @@ -27,10 +31,10 @@ class ChannelWrite(RunnablePassthrough): def __init__( self, *, - channels: Sequence[tuple[str, Optional[Runnable]]], + channels: Sequence[tuple[str, Optional[Runnable], bool]], ): super().__init__(func=self._write, afunc=self._awrite, channels=channels) - self.name = f"ChannelWrite<{','.join(chan for chan, _ in self.channels)}>" + self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>" def __repr_args__(self) -> Any: return [("channels", self.channels)] @@ -49,25 +53,28 @@ class ChannelWrite(RunnablePassthrough): def _write(self, input: Any, config: RunnableConfig) -> None: values = [ - (chan, r.invoke(input, config) if r else input) for chan, r in self.channels + (chan, r.invoke(input, config) if r else input) + for chan, r, _ in self.channels ] values = [ write - for write, chan in zip(values, self.channels) - if chan[1] is None or write[1] is not None + for write, (_, _, skip_none) in zip(values, self.channels) + if not skip_none or write[1] is not None ] self.do_write(config, **dict(values)) async def _awrite(self, input: Any, config: RunnableConfig) -> None: + values = await asyncio.gather( + *( + r.ainvoke(input, config) if r else _mk_future(input) + for _, r, _ in self.channels + ) + ) values = [ - (chan, await r.ainvoke(input, config) if r else input) - for chan, r in self.channels - ] - values = [ - write - for write, chan in zip(values, self.channels) - if chan[1] is None or write[1] is not None + (chan, val) + for val, (chan, _, skip_none) in zip(values, self.channels) + if not skip_none or val is not None ] self.do_write(config, **dict(values)) @@ -75,4 +82,10 @@ class ChannelWrite(RunnablePassthrough): @staticmethod def do_write(config: RunnableConfig, **values: Any) -> None: write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] - write([(chan, val) for chan, val in values.items()]) + write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE]) + + +def _mk_future(val: Any) -> asyncio.Future: + fut = asyncio.Future() + fut.set_result(val) + return fut