From a6a9def91a81387b8d3b230bea807e24ccd17c03 Mon Sep 17 00:00:00 2001
From: jacoblee93 <jacoblee93@gmail.com>
Date: Tue, 26 Mar 2024 12:46:15 -0700
Subject: [PATCH] Adds async conditional edge support

---
 langgraph/graph/graph.py   | 45 ++++++++++++++++++++++++++++++--------
 langgraph/graph/state.py   |  3 ++-
 tests/test_pregel_async.py |  2 +-
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index e3f15798b..a91097684 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -1,6 +1,7 @@
 import logging
-from asyncio import iscoroutinefunction
+from asyncio import get_running_loop, iscoroutinefunction
 from collections import defaultdict
+from functools import partial
 from typing import Any, Callable, Dict, NamedTuple, Optional, Sequence
 
 from langchain_core.runnables import Runnable
@@ -31,7 +32,14 @@ class Branch(NamedTuple):
     condition: Callable[..., str]
     ends: Optional[dict[str, str]]
 
+    def is_coroutine(self):
+        return iscoroutinefunction(self.condition)
+
     def runnable(self, input: Any) -> Runnable:
+        if self.is_coroutine():
+            raise ValueError(
+                "All conditions must be sync when invoking graphs synchronously."
+            )
         result = self.condition(input)
         if self.ends:
             destination = self.ends[result]
@@ -39,6 +47,20 @@ class Branch(NamedTuple):
             destination = result
         return Channel.write_to(f"{destination}:inbox" if destination != END else END)
 
+    async def arunnable(self, input: Any) -> Runnable:
+        if self.is_coroutine():
+            result = await self.condition(input)
+        else:
+            result = await get_running_loop().run_in_executor(
+                None,
+                partial(self.condition, input),
+            )
+        if self.ends:
+            destination = self.ends[result]
+        else:
+            destination = result
+        return Channel.write_to(f"{destination}:inbox" if destination != END else END)
+
 
 class Graph:
     def __init__(self) -> None:
@@ -100,8 +122,6 @@ class Graph:
             )
         if start_key not in self.nodes:
             raise ValueError(f"Need to add_node `{start_key}` first")
-        if iscoroutinefunction(condition):
-            raise ValueError("Condition cannot be a coroutine function")
         if conditional_edge_mapping and set(
             conditional_edge_mapping.values()
         ).difference([END]).difference(self.nodes):
@@ -134,8 +154,6 @@ class Graph:
                 "Setting the entry point of a graph that has already been compiled. "
                 "This will not be reflected in the compiled graph."
             )
-        if iscoroutinefunction(condition):
-            raise ValueError("Condition cannot be a coroutine function")
         if conditional_edge_mapping and set(
             conditional_edge_mapping.values()
         ).difference([END]).difference(self.nodes):
@@ -219,14 +237,18 @@ class Graph:
             if key in self.branches:
                 for branch in self.branches[key]:
                     nodes[edges_key] |= RunnableLambda(
-                        branch.runnable, name=f"{key}_condition"
+                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        name=f"{key}_condition",
                     )
 
         if self.entry_point_branch:
             nodes[f"{START}:edges"] = Channel.subscribe_to(
                 START, tags=["langsmith:hidden"]
             ) | RunnableLambda(
-                self.entry_point_branch.runnable, name=f"{START}_condition"
+                self.entry_point_branch.arunnable
+                if self.entry_point_branch.is_coroutine()
+                else self.entry_point_branch.runnable,
+                name=f"{START}_condition",
             )
         elif self.entry_point is None:
             raise ValueError("No entry point set")
@@ -289,7 +311,10 @@ class CompiledGraph(Pregel):
                 if i > 0:
                     name += f"_{i}"
                 cond = graph.add_node(
-                    RunnableLambda(branch.runnable, name=branch.condition.__name__),
+                    RunnableLambda(
+                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        name=branch.condition.__name__,
+                    ),
                     name,
                 )
                 graph.add_edge(start_nodes[start], cond)
@@ -302,7 +327,9 @@ class CompiledGraph(Pregel):
         if self.graph.entry_point_branch:
             cond = graph.add_node(
                 RunnableLambda(
-                    self.graph.entry_point_branch.runnable,
+                    self.graph.entry_point_branch.arunnable
+                    if self.graph.entry_point_branch.is_coroutine()
+                    else self.graph.entry_point_branch.runnable,
                     name=self.graph.entry_point_branch.condition.__name__,
                 ),
                 f"{START}_condition",
diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py
index 3bb0adf5c..b0d868a4f 100644
--- a/langgraph/graph/state.py
+++ b/langgraph/graph/state.py
@@ -160,7 +160,8 @@ class StateGraph(Graph):
             if key in self.branches:
                 for branch in self.branches[key]:
                     nodes[edges_key] |= RunnableLambda(
-                        branch.runnable, name=f"{key}_condition"
+                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        name=f"{key}_condition",
                     )
 
         nodes[START] = Channel.subscribe_to(
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index 5ada8b3c3..3024b1284 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -820,7 +820,7 @@ async def test_conditional_graph() -> None:
         return data
 
     # Define decision-making logic
-    def should_continue(data: dict) -> str:
+    async def should_continue(data: dict) -> str:
         # Logic to decide whether to continue in the loop or exit
         if isinstance(data["agent_outcome"], AgentFinish):
             return "exit"