From fd6f44a7e407c32eeaf911805ba8937feb894aa6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 8 Apr 2024 17:05:16 -0700 Subject: [PATCH] feat: Return multiple destinations from conditional edge --- langgraph/graph/graph.py | 34 +++++++++------ langgraph/graph/state.py | 9 ++-- tests/test_pregel.py | 87 ++++++++++++++++++++++++++++++++++++++ tests/test_pregel_async.py | 87 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 200 insertions(+), 17 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 1a4a9dfc7..d637beaaf 100644 --- a/langgraph/graph/graph.py +++ b/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: diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 655730b37..33114a795 100644 --- a/langgraph/graph/state.py +++ b/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], ) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index ac80fb981..8e091aa09 100644 --- a/tests/test_pregel.py +++ b/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] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6f591767a..6f4d70128 100644 --- a/tests/test_pregel_async.py +++ b/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"}}, + ]