chore(langgraph): fix typing of task decorator (#3670)

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.

I've tested it locally and confirmed it work on:

```python
@task 
def sync_fn(a: int) -> int: ...

@task 
def async_fn(a: int) -> int: ...
```

Let me know if you want me to add tests, just let me know how you test
type annotations.
This commit is contained in:
Nuno Campos
2025-03-04 17:05:35 -08:00
committed by GitHub
+13 -5
View File
@@ -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.