Handle Async Callable (#457)

langgraph/utils.py

@@ -4,7 +4,7 @@ import inspect
 import sys
 from contextvars import copy_context
 from functools import partial, wraps
-from typing import Any, Awaitable, Callable, Optional
+from typing import Any, AsyncIterator, Awaitable, Callable, Optional
 
 from langchain_core.runnables.base import (
     Runnable,
@@ -20,6 +20,7 @@ from langchain_core.runnables.config import (
 )
 from langchain_core.runnables.graph import Edge, Graph, Node, is_uuid
 from langchain_core.runnables.utils import accepts_config
+from typing_extensions import TypeGuard
 
 
 # Before Python 3.11 native StrEnum is not available
@@ -148,6 +149,28 @@ class DrawableGraph(Graph):
         )
 
 
+def is_async_callable(
+    func: Any,
+) -> TypeGuard[Callable[..., Awaitable]]:
+    """Check if a function is async."""
+    return (
+        asyncio.iscoroutinefunction(func)
+        or hasattr(func, "__call__")
+        and asyncio.iscoroutinefunction(func.__call__)
+    )
+
+
+def is_async_generator(
+    func: Any,
+) -> TypeGuard[Callable[..., AsyncIterator]]:
+    """Check if a function is an async generator."""
+    return (
+        inspect.isasyncgenfunction(func)
+        or hasattr(func, "__call__")
+        and inspect.isasyncgenfunction(func.__call__)
+    )
+
+
 def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnable:
     """Coerce a runnable-like object into a Runnable.
 
@@ -159,10 +182,10 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab
     """
     if isinstance(thing, Runnable):
         return thing
-    elif inspect.isasyncgenfunction(thing) or inspect.isgeneratorfunction(thing):
+    elif is_async_generator(thing) or inspect.isgeneratorfunction(thing):
         return RunnableLambda(thing, name=name)
     elif callable(thing):
-        if asyncio.iscoroutinefunction(thing):
+        if is_async_callable(thing):
             return RunnableCallable(None, thing, name=name, trace=trace)
         else:
             return RunnableCallable(

tests/test_utils.py

@@ -0,0 +1,72 @@
+import functools
+
+from langgraph.utils import is_async_callable, is_async_generator
+
+
+def test_is_async() -> None:
+    async def func() -> None:
+        pass
+
+    assert is_async_callable(func)
+    wrapped_func = functools.wraps(func)(func)
+    assert is_async_callable(wrapped_func)
+
+    def sync_func() -> None:
+        pass
+
+    assert not is_async_callable(sync_func)
+    wrapped_sync_func = functools.wraps(sync_func)(sync_func)
+    assert not is_async_callable(wrapped_sync_func)
+
+    class AsyncFuncCallable:
+        async def __call__(self) -> None:
+            pass
+
+    runnable = AsyncFuncCallable()
+    assert is_async_callable(runnable)
+    wrapped_runnable = functools.wraps(runnable)(runnable)
+    assert is_async_callable(wrapped_runnable)
+
+    class SyncFuncCallable:
+        def __call__(self) -> None:
+            pass
+
+    sync_runnable = SyncFuncCallable()
+    assert not is_async_callable(sync_runnable)
+    wrapped_sync_runnable = functools.wraps(sync_runnable)(sync_runnable)
+    assert not is_async_callable(wrapped_sync_runnable)
+
+
+def test_is_generator() -> None:
+    async def gen():
+        yield
+
+    assert is_async_generator(gen)
+
+    wrapped_gen = functools.wraps(gen)(gen)
+    assert is_async_generator(wrapped_gen)
+
+    def sync_gen():
+        yield
+
+    assert not is_async_generator(sync_gen)
+    wrapped_sync_gen = functools.wraps(sync_gen)(sync_gen)
+    assert not is_async_generator(wrapped_sync_gen)
+
+    class AsyncGenCallable:
+        async def __call__(self):
+            yield
+
+    runnable = AsyncGenCallable()
+    assert is_async_generator(runnable)
+    wrapped_runnable = functools.wraps(runnable)(runnable)
+    assert is_async_generator(wrapped_runnable)
+
+    class SyncGenCallable:
+        def __call__(self):
+            yield
+
+    sync_runnable = SyncGenCallable()
+    assert not is_async_generator(sync_runnable)
+    wrapped_sync_runnable = functools.wraps(sync_runnable)(sync_runnable)
+    assert not is_async_generator(wrapped_sync_runnable)