Merge pull request #345 from langchain-ai/nc/24apr/fix-write-runnable-bug

Fix bug where saving a runnable or function in state object would misbehave

langgraph/graph/state.py

@@ -174,10 +174,13 @@ class CompiledStateGraph(CompiledGraph):
         # state updaters
         state_write_entries = [
             (
-                ChannelWriteEntry(key, None, skip_none=True)
+                ChannelWriteEntry(key, skip_none=True)
                 if key == "__root__"
                 else ChannelWriteEntry(
-                    key, RunnableCallable(_get_state_key, key=key, trace=False)
+                    key,
+                    mapper=RunnableCallable(
+                        _get_state_key, key=key, trace=False, recurse=False
+                    ),
                 )
             )
             for key in state_keys

langgraph/pregel/__init__.py

@@ -167,7 +167,9 @@ class Channel:
         return ChannelWrite(
             [ChannelWriteEntry(c) for c in channels]
             + [
-                ChannelWriteEntry(k, _coerce_write_value(v), True)
+                ChannelWriteEntry(k, skip_none=True, mapper=coerce_to_runnable(v))
+                if isinstance(v, Runnable) or callable(v)
+                else ChannelWriteEntry(k, value=v)
                 for k, v in kwargs.items()
             ]
         )

langgraph/pregel/write.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar, Union
+from typing import Any, Callable, NamedTuple, Optional, Sequence, TypeVar
 
 from langchain_core.runnables import Runnable, RunnableConfig
 from langchain_core.runnables.utils import ConfigurableFieldSpec
@@ -14,12 +14,14 @@ R = TypeVar("R", bound=Runnable)
 
 
 SKIP_WRITE = object()
+PASSTHROUGH = object()
 
 
 class ChannelWriteEntry(NamedTuple):
     channel: str
-    value: Optional[Union[Any, Runnable]] = None
+    value: Any = PASSTHROUGH
     skip_none: bool = False
+    mapper: Optional[Runnable] = None
 
 
 class ChannelWrite(RunnableCallable):
@@ -44,7 +46,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(chan for chan, _, _, _ in self.writes)}>"
         return super().get_name(suffix, name=name)
 
     @property
@@ -61,39 +63,38 @@ class ChannelWrite(RunnableCallable):
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
         values = [
-            (
-                chan,
-                r.invoke(input, config)
-                if isinstance(r, Runnable)
-                else r
-                if r is not None
-                else input,
-            )
-            for chan, r, _ in self.writes
+            input if write.value is PASSTHROUGH else write.value
+            for write in self.writes
         ]
         values = [
-            write
-            for write, (_, _, skip_none) in zip(values, self.writes)
-            if not skip_none or write[1] is not None
+            val if write.mapper is None else write.mapper.invoke(val, config)
+            for val, write in zip(values, self.writes)
+        ]
+        values = [
+            (write.channel, val)
+            for val, write in zip(values, self.writes)
+            if not write.skip_none or val is not None
         ]
         self.do_write(config, **dict(values))
         return input
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
+        values = [
+            input if write.value is PASSTHROUGH else write.value
+            for write in self.writes
+        ]
         values = await asyncio.gather(
             *(
-                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.writes
+                _mk_future(val)
+                if write.mapper is None
+                else write.mapper.ainvoke(val, config)
+                for val, write in zip(values, self.writes)
             )
         )
         values = [
-            (chan, val)
-            for val, (chan, _, skip_none) in zip(values, self.writes)
-            if not skip_none or val is not None
+            (write.channel, val)
+            for val, write in zip(values, self.writes)
+            if not write.skip_none or val is not None
         ]
         self.do_write(config, **dict(values))
         return input

langgraph/utils.py

@@ -23,6 +23,7 @@ class RunnableCallable(Runnable):
         name: Optional[str] = None,
         tags: Optional[list[str]] = None,
         trace: bool = True,
+        recurse: bool = True,
         **kwargs: Any,
     ) -> None:
         self.name = name or func.__name__
@@ -31,6 +32,7 @@ class RunnableCallable(Runnable):
         self.config = {"tags": tags} if tags else None
         self.kwargs = kwargs
         self.trace = trace
+        self.recurse = recurse
 
     def __repr__(self) -> str:
         repr_args = {
@@ -47,7 +49,7 @@ class RunnableCallable(Runnable):
             )
         else:
             ret = self.func(input, merge_configs(self.config, config), **self.kwargs)
-        if isinstance(ret, Runnable):
+        if isinstance(ret, Runnable) and self.recurse:
             return ret.invoke(input, config)
         return ret
 
@@ -62,6 +64,6 @@ class RunnableCallable(Runnable):
             ret = await self.afunc(
                 input, merge_configs(self.config, config), **self.kwargs
             )
-        if isinstance(ret, Runnable):
+        if isinstance(ret, Runnable) and self.recurse:
             return await ret.ainvoke(input, config)
         return ret

tests/__snapshots__/test_pregel.ambr

@@ -1884,6 +1884,23 @@
    +---------+   
   '''
 # ---
+# name: test_nested_graph.1
+  '''
+  %%{init: {'flowchart': {'curve': 'linear'}}}%%
+  graph TD;
+  	__start__[__start__]:::startclass;
+  	__end__[__end__]:::endclass;
+  	inner([inner]):::otherclass;
+  	side([side]):::otherclass;
+  	__start__ --> inner;
+  	inner --> side;
+  	side --> __end__;
+  	classDef startclass fill:#ffdfba;
+  	classDef endclass fill:#baffc9;
+  	classDef otherclass fill:#fad7de;
+  
+  '''
+# ---
 # name: test_prebuilt_chat
   '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}'
 # ---

tests/test_pregel.py

@@ -3849,17 +3849,27 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
 
 
 def test_nested_graph(snapshot: SnapshotAssertion) -> None:
-    class State(TypedDict):
+    def never_called_fn(state: Any):
+        assert 0, "This function should never be called"
+
+    never_called = RunnableLambda(never_called_fn)
+
+    class InnerState(TypedDict):
         my_key: str
+        my_other_key: str
 
-    def up(state: State):
-        return {"my_key": state["my_key"] + " there"}
+    def up(state: InnerState):
+        return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]}
 
-    inner = StateGraph(State)
+    inner = StateGraph(InnerState)
     inner.add_node("up", up)
     inner.set_entry_point("up")
     inner.set_finish_point("up")
 
+    class State(TypedDict):
+        my_key: str
+        never_called: Any
+
     def side(state: State):
         return {"my_key": state["my_key"] + " and back again"}
 
@@ -3873,9 +3883,35 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
     app = graph.compile()
 
     assert app.get_graph().draw_ascii() == snapshot
-    assert app.invoke({"my_key": "my value"}) == {
-        "my_key": "my value there and back again"
+    assert app.get_graph(xray=True).draw_mermaid() == snapshot
+    assert app.invoke(
+        {"my_key": "my value", "never_called": never_called}, debug=True
+    ) == {
+        "my_key": "my value there and back again",
+        "never_called": never_called,
     }
+    assert [*app.stream({"my_key": "my value", "never_called": never_called})] == [
+        {"inner": {"my_key": "my value there"}},
+        {"side": {"my_key": "my value there and back again"}},
+    ]
+    assert [
+        *app.stream(
+            {"my_key": "my value", "never_called": never_called}, stream_mode="values"
+        )
+    ] == [
+        {
+            "my_key": "my value",
+            "never_called": never_called,
+        },
+        {
+            "my_key": "my value there",
+            "never_called": never_called,
+        },
+        {
+            "my_key": "my value there and back again",
+            "never_called": never_called,
+        },
+    ]
 
 
 def test_repeat_condition(snapshot: SnapshotAssertion) -> None:

tests/test_pregel_async.py

@@ -3475,17 +3475,27 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N
 
 
 async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
-    class State(TypedDict):
+    def never_called_fn(state: Any):
+        assert 0, "This function should never be called"
+
+    never_called = RunnableLambda(never_called_fn)
+
+    class InnerState(TypedDict):
         my_key: str
+        my_other_key: str
 
-    async def up(state: State):
-        return {"my_key": state["my_key"] + " there"}
+    def up(state: InnerState):
+        return {"my_key": state["my_key"] + " there", "my_other_key": state["my_key"]}
 
-    inner = StateGraph(State)
+    inner = StateGraph(InnerState)
     inner.add_node("up", up)
     inner.set_entry_point("up")
     inner.set_finish_point("up")
 
+    class State(TypedDict):
+        my_key: str
+        never_called: Any
+
     async def side(state: State):
         return {"my_key": state["my_key"] + " and back again"}
 
@@ -3499,24 +3509,32 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
     app = graph.compile()
 
     assert app.get_graph().draw_ascii() == snapshot
-    assert await app.ainvoke({"my_key": "my value"}) == {
-        "my_key": "my value there and back again"
+    assert await app.ainvoke({"my_key": "my value", "never_called": never_called}) == {
+        "my_key": "my value there and back again",
+        "never_called": never_called,
     }
-    assert [chunk async for chunk in app.astream({"my_key": "my value"})] == [
+    assert [
+        chunk
+        async for chunk in app.astream(
+            {"my_key": "my value", "never_called": never_called}
+        )
+    ] == [
         {"inner": {"my_key": "my value there"}},
         {"side": {"my_key": "my value there and back again"}},
     ]
     assert [
         chunk
-        async for chunk in app.astream({"my_key": "my value"}, stream_mode="values")
+        async for chunk in app.astream(
+            {"my_key": "my value", "never_called": never_called}, stream_mode="values"
+        )
     ] == [
-        {"my_key": "my value"},
-        {"my_key": "my value there"},
-        {"my_key": "my value there and back again"},
+        {"my_key": "my value", "never_called": never_called},
+        {"my_key": "my value there", "never_called": never_called},
+        {"my_key": "my value there and back again", "never_called": never_called},
     ]
     times_called = 0
     async for event in app.astream_events(
-        {"my_key": "my value"},
+        {"my_key": "my value", "never_called": never_called},
         version="v1",
         config={"run_id": UUID(int=0)},
         stream_mode="values",
@@ -3524,12 +3542,15 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
         if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)):
             times_called += 1
             assert event["data"] == {
-                "output": {"my_key": "my value there and back again"}
+                "output": {
+                    "my_key": "my value there and back again",
+                    "never_called": never_called,
+                }
             }
     assert times_called == 1
     times_called = 0
     async for event in app.astream_events(
-        {"my_key": "my value"},
+        {"my_key": "my value", "never_called": never_called},
         version="v1",
         config={"run_id": UUID(int=0)},
     ):