From 3926e83884149ddbc2f7b88ffd4dc3dc774743b2 Mon Sep 17 00:00:00 2001 From: JP-Ellis Date: Tue, 4 Mar 2025 16:54:53 +1100 Subject: [PATCH] chore(langgraph): fix typing of task decorator An async function of the form `def foo(P) -> T` has type `Callable[[P], Awaitable[T]]`. The old type annotations then converted the function into a `Callable[[P], SyncAsyncFuture[Awaitable[T]]]` which is incorrect. The change introduced in this commit updates the type annotations to ensure the `Awaitable[T]` is correctly unwrapped. Signed-off-by: JP-Ellis --- libs/langgraph/langgraph/func/__init__.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index a3d012e55..335d7b78e 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -36,23 +36,31 @@ from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode @overload def task( - *, name: Optional[str] = None, retry: Optional[RetryPolicy] = None -) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ... + *, + name: Optional[str] = None, + retry: Optional[RetryPolicy] = None, +) -> Callable[ + [Union[Callable[P, Awaitable[T]], Callable[P, T]]], + Callable[P, SyncAsyncFuture[T]], +]: ... @overload def task( - __func_or_none__: Callable[P, T], + __func_or_none__: Union[Callable[P, Awaitable[T]], Callable[P, T]], ) -> Callable[P, SyncAsyncFuture[T]]: ... def task( - __func_or_none__: Optional[Union[Callable[P, T], Callable[P, Awaitable[T]]]] = None, + __func_or_none__: Optional[Union[Callable[P, Awaitable[T]], Callable[P, T]]] = None, *, name: Optional[str] = None, retry: Optional[RetryPolicy] = None, ) -> Union[ - Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]], + Callable[ + [Union[Callable[P, Awaitable[T]], Callable[P, T]]], + Callable[P, SyncAsyncFuture[T]], + ], Callable[P, SyncAsyncFuture[T]], ]: """Define a LangGraph task using the `task` decorator.