From db3219f77c45e51f0e934c1d816262da7726ae36 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 3 Apr 2024 13:13:04 -0700 Subject: [PATCH] Ensure node tracing output is exactly what was returned from node --- langgraph/graph/graph.py | 15 +++++++++--- langgraph/graph/message.py | 3 ++- langgraph/graph/state.py | 37 +++++++++++++++-------------- langgraph/pregel/read.py | 48 +++++++++++++++----------------------- tests/test_pregel.py | 8 +++---- 5 files changed, 56 insertions(+), 55 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 9fdc3d8af..1a4a9dfc7 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -40,12 +40,17 @@ class Branch(NamedTuple): condition: Runnable[Any, str] ends: Optional[dict[str, str]] - def run(self, writer: Callable[[str], Optional[Runnable]]) -> None: + def run( + self, + writer: Callable[[str], Optional[Runnable]], + reader: Optional[Callable[[RunnableConfig], Any]] = None, + ) -> None: return ChannelWrite.register_writer( RunnableCallable( func=self._route, afunc=self._aroute, writer=writer, + reader=reader, name=None, trace=False, ) @@ -56,9 +61,10 @@ class Branch(NamedTuple): input: Any, config: RunnableConfig, *, + reader: Optional[Callable[[], Any]], writer: Callable[[str], Optional[Runnable]], ) -> Runnable: - result = self.condition.invoke(input, config) + result = self.condition.invoke(reader(config) if reader else input, config) if self.ends: destination = self.ends[result] else: @@ -70,9 +76,12 @@ class Branch(NamedTuple): input: Any, config: RunnableConfig, *, + reader: Optional[Callable[[], Any]], writer: Callable[[str], Optional[Runnable]], ) -> Runnable: - result = await self.condition.ainvoke(input, config) + result = await self.condition.ainvoke( + reader(config) if reader else input, config + ) if self.ends: destination = self.ends[result] else: diff --git a/langgraph/graph/message.py b/langgraph/graph/message.py index 42b7b2a49..fa2b29ff1 100644 --- a/langgraph/graph/message.py +++ b/langgraph/graph/message.py @@ -3,13 +3,14 @@ from typing import Annotated, Union from langchain_core.messages import ( AnyMessage, + MessageLikeRepresentation, convert_to_messages, message_chunk_to_message, ) from langgraph.graph.state import StateGraph -Messages = Union[list[AnyMessage], AnyMessage] +Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation] def add_messages(left: Messages, right: Messages) -> Messages: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 9815816ea..655730b37 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -142,18 +142,6 @@ class CompiledStateGraph(CompiledGraph): ) for key in state_keys ] - # node that reads current state with (this node's) updates applied - state_reader = ChannelRead( - state_keys[0] if state_keys == ["__root__"] else state_keys, - tags=[TAG_HIDDEN], - fresh=True, - # coerce state dict to schema class (eg. pydantic model) - mapper=( - None - if state_keys == ["__root__"] - else partial(_coerce_state, self.graph.schema) - ), - ) # add node and output channel if key == START: @@ -163,8 +151,6 @@ class CompiledStateGraph(CompiledGraph): channels=[START], writers=[ ChannelWrite(state_write_entries, tags=[TAG_HIDDEN]), - # read back state with updates applied - state_reader, ], ) else: @@ -178,15 +164,17 @@ class CompiledStateGraph(CompiledGraph): else {chan: chan for chan in state_keys} ), # coerce state dict to schema class (eg. pydantic model) - mapper=state_reader.mapper, + mapper=( + None + if state_keys == ["__root__"] + else partial(_coerce_state, self.graph.schema) + ), writers=[ # publish to this channel and state keys ChannelWrite( [ChannelWriteEntry(key)] + state_write_entries, tags=[TAG_HIDDEN], ), - # read back state with updates applied - state_reader, ], ).pipe(node) @@ -226,7 +214,7 @@ class CompiledStateGraph(CompiledGraph): ) # attach branch publisher - self.nodes[start] |= branch.run(branch_writer) + self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph)) # attach branch subscribers ends = branch.ends.values() if branch.ends else [node for node in self.nodes] @@ -237,6 +225,19 @@ class CompiledStateGraph(CompiledGraph): self.nodes[end].triggers.append(channel_name) +def _get_state_reader(graph: StateGraph) -> ChannelRead: + state_keys = list(graph.channels) + return partial( + ChannelRead.do_read, + channel=state_keys[0] if state_keys == ["__root__"] else state_keys, + fresh=True, + # coerce state dict to schema class (eg. pydantic model) + mapper=( + None if state_keys == ["__root__"] else partial(_coerce_state, graph.schema) + ), + ) + + def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 6e6049e24..6398a1e2c 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -65,30 +65,34 @@ class ChannelRead(RunnableCallable): return super().get_name(suffix, name=name) def _read(self, _: Any, config: RunnableConfig) -> Any: - try: - read: READ_TYPE = config["configurable"][CONFIG_KEY_READ] - except KeyError: - raise RuntimeError( - f"Runnable {self} is not configured with a read function" - "Make sure to call in the context of a Pregel process" - ) - if self.mapper: - return self.mapper(read(self.channel, self.fresh)) - else: - return read(self.channel, self.fresh) + return self.do_read( + config, channel=self.channel, fresh=self.fresh, mapper=self.mapper + ) async def _aread(self, _: Any, config: RunnableConfig) -> Any: + return self.do_read( + config, channel=self.channel, fresh=self.fresh, mapper=self.mapper + ) + + @staticmethod + def do_read( + config: RunnableConfig, + *, + channel: Union[str, list[str]], + fresh: bool = False, + mapper: Optional[Callable[[Any], Any]] = None, + ) -> Any: try: read: READ_TYPE = config["configurable"][CONFIG_KEY_READ] except KeyError: raise RuntimeError( - f"Runnable {self} is not configured with a read function" + "Not configured with a read function" "Make sure to call in the context of a Pregel process" ) - if self.mapper: - return self.mapper(read(self.channel, self.fresh)) + if mapper: + return mapper(read(channel, fresh)) else: - return read(self.channel, self.fresh) + return read(channel, fresh) DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() @@ -110,20 +114,6 @@ class PregelNode(RunnableBindingBase): def get_writers(self) -> list[Runnable]: """Get writers with optimizations applied.""" writers = self.writers.copy() - while writers and isinstance(writers[-1], ChannelRead): - # we can avoid reads if no writers would be called after them - writers.pop() - while ( - len(writers) > 1 - and isinstance(writers[-1], ChannelWrite) - and all( - write.value is not None and not isinstance(write.value, Runnable) - for write in writers[-1].writes - ) - and isinstance(writers[-2], ChannelRead) - ): - # we can avoid reads if all subsequent write values don't use the input - writers.pop(-2) while ( len(writers) > 1 and isinstance(writers[-1], ChannelWrite) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 903d51dc2..ac80fb981 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2307,7 +2307,7 @@ def test_message_graph( FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000013", + id="00000000-0000-4000-8000-000000000012", ), AIMessage( content="", @@ -2319,7 +2319,7 @@ def test_message_graph( FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000024", + id="00000000-0000-4000-8000-000000000022", ), AIMessage(content="answer", id="ai3"), ] @@ -2338,7 +2338,7 @@ def test_message_graph( "action": FunctionMessage( content="result for query", name="search_api", - id="00000000-0000-4000-8000-000000000043", + id="00000000-0000-4000-8000-000000000039", ) }, { @@ -2354,7 +2354,7 @@ def test_message_graph( "action": FunctionMessage( content="result for another", name="search_api", - id="00000000-0000-4000-8000-000000000054", + id="00000000-0000-4000-8000-000000000049", ) }, {"agent": AIMessage(content="answer", id="ai3")},