From 01a3c23a2901d06854d2cf9171452a174c0ce417 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Mon, 11 Nov 2024 10:05:34 -0800
Subject: [PATCH] WIP

---
 libs/langgraph/langgraph/func/__init__.py | 85 +++++++++++++++++++++++
 libs/langgraph/tests/test_pregel.py       | 24 +++++++
 2 files changed, 109 insertions(+)
 create mode 100644 libs/langgraph/langgraph/func/__init__.py

diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
new file mode 100644
index 000000000..66acba23c
--- /dev/null
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -0,0 +1,85 @@
+import asyncio
+import concurrent
+import concurrent.futures
+import types
+from functools import partial, update_wrapper
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Optional,
+    ParamSpec,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from langgraph.channels.ephemeral_value import EphemeralValue
+from langgraph.channels.last_value import LastValue
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.constants import END, START, TAG_HIDDEN
+from langgraph.pregel import Pregel
+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
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+@overload
+def task(
+    *, retry: Optional[RetryPolicy] = None
+) -> Callable[
+    [Callable[P, Coroutine[None, None, T]]], Callable[P, asyncio.Future[T]]
+]: ...
+
+
+@overload
+def task(
+    *, retry: Optional[RetryPolicy] = None
+) -> Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]]: ...
+
+
+def task(
+    *, retry: Optional[RetryPolicy] = None
+) -> Callable[
+    [Callable[P, Union[T, Awaitable[T]]]],
+    Callable[P, Union[concurrent.futures.Future[T], asyncio.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 _task
+
+
+def imp(
+    *,
+    checkpointer: Optional[BaseCheckpointSaver] = None,
+    store: Optional[BaseStore] = None,
+) -> Callable[[types.FunctionType], Pregel]:
+    def _imp(func: types.FunctionType):
+        return Pregel(
+            nodes={
+                func.__name__: PregelNode(
+                    bound=get_runnable_for_func(func),
+                    triggers=[START],
+                    channels=[START],
+                    writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])],
+                )
+            },
+            channels={START: EphemeralValue(Any, START), END: LastValue(Any, END)},
+            input_channels=START,
+            output_channels=END,
+            stream_mode="updates",
+            checkpointer=checkpointer,
+            store=store,
+        )
+
+    return _imp
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 120b3deac..356ecbadf 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -65,6 +65,7 @@ from langgraph.constants import (
     START,
 )
 from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
+from langgraph.func import imp, task
 from langgraph.graph import END, Graph, StateGraph
 from langgraph.graph.message import MessageGraph, MessagesState, add_messages
 from langgraph.managed.shared_value import SharedValue
@@ -1969,6 +1970,29 @@ def test_send_sequences() -> None:
     )
 
 
+@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
+    checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
+
+    @task()
+    def mapper(input: str) -> str:
+        print(f"mapper {input}")
+        return input * 2
+
+    @imp(checkpointer=checkpointer)
+    def graph(input: list[str]) -> list[str]:
+        futures = [mapper(i) for i in input]
+        mapped = [f.result() for f in futures]
+        # answer = interrupt("question")
+        # TODO raises NodeInterrupt if no answer provided yet
+        # returns answer (saved in writes?) if provided
+        # what is the API for passing the answer?
+        return mapped
+
+    thread1 = {"configurable": {"thread_id": "1"}}
+    assert graph.invoke(["0", "1"], thread1) == ["00", "11"]
+
+
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_send_dedupe_on_resume(
     request: pytest.FixtureRequest, checkpointer_name: str