From 64b828def6a4fe696d654dd2a7fc36cd7d9128ca Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Mon, 27 Jan 2025 12:05:39 -0500
Subject: [PATCH] langgraph: handle task naming for reused class methods
 (#3216)

---
 libs/langgraph/langgraph/func/__init__.py | 6 +++++-
 libs/langgraph/tests/test_pregel.py       | 3 +++
 libs/langgraph/tests/test_pregel_async.py | 3 +++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py
index cf6341387..8e392c3ea 100644
--- a/libs/langgraph/langgraph/func/__init__.py
+++ b/libs/langgraph/langgraph/func/__init__.py
@@ -123,7 +123,11 @@ def task(
         if name is not None:
             if hasattr(func, "__func__"):
                 # handle class methods
-                func.__func__.__name__ = name
+                # NOTE: we're modifying the instance method to avoid modifying
+                # the original class method in case it's shared across multiple tasks
+                instance_method = functools.partial(func.__func__, func.__self__)  # type: ignore [union-attr]
+                instance_method.__name__ = name  # type: ignore [attr-defined]
+                func = instance_method
             else:
                 # handle regular functions / partials / callable classes, etc.
                 func.__name__ = name
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 53a67ddc6..5716c65fb 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -6277,6 +6277,7 @@ def test_named_tasks_functional() -> None:
 
     # class method task
     foo = task(f.foo, name="custom_foo")
+    other_foo = task(f.foo, name="other_foo")
 
     # regular function task
     @task(name="custom_bar")
@@ -6300,6 +6301,7 @@ def test_named_tasks_functional() -> None:
     @entrypoint()
     def workflow(inputs: dict) -> dict:
         fut_foo = foo(inputs)
+        other_foo(inputs)
         fut_bar = bar(fut_foo.result())
         fut_baz = baz_task(fut_bar.result())
         fut_custom_baz = custom_baz_task(fut_baz.result())
@@ -6308,6 +6310,7 @@ def test_named_tasks_functional() -> None:
 
     assert list(workflow.stream("", stream_mode="updates")) == [
         {"custom_foo": "foo"},
+        {"other_foo": "foo"},
         {"custom_bar": "foo|bar"},
         {"baz": "foo|bar|baz"},
         {"custom_baz": "foo|bar|baz|custom_baz"},
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 192d31c11..24e53345f 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -7338,6 +7338,7 @@ async def test_named_tasks_functional() -> None:
 
     # class method task
     foo = task(f.foo, name="custom_foo")
+    other_foo = task(f.foo, name="other_foo")
 
     # regular function task
     @task(name="custom_bar")
@@ -7361,6 +7362,7 @@ async def test_named_tasks_functional() -> None:
     @entrypoint()
     async def workflow(inputs: dict) -> dict:
         foo_result = await foo(inputs)
+        await other_foo(inputs)
         bar_result = await bar(foo_result)
         baz_result = await baz_task(bar_result)
         custom_baz_result = await custom_baz_task(baz_result)
@@ -7369,6 +7371,7 @@ async def test_named_tasks_functional() -> None:
 
     assert [c async for c in workflow.astream("", stream_mode="updates")] == [
         {"custom_foo": "foo"},
+        {"other_foo": "foo"},
         {"custom_bar": "foo|bar"},
         {"baz": "foo|bar|baz"},
         {"custom_baz": "foo|bar|baz|custom_baz"},