Merge pull request #288 from langchain-ai/nc/8apr/cond-edge-multiple-destinations

feat: Return multiple destinations from conditional edge

langgraph/graph/graph.py

@@ -37,12 +37,12 @@ END = "__end__"
 
 
 class Branch(NamedTuple):
-    condition: Runnable[Any, str]
+    condition: Runnable[Any, Union[str, list[str]]]
     ends: Optional[dict[str, str]]
 
     def run(
         self,
-        writer: Callable[[str], Optional[Runnable]],
+        writer: Callable[[list[str]], Optional[Runnable]],
         reader: Optional[Callable[[RunnableConfig], Any]] = None,
     ) -> None:
         return ChannelWrite.register_writer(
@@ -65,11 +65,13 @@ class Branch(NamedTuple):
         writer: Callable[[str], Optional[Runnable]],
     ) -> Runnable:
         result = self.condition.invoke(reader(config) if reader else input, config)
+        if isinstance(result, str):
+            result = [result]
         if self.ends:
-            destination = self.ends[result]
+            destinations = [self.ends[r] for r in result]
         else:
-            destination = result
-        return writer(destination)
+            destinations = result
+        return writer(destinations)
 
     async def _aroute(
         self,
@@ -82,11 +84,13 @@ class Branch(NamedTuple):
         result = await self.condition.ainvoke(
             reader(config) if reader else input, config
         )
+        if isinstance(result, str):
+            result = [result]
         if self.ends:
-            destination = self.ends[result]
+            destinations = [self.ends[r] for r in result]
         else:
-            destination = result
-        return writer(destination)
+            destinations = result
+        return writer(destinations)
 
 
 class Graph:
@@ -143,7 +147,9 @@ class Graph:
         self,
         start_key: str,
         condition: Union[
-            Callable[..., str], Callable[..., Awaitable[str]], Runnable[Any, str]
+            Callable[..., Union[str, list[str]]],
+            Callable[..., Awaitable[Union[str, list[str]]]],
+            Runnable[Any, Union[str, list[str]]],
         ],
         conditional_edge_mapping: Optional[dict[str, str]] = None,
     ) -> None:
@@ -286,11 +292,11 @@ class CompiledGraph(Pregel):
             self.nodes[end].channels.append(start)
 
     def attach_branch(self, start: str, name: str, branch: Branch) -> None:
-        def branch_writer(end: str) -> Optional[ChannelWrite]:
-            return Channel.write_to(
-                f"branch:{start}:{name}:{end}" if end != END else END,
-                tags=[TAG_HIDDEN],
-            )
+        def branch_writer(ends: list[str]) -> Optional[ChannelWrite]:
+            channels = [
+                f"branch:{start}:{name}:{end}" if end != END else END for end in ends
+            ]
+            return Channel.write_to(*channels, tags=[TAG_HIDDEN])
 
         # add hidden start node
         if start == START and start not in self.nodes:

langgraph/graph/state.py

@@ -206,10 +206,13 @@ class CompiledStateGraph(CompiledGraph):
                 )
 
     def attach_branch(self, start: str, name: str, branch: Branch) -> None:
-        def branch_writer(end: str) -> Optional[ChannelWrite]:
-            if end != END:
+        def branch_writer(ends: list[str]) -> Optional[ChannelWrite]:
+            if filtered_ends := [end for end in ends if end != END]:
                 return ChannelWrite(
-                    [ChannelWriteEntry(f"branch:{start}:{name}:{end}", start)],
+                    [
+                        ChannelWriteEntry(f"branch:{start}:{name}:{end}", start)
+                        for end in filtered_ends
+                    ],
                     tags=[TAG_HIDDEN],
                 )
 

tests/test_pregel.py

@@ -3317,6 +3317,93 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
     ]
 
 
+def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        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_picker(data: State) -> list[str]:
+        return ["analyzer_one", "retriever_two"]
+
+    def analyzer_one(data: State) -> State:
+        return {"query": f'analyzed: {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"])}
+
+    def decider(data: State) -> None:
+        return None
+
+    def decider_cond(data: State) -> str:
+        if data["query"].count("analyzed") > 1:
+            return "qa"
+        else:
+            return "rewrite_query"
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("analyzer_one", analyzer_one)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("decider", decider)
+    workflow.add_node("qa", qa)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_conditional_edges("rewrite_query", retriever_picker)
+    workflow.add_edge("analyzer_one", "retriever_one")
+    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
+    workflow.add_conditional_edges("decider", decider_cond)
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert app.invoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: analyzed: query: what is weather in sf",
+        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+    }
+
+    assert [*app.stream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "analyzer_one": {"query": "analyzed: query: what is weather in sf"},
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "analyzer_one": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf"
+            },
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
+    ]
+
+
 def test_simple_multi_edge() -> None:
     class State(TypedDict):
         my_key: Annotated[str, operator.add]

tests/test_pregel_async.py

@@ -3099,3 +3099,90 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         },
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
+
+
+async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
+    def sorted_add(
+        x: list[str], y: Union[list[str], list[tuple[str, str]]]
+    ) -> list[str]:
+        if isinstance(y[0], tuple):
+            for rem, _ in y:
+                x.remove(rem)
+            y = [t[1] for t in y]
+        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_picker(data: State) -> list[str]:
+        return ["analyzer_one", "retriever_two"]
+
+    async def analyzer_one(data: State) -> State:
+        return {"query": f'analyzed: {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"])}
+
+    async def decider(data: State) -> None:
+        return None
+
+    def decider_cond(data: State) -> str:
+        if data["query"].count("analyzed") > 1:
+            return "qa"
+        else:
+            return "rewrite_query"
+
+    workflow = StateGraph(State)
+
+    workflow.add_node("rewrite_query", rewrite_query)
+    workflow.add_node("analyzer_one", analyzer_one)
+    workflow.add_node("retriever_one", retriever_one)
+    workflow.add_node("retriever_two", retriever_two)
+    workflow.add_node("decider", decider)
+    workflow.add_node("qa", qa)
+
+    workflow.set_entry_point("rewrite_query")
+    workflow.add_conditional_edges("rewrite_query", retriever_picker)
+    workflow.add_edge("analyzer_one", "retriever_one")
+    workflow.add_edge(["retriever_one", "retriever_two"], "decider")
+    workflow.add_conditional_edges("decider", decider_cond)
+    workflow.set_finish_point("qa")
+
+    app = workflow.compile()
+
+    assert await app.ainvoke({"query": "what is weather in sf"}) == {
+        "query": "analyzed: query: analyzed: query: what is weather in sf",
+        "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4",
+        "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"],
+    }
+
+    assert [c async for c in app.astream({"query": "what is weather in sf"})] == [
+        {"rewrite_query": {"query": "query: what is weather in sf"}},
+        {
+            "analyzer_one": {"query": "analyzed: query: what is weather in sf"},
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
+        {
+            "analyzer_one": {
+                "query": "analyzed: query: analyzed: query: what is weather in sf"
+            },
+            "retriever_two": {"docs": ["doc3", "doc4"]},
+        },
+        {
+            "retriever_one": {"docs": ["doc1", "doc2"]},
+        },
+        {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
+    ]