From e79453633d94d05463f2dc0505a5c055ca44104c Mon Sep 17 00:00:00 2001
From: jacoblee93 <jacoblee93@gmail.com>
Date: Tue, 26 Mar 2024 12:59:26 -0700
Subject: [PATCH] Simplify

---
 langgraph/graph/graph.py | 39 ++++++++++++++-------------------------
 langgraph/graph/state.py |  2 +-
 2 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py
index a91097684..baf3d4e19 100644
--- a/langgraph/graph/graph.py
+++ b/langgraph/graph/graph.py
@@ -1,7 +1,6 @@
 import logging
-from asyncio import get_running_loop, iscoroutinefunction
+from asyncio import iscoroutinefunction
 from collections import defaultdict
-from functools import partial
 from typing import Any, Callable, Dict, NamedTuple, Optional, Sequence
 
 from langchain_core.runnables import Runnable
@@ -32,14 +31,14 @@ class Branch(NamedTuple):
     condition: Callable[..., str]
     ends: Optional[dict[str, str]]
 
-    def is_coroutine(self):
-        return iscoroutinefunction(self.condition)
+    @property
+    def runnable(self):
+        if iscoroutinefunction(self.condition):
+            return self._arunnable
+        else:
+            return self._runnable
 
-    def runnable(self, input: Any) -> Runnable:
-        if self.is_coroutine():
-            raise ValueError(
-                "All conditions must be sync when invoking graphs synchronously."
-            )
+    def _runnable(self, input: Any) -> Runnable:
         result = self.condition(input)
         if self.ends:
             destination = self.ends[result]
@@ -47,14 +46,8 @@ 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),
-            )
+    async def _arunnable(self, input: Any) -> Runnable:
+        result = await self.condition(input)
         if self.ends:
             destination = self.ends[result]
         else:
@@ -237,7 +230,7 @@ class Graph:
             if key in self.branches:
                 for branch in self.branches[key]:
                     nodes[edges_key] |= RunnableLambda(
-                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        branch.runnable,
                         name=f"{key}_condition",
                     )
 
@@ -245,9 +238,7 @@ class Graph:
             nodes[f"{START}:edges"] = Channel.subscribe_to(
                 START, tags=["langsmith:hidden"]
             ) | RunnableLambda(
-                self.entry_point_branch.arunnable
-                if self.entry_point_branch.is_coroutine()
-                else self.entry_point_branch.runnable,
+                self.entry_point_branch.runnable,
                 name=f"{START}_condition",
             )
         elif self.entry_point is None:
@@ -312,7 +303,7 @@ class CompiledGraph(Pregel):
                     name += f"_{i}"
                 cond = graph.add_node(
                     RunnableLambda(
-                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        branch.runnable,
                         name=branch.condition.__name__,
                     ),
                     name,
@@ -327,9 +318,7 @@ class CompiledGraph(Pregel):
         if self.graph.entry_point_branch:
             cond = graph.add_node(
                 RunnableLambda(
-                    self.graph.entry_point_branch.arunnable
-                    if self.graph.entry_point_branch.is_coroutine()
-                    else self.graph.entry_point_branch.runnable,
+                    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 b0d868a4f..bc12713d4 100644
--- a/langgraph/graph/state.py
+++ b/langgraph/graph/state.py
@@ -160,7 +160,7 @@ class StateGraph(Graph):
             if key in self.branches:
                 for branch in self.branches[key]:
                     nodes[edges_key] |= RunnableLambda(
-                        branch.arunnable if branch.is_coroutine() else branch.runnable,
+                        branch.runnable,
                         name=f"{key}_condition",
                     )