Fix stream order

libs/langgraph/langgraph/func/__init__.py

@@ -23,7 +23,7 @@ from langgraph.pregel.call import get_runnable_for_func
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
-from langgraph.types import RetryPolicy, acall, call
+from langgraph.types import RetryPolicy, call
 
 P = ParamSpec("P")
 T = TypeVar("T")
@@ -48,10 +48,7 @@ def task(
     Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]],
 ]:
     def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]:
-        if asyncio.iscoroutinefunction(func):
-            return update_wrapper(partial(acall, func), func)
-        else:
-            return update_wrapper(partial(call, func), func)
+        return update_wrapper(partial(call, func), func)
 
     return _task
 

libs/langgraph/langgraph/pregel/call.py

@@ -1,11 +1,10 @@
-import asyncio
 import sys
 import types
 from typing import Any, Callable, Optional
 
 from langgraph.constants import RETURN
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
-from langgraph.utils.runnable import RunnableCallable, RunnableSeq
+from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable
 
 """
 Utilities borrowed from cloudpickle.
@@ -110,25 +109,15 @@ def _lookup_module_and_qualname(
 def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq:
     if func in CACHE:
         return CACHE[func]
-    elif not _lookup_module_and_qualname(func):
-        return RunnableSeq(
-            RunnableCallable(None, func, trace=False)
-            if asyncio.iscoroutinefunction(func)
-            else RunnableCallable(func, trace=False),
+    else:
+        seq = RunnableSeq(
+            coerce_to_runnable(func, name=None, trace=False),
             ChannelWrite([ChannelWriteEntry(RETURN)]),
             name=func.__name__,
         )
-    else:
-        return CACHE.setdefault(
-            func,
-            RunnableSeq(
-                RunnableCallable(None, func, trace=False)
-                if asyncio.iscoroutinefunction(func)
-                else RunnableCallable(func, trace=False),
-                ChannelWrite([ChannelWriteEntry(RETURN)]),
-                name=func.__name__,
-            ),
-        )
+        if not _lookup_module_and_qualname(func):
+            return seq
+        return CACHE.setdefault(func, seq)
 
 
 CACHE: dict[Callable[..., Any], RunnableSeq] = {}

libs/langgraph/langgraph/pregel/runner.py

@@ -34,6 +34,7 @@ from langgraph.pregel.algo import Call
 from langgraph.pregel.executor import Submit
 from langgraph.pregel.retry import arun_with_retry, run_with_retry
 from langgraph.types import PregelExecutableTask, RetryPolicy
+from langgraph.utils.future import chain_future
 
 
 class PregelRunner:
@@ -133,6 +134,7 @@ class PregelRunner:
                             },
                             __reraise_on_exit__=reraise,
                         )
+                        fut.add_done_callback(partial(self.commit, next_task))
                         futures[fut] = next_task
                         rtn[idx - prev_length] = fut
             return [rtn.get(i) for i in range(len(writes))]
@@ -165,7 +167,7 @@ class PregelRunner:
                 )
                 self.commit(t, None)
             except Exception as exc:
-                self.commit(t, exc)
+                self.commit(t, None, exc)
                 if reraise and futures:
                     # will be re-raised after futures are done
                     fut: concurrent.futures.Future = concurrent.futures.Future()
@@ -183,18 +185,18 @@ class PregelRunner:
         # yield updates/debug output as each task finishes
         for t in tasks:
             if not t.writes:
-                futures[
-                    self.submit(
-                        run_with_retry,
-                        t,
-                        retry_policy,
-                        configurable={
-                            CONFIG_KEY_SEND: partial(writer, t),
-                            CONFIG_KEY_CALL: partial(call, t),
-                        },
-                        __reraise_on_exit__=reraise,
-                    )
-                ] = t
+                fut = self.submit(
+                    run_with_retry,
+                    t,
+                    retry_policy,
+                    configurable={
+                        CONFIG_KEY_SEND: partial(writer, t),
+                        CONFIG_KEY_CALL: partial(call, t),
+                    },
+                    __reraise_on_exit__=reraise,
+                )
+                fut.add_done_callback(partial(self.commit, t))
+                futures[fut] = t
         end_time = timeout + time.monotonic() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = concurrent.futures.wait(
@@ -213,8 +215,6 @@ class PregelRunner:
                 else:
                     # store for panic check
                     done_futures.add(fut)
-                    # task finished, commit writes
-                    self.commit(task, _exception(fut))
             else:
                 # remove references to loop vars
                 del fut, task
@@ -264,11 +264,8 @@ class PregelRunner:
                 if w[0] != PUSH:
                     continue
                 # schedule the next task, if the callback returns one
-                if next_task := self.schedule_task(
-                    task,
-                    idx,
-                    calls[idx - prev_length] if calls is not None else None,
-                ):
+                wcall = calls[idx - prev_length] if calls is not None else None
+                if next_task := self.schedule_task(task, idx, wcall):
                     # if the parent task was retried,
                     # the next task might already be running
                     if fut := next(
@@ -314,6 +311,7 @@ class PregelRunner:
                                 __reraise_on_exit__=reraise,
                             ),
                         )
+                        fut.add_done_callback(partial(self.commit, next_task))
                         futures[fut] = next_task
                         rtn[idx - prev_length] = fut
             return [rtn.get(i) for i in range(len(writes))]
@@ -322,10 +320,15 @@ class PregelRunner:
             task: PregelExecutableTask,
             func: Callable[[Any], Union[Awaitable[Any], Any]],
             input: Any,
-        ) -> asyncio.Future[Any]:
+        ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]:
             (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)])
             assert fut is not None, "writer did not return a future for call"
-            return fut
+            if asyncio.iscoroutinefunction(func):
+                return fut
+            # adapted from asyncio.run_coroutine_threadsafe
+            sfut = concurrent.futures.Future()
+            loop.call_soon_threadsafe(chain_future, fut, sfut)
+            return sfut
 
         loop = asyncio.get_event_loop()
         tasks = tuple(tasks)
@@ -348,7 +351,7 @@ class PregelRunner:
                 )
                 self.commit(t, None)
             except Exception as exc:
-                self.commit(t, exc)
+                self.commit(t, None, exc)
                 if reraise and futures:
                     # will be re-raised after futures are done
                     fut: asyncio.Future = loop.create_future()
@@ -366,24 +369,24 @@ class PregelRunner:
         # yield updates/debug output as each task finishes
         for t in tasks:
             if not t.writes:
-                futures[
-                    cast(
-                        asyncio.Future,
-                        self.submit(
-                            arun_with_retry,
-                            t,
-                            retry_policy,
-                            stream=self.use_astream,
-                            configurable={
-                                CONFIG_KEY_SEND: partial(writer, t),
-                                CONFIG_KEY_CALL: partial(call, t),
-                            },
-                            __name__=t.name,
-                            __cancel_on_exit__=True,
-                            __reraise_on_exit__=reraise,
-                        ),
-                    )
-                ] = t
+                fut = cast(
+                    asyncio.Future,
+                    self.submit(
+                        arun_with_retry,
+                        t,
+                        retry_policy,
+                        stream=self.use_astream,
+                        configurable={
+                            CONFIG_KEY_SEND: partial(writer, t),
+                            CONFIG_KEY_CALL: partial(call, t),
+                        },
+                        __name__=t.name,
+                        __cancel_on_exit__=True,
+                        __reraise_on_exit__=reraise,
+                    ),
+                )
+                fut.add_done_callback(partial(self.commit, t))
+                futures[fut] = t
         end_time = timeout + loop.time() if timeout else None
         while len(futures) > (1 if get_waiter is not None else 0):
             done, inflight = await asyncio.wait(
@@ -402,8 +405,6 @@ class PregelRunner:
                 else:
                     # store for panic check
                     done_futures.add(fut)
-                    # task finished, commit writes
-                    self.commit(task, _exception(fut))
             else:
                 # remove references to loop vars
                 del fut, task
@@ -423,8 +424,13 @@ class PregelRunner:
         )
 
     def commit(
-        self, task: PregelExecutableTask, exception: Optional[BaseException]
+        self,
+        task: PregelExecutableTask,
+        fut: Union[None, concurrent.futures.Future[Any], asyncio.Future[Any]],
+        exception: Optional[BaseException] = None,
     ) -> None:
+        if fut is not None:
+            exception = _exception(fut)
         if exception:
             if isinstance(exception, GraphInterrupt):
                 # save interrupt to checkpointer

libs/langgraph/langgraph/types.py

@@ -1,4 +1,3 @@
-import asyncio
 import concurrent
 import concurrent.futures
 import dataclasses
@@ -7,7 +6,6 @@ from collections import deque
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
     ClassVar,
     Generic,
@@ -382,20 +380,4 @@ def call(
     conf = get_configurable()
     impl = conf[CONFIG_KEY_CALL]
     fut = impl(func, *args, **kwargs)
-    if not isinstance(fut, concurrent.futures.Future):
-        raise RuntimeError("In an async context, use acall() instead of call()")
-    return fut
-
-
-def acall(
-    func: str | Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs
-) -> asyncio.Future[T]:
-    from langgraph.constants import CONFIG_KEY_CALL
-    from langgraph.utils.config import get_configurable
-
-    conf = get_configurable()
-    impl = conf[CONFIG_KEY_CALL]
-    fut = impl(func, *args, **kwargs)
-    if isinstance(fut, concurrent.futures.Future):
-        fut = asyncio.wrap_future(fut)
     return fut

libs/langgraph/langgraph/utils/future.py

@@ -0,0 +1,121 @@
+import asyncio
+import concurrent.futures
+from typing import Union
+
+AnyFuture = Union[asyncio.Future, concurrent.futures.Future]
+
+
+def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop:
+    # Tries to call Future.get_loop() if it's available.
+    # Otherwise fallbacks to using the old '_loop' property.
+    try:
+        get_loop = fut.get_loop
+    except AttributeError:
+        pass
+    else:
+        return get_loop()
+    return fut._loop
+
+
+def _convert_future_exc(exc):
+    exc_class = type(exc)
+    if exc_class is concurrent.futures.CancelledError:
+        return asyncio.CancelledError(*exc.args)
+    elif exc_class is concurrent.futures.TimeoutError:
+        return asyncio.TimeoutError(*exc.args)
+    elif exc_class is concurrent.futures.InvalidStateError:
+        return asyncio.InvalidStateError(*exc.args)
+    else:
+        return exc
+
+
+def _set_concurrent_future_state(concurrent, source):
+    """Copy state from a future to a concurrent.futures.Future."""
+    assert source.done()
+    if source.cancelled():
+        concurrent.cancel()
+    if not concurrent.set_running_or_notify_cancel():
+        return
+    exception = source.exception()
+    if exception is not None:
+        concurrent.set_exception(_convert_future_exc(exception))
+    else:
+        result = source.result()
+        concurrent.set_result(result)
+
+
+def _copy_future_state(source, dest):
+    """Internal helper to copy state from another Future.
+
+    The other Future may be a concurrent.futures.Future.
+    """
+    assert source.done()
+    if dest.cancelled():
+        return
+    assert not dest.done()
+    if source.cancelled():
+        dest.cancel()
+    else:
+        exception = source.exception()
+        if exception is not None:
+            dest.set_exception(_convert_future_exc(exception))
+        else:
+            result = source.result()
+            dest.set_result(result)
+
+
+def _chain_future(source: AnyFuture, destination: AnyFuture) -> None:
+    """Chain two futures so that when one completes, so does the other.
+
+    The result (or exception) of source will be copied to destination.
+    If destination is cancelled, source gets cancelled too.
+    Compatible with both asyncio.Future and concurrent.futures.Future.
+    """
+    if not asyncio.isfuture(source) and not isinstance(
+        source, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for source argument")
+    if not asyncio.isfuture(destination) and not isinstance(
+        destination, concurrent.futures.Future
+    ):
+        raise TypeError("A future is required for destination argument")
+    source_loop = _get_loop(source) if asyncio.isfuture(source) else None
+    dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None
+
+    def _set_state(future, other):
+        if asyncio.isfuture(future):
+            _copy_future_state(other, future)
+        else:
+            _set_concurrent_future_state(future, other)
+
+    def _call_check_cancel(destination):
+        if destination.cancelled():
+            if source_loop is None or source_loop is dest_loop:
+                source.cancel()
+            else:
+                source_loop.call_soon_threadsafe(source.cancel)
+
+    def _call_set_state(source):
+        if destination.cancelled() and dest_loop is not None and dest_loop.is_closed():
+            return
+        if dest_loop is None or dest_loop is source_loop:
+            _set_state(destination, source)
+        else:
+            if dest_loop.is_closed():
+                return
+            dest_loop.call_soon_threadsafe(_set_state, destination, source)
+
+    destination.add_done_callback(_call_check_cancel)
+    source.add_done_callback(_call_set_state)
+
+
+def chain_future(source: AnyFuture, destination: AnyFuture) -> None:
+    # adapted from asyncio.run_coroutine_threadsafe
+    try:
+        _chain_future(source, destination)
+    except (SystemExit, KeyboardInterrupt):
+        raise
+    except BaseException as exc:
+        if destination.set_running_or_notify_cancel():
+            destination.set_exception(exc)
+        raise

libs/langgraph/tests/test_pregel.py

@@ -1976,21 +1976,21 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non
     mapper_calls = 0
 
     @task()
-    def mapper(input: str) -> str:
+    def mapper(input: int) -> str:
         nonlocal mapper_calls
         mapper_calls += 1
-        return input * 2
+        time.sleep(input / 100)
+        return str(input) * 2
 
     @imp(checkpointer=checkpointer)
-    def graph(input: list[str]) -> list[str]:
+    def graph(input: list[int]) -> list[str]:
         futures = [mapper(i) for i in input]
         mapped = [f.result() for f in futures]
         answer = interrupt("question")
         return [m + answer for m in mapped]
 
     thread1 = {"configurable": {"thread_id": "1"}}
-    assert [*graph.stream(["0", "1"], thread1)] == [
-        # TODO make test not depend on order of execution (which is not guaranteed)
+    assert [*graph.stream([0, 1], thread1)] == [
         {"mapper": "00"},
         {"mapper": "11"},
         {
@@ -2013,6 +2013,40 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non
     assert mapper_calls == 2
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_imp_stream_order(
+    request: pytest.FixtureRequest, checkpointer_name: str
+) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    @task()
+    def foo(state: dict) -> dict:
+        return {"a": state["a"] + "foo", "b": "bar"}
+
+    @task()
+    def bar(state: dict) -> dict:
+        return {"a": state["a"] + state["b"], "c": "bark"}
+
+    @task()
+    def baz(state: dict) -> dict:
+        return {"a": state["a"] + "baz", "c": "something else"}
+
+    @imp(checkpointer=checkpointer)
+    def graph(state: dict) -> dict:
+        fut_foo = foo(state)
+        fut_bar = bar(fut_foo.result())
+        fut_baz = baz(fut_bar.result())
+        return fut_baz.result()
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert [c for c in graph.stream({"a": "0"}, thread1)] == [
+        {"foo": {"a": "0foo", "b": "bar"}},
+        {"bar": {"a": "0foobar", "c": "bark"}},
+        {"baz": {"a": "0foobarbaz", "c": "something else"}},
+        {"graph": {"a": "0foobarbaz", "c": "something else"}},
+    ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_send_dedupe_on_resume(
     request: pytest.FixtureRequest, checkpointer_name: str

libs/langgraph/tests/test_pregel_async.py

@@ -2650,26 +2650,24 @@ async def test_send_sequences(checkpointer_name: str) -> None:
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_imp_task(checkpointer_name: str) -> None:
-    mapper_calls = 0
-
-    @task()
-    async def mapper(input: str) -> str:
-        nonlocal mapper_calls
-        mapper_calls += 1
-        return input * 2
-
     async with awith_checkpointer(checkpointer_name) as checkpointer:
+        mapper_calls = 0
+
+        @task()
+        async def mapper(input: int) -> str:
+            nonlocal mapper_calls
+            mapper_calls += 1
+            return str(input) * 2
 
         @imp(checkpointer=checkpointer)
-        async def graph(input: list[str]) -> list[str]:
+        async def graph(input: list[int]) -> list[str]:
             futures = [mapper(i) for i in input]
             mapped = await asyncio.gather(*futures)
             answer = interrupt("question")
             return [m + answer for m in mapped]
 
         thread1 = {"configurable": {"thread_id": "1"}}
-        assert [c async for c in graph.astream(["0", "1"], thread1)] == [
-            # TODO make test not depend on order of execution (which is not guaranteed)
+        assert [c async for c in graph.astream([0, 1], thread1)] == [
             {"mapper": "00"},
             {"mapper": "11"},
             {
@@ -2692,6 +2690,70 @@ async def test_imp_task(checkpointer_name: str) -> None:
         assert mapper_calls == 2
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_sync_from_async(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @task()
+        def foo(state: dict) -> dict:
+            return {"a": state["a"] + "foo", "b": "bar"}
+
+        @task()
+        def bar(state: dict) -> dict:
+            return {"a": state["a"] + state["b"], "c": "bark"}
+
+        @task()
+        def baz(state: dict) -> dict:
+            return {"a": state["a"] + "baz", "c": "something else"}
+
+        @imp(checkpointer=checkpointer)
+        def graph(state: dict) -> dict:
+            fut_foo = foo(state)
+            fut_bar = bar(fut_foo.result())
+            fut_baz = baz(fut_bar.result())
+            return fut_baz.result()
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
+            {"foo": {"a": "0foo", "b": "bar"}},
+            {"bar": {"a": "0foobar", "c": "bark"}},
+            {"baz": {"a": "0foobarbaz", "c": "something else"}},
+            {"graph": {"a": "0foobarbaz", "c": "something else"}},
+        ]
+
+
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
+async def test_imp_stream_order(checkpointer_name: str) -> None:
+    async with awith_checkpointer(checkpointer_name) as checkpointer:
+
+        @task()
+        async def foo(state: dict) -> dict:
+            return {"a": state["a"] + "foo", "b": "bar"}
+
+        @task()
+        async def bar(state: dict) -> dict:
+            return {"a": state["a"] + state["b"], "c": "bark"}
+
+        @task()
+        async def baz(state: dict) -> dict:
+            return {"a": state["a"] + "baz", "c": "something else"}
+
+        @imp(checkpointer=checkpointer)
+        async def graph(state: dict) -> dict:
+            fut_foo = foo(state)
+            fut_bar = bar(await fut_foo)
+            fut_baz = baz(await fut_bar)
+            return await fut_baz
+
+        thread1 = {"configurable": {"thread_id": "1"}}
+        assert [c async for c in graph.astream({"a": "0"}, thread1)] == [
+            {"foo": {"a": "0foo", "b": "bar"}},
+            {"bar": {"a": "0foobar", "c": "bark"}},
+            {"baz": {"a": "0foobarbaz", "c": "something else"}},
+            {"graph": {"a": "0foobarbaz", "c": "something else"}},
+        ]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
     if not FF_SEND_V2: