Merge pull request #41 from langchain-ai/nc/fan-out

StateGraph/MessageGraph: Add support for multiple incoming edges

langgraph/channels/any_value.py

@@ -0,0 +1,55 @@
+from contextlib import contextmanager
+from typing import Generator, Generic, Optional, Sequence, Type
+
+from typing_extensions import Self
+
+from langgraph.channels.base import BaseChannel, EmptyChannelError, Value
+
+
+class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the last value received, assumes that if multiple values are
+    received, they are all equal."""
+
+    def __init__(self, typ: Type[Value]) -> None:
+        self.typ = typ
+
+    @property
+    def ValueType(self) -> Type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> Type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ)
+        if checkpoint is not None:
+            empty.value = checkpoint
+        try:
+            yield empty
+        finally:
+            try:
+                del empty.value
+            except AttributeError:
+                pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        if len(values) == 0:
+            return
+
+        self.value = values[-1]
+
+    def get(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()

langgraph/channels/ephemeral_value.py

@@ -0,0 +1,67 @@
+from contextlib import contextmanager
+from typing import Generator, Generic, Optional, Sequence, Type
+
+from typing_extensions import Self
+
+from langgraph.channels.base import (
+    BaseChannel,
+    EmptyChannelError,
+    InvalidUpdateError,
+    Value,
+)
+
+
+class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
+    """Stores the value received in the step immediately preceding, clears after."""
+
+    def __init__(self, typ: Type[Value], guard: bool = True) -> None:
+        self.typ = typ
+        self.guard = guard
+
+    @property
+    def ValueType(self) -> Type[Value]:
+        """The type of the value stored in the channel."""
+        return self.typ
+
+    @property
+    def UpdateType(self) -> Type[Value]:
+        """The type of the update received by the channel."""
+        return self.typ
+
+    @contextmanager
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
+        empty = self.__class__(self.typ, self.guard)
+        if checkpoint is not None:
+            empty.value = checkpoint
+        try:
+            yield empty
+        finally:
+            try:
+                del empty.value
+            except AttributeError:
+                pass
+
+    def update(self, values: Sequence[Value]) -> None:
+        if len(values) == 0:
+            try:
+                del self.value
+            except AttributeError:
+                pass
+            finally:
+                return
+        if len(values) != 1 and self.guard:
+            raise InvalidUpdateError("LastValue can only receive one value per step.")
+
+        self.value = values[-1]
+
+    def get(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()
+
+    def checkpoint(self) -> Value:
+        try:
+            return self.value
+        except AttributeError:
+            raise EmptyChannelError()

langgraph/graph/state.py

@@ -6,8 +6,10 @@ from typing import Any, Optional, Sequence, Type
 from langchain_core.runnables import RunnableLambda, RunnablePassthrough
 from langchain_core.runnables.base import RunnableLike
 
+from langgraph.channels.any_value import AnyValue
 from langgraph.channels.base import BaseChannel, InvalidUpdateError
 from langgraph.channels.binop import BinaryOperatorAggregate
+from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, Graph
@@ -73,6 +75,17 @@ class StateGraph(Graph):
             )
             for key, node in self.nodes.items()
         }
+        node_inboxes = {
+            # we take any value written to channel because all writers
+            # write the entire state as of that step, which is equal for all
+            f"{key}:inbox": AnyValue(Any)
+            for key in self.nodes
+        }
+        node_outboxes = {
+            # we clear outbox channels after each step
+            key: EphemeralValue(Any)
+            for key in self.nodes
+        }
 
         for key in self.nodes:
             outgoing = outgoing_edges[key]
@@ -100,7 +113,7 @@ class StateGraph(Graph):
 
         return Pregel(
             nodes=nodes,
-            channels=self.channels,
+            channels={**self.channels, **node_inboxes, **node_outboxes},
             input=f"{START}:inbox",
             output=END,
             hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,

langgraph/pregel/__init__.py

@@ -789,7 +789,7 @@ def _prepare_next_tasks(
                 checkpoint["channel_versions"][chan] > seen[chan]
                 for chan in proc.triggers
             ):
-                # If all channels subscribed by this process have been initialized
+                # If all channels subscribed by this process are not empty
                 try:
                     val: Any = {
                         k: _read_channel(
@@ -819,9 +819,11 @@ def _prepare_next_tasks(
         elif isinstance(proc, ChannelBatch):
             # If the channel read by this process was updated
             if checkpoint["channel_versions"][proc.channel] > seen[proc.channel]:
-                # Here we don't catch EmptyChannelError because the channel
-                # must be intialized if the previous `if` condition is true
-                val = channels[proc.channel].get()
+                # If the channel subscribed by this process is not empty
+                try:
+                    val = channels[proc.channel].get()
+                except EmptyChannelError:
+                    continue
                 if proc.key is not None:
                     val = [{proc.key: v} for v in val]
 

tests/test_pregel.py

@@ -1628,3 +1628,63 @@ def test_message_graph() -> None:
             ]
         },
     ]
+
+
+def test_in_one_fan_out_out_one_graph_state() -> None:
+    def sorted_add(x: list[str], y: list[str]) -> list[str]:
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    def rewrite_query(data: State) -> State:
+        return {"query": f'query: {data["query"]}'}
+
+    def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    def retriever_two(data: State) -> State:
+        return {"docs": ["doc3", "doc4"]}
+
+    def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("qa", qa)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_edge("rewrite_query", "retriever_one")
+    workflow.add_edge("rewrite_query", "retriever_two")
+    workflow.add_edge("retriever_one", "qa")
+    workflow.add_edge("retriever_two", "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert app.invoke({"query": "what is weather in sf"}) == {
+        "query": "query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [*app.stream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]

tests/test_pregel_async.py

@@ -1685,3 +1685,63 @@ async def test_message_graph() -> None:
             ]
         },
     ]
+
+
+async def test_in_one_fan_out_out_one_graph_state() -> None:
+    def sorted_add(x: list[str], y: list[str]) -> list[str]:
+        return sorted(operator.add(x, y))
+
+    class State(TypedDict, total=False):
+        query: str
+        answer: str
+        docs: Annotated[list[str], sorted_add]
+
+    async def rewrite_query(data: State) -> State:
+        return {"query": f'query: {data["query"]}'}
+
+    async def retriever_one(data: State) -> State:
+        return {"docs": ["doc1", "doc2"]}
+
+    async def retriever_two(data: State) -> State:
+        return {"docs": ["doc3", "doc4"]}
+
+    async def qa(data: State) -> State:
+        return {"answer": ",".join(data["docs"])}
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("qa", qa)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_edge("rewrite_query", "retriever_one")
+    workflow.add_edge("rewrite_query", "retriever_two")
+    workflow.add_edge("retriever_one", "qa")
+    workflow.add_edge("retriever_two", "qa")
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"query": "what is weather in sf"}) == {
+        "query": "query: what is weather in sf",
+        "docs": ["doc1", "doc2", "doc3", "doc4"],
+        "answer": "doc1,doc2,doc3,doc4",
+    }
+
+    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"qa": {"answer": "doc1,doc2,doc3,doc4"}},
+        {
+            "__end__": {
+                "query": "query: what is weather in sf",
+                "answer": "doc1,doc2,doc3,doc4",
+                "docs": ["doc1", "doc2", "doc3", "doc4"],
+            }
+        },
+    ]