langgraph: add names for tasks

libs/langgraph/langgraph/func/__init__.py

@@ -1,4 +1,4 @@
 import asyncio
 import concurrent.futures
 import functools
 import inspect
@@ -37,7 +37,7 @@
 
 @overload
 def task(
-    *, retry: Optional[RetryPolicy] = None
+    *, name: Optional[str] = None, retry: Optional[RetryPolicy] = None
 ) -> Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]]: ...
 
 
@@ -50,6 +50,7 @@
 def task(
     __func_or_none__: Optional[Union[Callable[P, T], Callable[P, Awaitable[T]]]] = None,
     *,
+    name: Optional[str] = None,
     retry: Optional[RetryPolicy] = None,
 ) -> Union[
     Callable[[Callable[P, T]], Callable[P, SyncAsyncFuture[T]]],
@@ -120,6 +121,9 @@
     ) -> Union[
         Callable[P, concurrent.futures.Future[T]], Callable[P, asyncio.Future[T]]
     ]:
+        if name is not None:
+            func.__name__ = name
+
         call_func = functools.partial(call, func, retry=retry)
         object.__setattr__(call_func, "_is_pregel_task", True)
         return functools.update_wrapper(call_func, func)

libs/langgraph/tests/test_pregel.py

@@ -5570,10 +5570,10 @@ def graph(state: dict) -> dict:
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
-def test_multiple_interrupts_imperative(
+def test_multiple_interrupts_functional(
     request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
 ):
-    """Test multiple interrupts with an imperative API."""
+    """Test multiple interrupts with an functional API."""
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
 
     counter = 0
@@ -6265,3 +6265,25 @@ async def workflow(inputs: dict, *, previous: Any):
     # Test with another thread
     assert await workflow.ainvoke({}, {"configurable": {"thread_id": "2"}}) == "!"
     assert previous_ is None
+
+
+def test_named_tasks_functional() -> None:
+    @task(name="custom_name")
+    def foo(state: dict) -> dict:
+        return "foo"
+
+    @task
+    def bar(state: dict) -> dict:
+        return "bar"
+
+    @entrypoint()
+    def workflow(inputs: dict) -> dict:
+        fut_foo = foo(inputs)
+        fut_bar = bar(fut_foo.result())
+        return fut_bar.result()
+
+    assert list(workflow.stream({}, stream_mode="updates")) == [
+        {"custom_name": "foo"},
+        {"bar": "bar"},
+        {"workflow": "bar"},
+    ]

libs/langgraph/tests/test_pregel_async.py

@@ -6812,8 +6812,8 @@ async def graph(state: dict) -> dict:
 
 @NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
-async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
-    """Test multiple interrupts with an imperative API."""
+async def test_multiple_interrupts_functional(checkpointer_name: str) -> None:
+    """Test multiple interrupts with an functional API."""
     from langgraph.func import entrypoint, task
 
     counter = 0
@@ -7325,3 +7325,25 @@ async def foo(inputs, previous=None) -> Any:
 
     assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"]
     assert previous_return_values == [None]
+
+
+async def test_named_tasks_functional() -> None:
+    @task(name="custom_name")
+    async def foo(state: dict) -> dict:
+        return "foo"
+
+    @task
+    async def bar(state: dict) -> dict:
+        return "bar"
+
+    @entrypoint()
+    async def workflow(inputs: dict) -> dict:
+        foo_result = await foo(inputs)
+        bar_result = await bar(foo_result)
+        return bar_result
+
+    assert [c async for c in workflow.astream({}, stream_mode="updates")] == [
+        {"custom_name": "foo"},
+        {"bar": "bar"},
+        {"workflow": "bar"},
+    ]