diff --git a/langgraph/utils.py b/langgraph/utils.py index 941b06b00..f04388677 100644 --- a/langgraph/utils.py +++ b/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( diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 000000000..ca16e4c4c --- /dev/null +++ b/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)