From 1093dd55c87c17b3e6139e08dd236ac60b839226 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 23 Jan 2025 16:57:34 -0800 Subject: [PATCH 1/2] Improve typings for task, it now returns a mixed sync/async future protocol - note this type is never instantiated, it is only used for typing (we cannot make it a protocol as it inherits from concurrent.futures.Future) --- libs/langgraph/langgraph/func/__init__.py | 30 +++++++------------ libs/langgraph/langgraph/pregel/call.py | 10 +++++-- libs/langgraph/tests/test_pregel_async.py | 35 +++++++++++++++++++++++ 3 files changed, 53 insertions(+), 22 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 40cff2ada..55ea88e02 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"]) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 35d6143dd..9b702f29e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -3,6 +3,7 @@ import logging import operator import random import sys +import time import uuid from collections import Counter, deque from contextlib import asynccontextmanager, contextmanager @@ -2553,6 +2554,40 @@ async def test_imp_nested(checkpointer_name: str) -> None: "11answera", ] + def syncmapper(input: int) -> str: + time.sleep(input / 100) + return submapper(input).result() * 2 + + @entrypoint(checkpointer=checkpointer) + async def graph(input: list[int]) -> list[str]: + mapped = [syncmapper(i) for i in input] + answer = interrupt("question") + final = [m + answer for m in mapped] + return await add_a.ainvoke(final) + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream([0, 1], thread1)] == [ + {"submapper": "0"}, + {"syncmapper": "00"}, + {"submapper": "1"}, + {"syncmapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), + ) + }, + ] + + assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + "00answera", + "11answera", + ] + @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) From 996b1206134c66cd9e7e3731641ee14eea689033 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 23 Jan 2025 16:59:27 -0800 Subject: [PATCH 2/2] Undo --- libs/langgraph/tests/test_pregel_async.py | 35 ----------------------- 1 file changed, 35 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9b702f29e..35d6143dd 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -3,7 +3,6 @@ import logging import operator import random import sys -import time import uuid from collections import Counter, deque from contextlib import asynccontextmanager, contextmanager @@ -2554,40 +2553,6 @@ async def test_imp_nested(checkpointer_name: str) -> None: "11answera", ] - def syncmapper(input: int) -> str: - time.sleep(input / 100) - return submapper(input).result() * 2 - - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - mapped = [syncmapper(i) for i in input] - answer = interrupt("question") - final = [m + answer for m in mapped] - return await add_a.ainvoke(final) - - thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream([0, 1], thread1)] == [ - {"submapper": "0"}, - {"syncmapper": "00"}, - {"submapper": "1"}, - {"syncmapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - - assert await graph.ainvoke(Command(resume="answer"), thread1) == [ - "00answera", - "11answera", - ] - @NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)