diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index ee8ece166..d0f429776 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -22,7 +22,13 @@ from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN from langgraph.pregel import Pregel -from langgraph.pregel.call import P, T, call, get_runnable_for_entrypoint +from langgraph.pregel.call import ( + P, + SyncAsyncFuture, + T, + call, + get_runnable_for_entrypoint, +) from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore @@ -32,25 +38,13 @@ from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode, StreamWriter @overload def task( *, retry: Optional[RetryPolicy] = None -) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]]: ... - - -@overload -def task( # type: ignore[overload-cannot-match] - *, retry: Optional[RetryPolicy] = None -) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ... +) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ... @overload def task( __func_or_none__: Callable[P, T], -) -> Callable[P, concurrent.futures.Future[T]]: ... - - -@overload -def task( - __func_or_none__: Callable[P, Awaitable[T]], -) -> Callable[P, asyncio.Future[T]]: ... +) -> Callable[P, SyncAsyncFuture[T]]: ... def task( @@ -58,10 +52,8 @@ def task( *, retry: Optional[RetryPolicy] = None, ) -> Union[ - Callable[[Callable[P, Awaitable[T]]], Callable[P, asyncio.Future[T]]], - Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]], - Callable[P, asyncio.Future[T]], - Callable[P, concurrent.futures.Future[T]], + Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]], + Callable[P, SyncAsyncFuture[T]], ]: """Define a LangGraph task using the `task` decorator. diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index a4bba63fe..1ddad1965 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,12 +1,11 @@ """Utility to convert a user provided function into a Runnable with a ChannelWrite.""" -import asyncio import concurrent.futures import functools import inspect import sys import types -from typing import Any, Callable, Optional, TypeVar, Union +from typing import Any, Callable, Generator, Generic, Optional, TypeVar, cast from langchain_core.runnables import Runnable from typing_extensions import ParamSpec @@ -208,12 +207,17 @@ P1 = TypeVar("P1") T = TypeVar("T") +class SyncAsyncFuture(Generic[T], concurrent.futures.Future[T]): + def __await__(self) -> Generator[T, None, T]: + yield cast(T, ...) + + def call( func: Callable[P, T], *args: Any, retry: Optional[RetryPolicy] = None, **kwargs: Any, -) -> Union[concurrent.futures.Future[T], asyncio.Future[T]]: +) -> SyncAsyncFuture[T]: config = get_config() impl = config[CONF][CONFIG_KEY_CALL] fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"])