diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index fc16e5a5a..d5ac32346 100644 --- a/langgraph/graph/state.py +++ b/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 diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 715d20299..5d5f90633 100644 --- a/langgraph/pregel/__init__.py +++ b/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() ] ) diff --git a/langgraph/pregel/write.py b/langgraph/pregel/write.py index 96b3b4b2c..4b2ad9ea5 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, 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 diff --git a/langgraph/utils.py b/langgraph/utils.py index dc7ec7bfd..0604fd81e 100644 --- a/langgraph/utils.py +++ b/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 diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 48cf24ecd..4d4a0df6e 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/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"]}}}' # --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f7d42269b..d184b3754 100644 --- a/tests/test_pregel.py +++ b/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: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3140cbe1a..1196de6d6 100644 --- a/tests/test_pregel_async.py +++ b/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)}, ):