Fix tracing of args for @task decorated functions (#3107)

- now using same logic as in langsmith sdk, treating as single args
dict, based on function signature

libs/langgraph/langgraph/func/__init__.py

@@ -1,5 +1,4 @@
 import asyncio
-import concurrent
 import concurrent.futures
 import functools
 import inspect
@@ -10,7 +9,6 @@ from typing import (
     Awaitable,
     Callable,
     Optional,
-    TypeVar,
     Union,
     overload,
 )
@@ -18,39 +16,19 @@ from typing import (
 from langchain_core.runnables.base import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.runnables.graph import Graph, Node
-from typing_extensions import ParamSpec
 
 from langgraph.channels.ephemeral_value import EphemeralValue
 from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langgraph.constants import CONF, END, START, TAG_HIDDEN
+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.call import P, T, call, get_runnable_for_entrypoint
 from langgraph.pregel.protocol import PregelProtocol
 from langgraph.pregel.read import PregelNode
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.store.base import BaseStore
 from langgraph.types import RetryPolicy, StreamMode, StreamWriter
 
-P = ParamSpec("P")
-P1 = TypeVar("P1")
-T = TypeVar("T")
-
-
-def call(
-    func: Callable[P, T],
-    *args: Any,
-    retry: Optional[RetryPolicy] = None,
-    **kwargs: Any,
-) -> concurrent.futures.Future[T]:
-    from langgraph.constants import CONFIG_KEY_CALL
-    from langgraph.utils.config import get_config
-
-    config = get_config()
-    impl = config[CONF][CONFIG_KEY_CALL]
-    fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"])
-    return fut
-
 
 @overload
 def task(
@@ -149,22 +127,12 @@ def task(
 
     def decorator(
         func: Union[Callable[P, Awaitable[T]], Callable[P, T]],
-    ) -> Callable[P, concurrent.futures.Future[T]]:
-        if asyncio.iscoroutinefunction(func):
-
-            @functools.wraps(func)
-            async def _tick(__allargs__: tuple) -> T:
-                return await func(*__allargs__[0], **__allargs__[1])
-
-        else:
-
-            @functools.wraps(func)
-            def _tick(__allargs__: tuple) -> T:
-                return func(*__allargs__[0], **__allargs__[1])
-
-        wrapper = functools.partial(call, _tick, retry=retry)
-        object.__setattr__(wrapper, "_is_pregel_task", True)
-        return functools.update_wrapper(wrapper, func)
+    ) -> Union[
+        Callable[P, concurrent.futures.Future[T]], Callable[P, asyncio.Future[T]]
+    ]:
+        call_func = functools.partial(call, func, retry=retry)
+        object.__setattr__(call_func, "_is_pregel_task", True)
+        return functools.update_wrapper(call_func, func)
 
     if __func_or_none__ is not None:
         return decorator(__func_or_none__)
@@ -347,7 +315,8 @@ def entrypoint(
                 new_sig = original_sig.replace(parameters=new_params)
                 # Update the signature of the wrapper function
                 gen_wrapper.__signature__ = new_sig  # type: ignore
-            bound = get_runnable_for_func(gen_wrapper)
+
+            bound = get_runnable_for_entrypoint(gen_wrapper)
             stream_mode: StreamMode = "custom"
         elif inspect.isasyncgenfunction(func):
             original_sig = inspect.signature(func)
@@ -390,10 +359,10 @@ def entrypoint(
                 # Update the signature of the wrapper function
                 agen_wrapper.__signature__ = new_sig  # type: ignore
 
-            bound = get_runnable_for_func(agen_wrapper)
+            bound = get_runnable_for_entrypoint(agen_wrapper)
             stream_mode = "custom"
         else:
-            bound = get_runnable_for_func(func)
+            bound = get_runnable_for_entrypoint(func)
             stream_mode = "updates"
 
         # get input and output types

libs/langgraph/langgraph/pregel/algo.py

@@ -62,7 +62,7 @@ from langgraph.constants import (
 )
 from langgraph.errors import EmptyChannelError, InvalidUpdateError
 from langgraph.managed.base import ManagedValueMapping
-from langgraph.pregel.call import get_runnable_for_func
+from langgraph.pregel.call import get_runnable_for_task
 from langgraph.pregel.io import read_channel, read_channels
 from langgraph.pregel.log import logger
 from langgraph.pregel.manager import ChannelsManager
@@ -439,7 +439,7 @@ def prepare_single_task(
         # (PUSH, parent task path, idx of PUSH write, id of parent task, Call)
         task_path_t = cast(tuple[str, tuple, int, str, Call], task_path)
         call = task_path_t[-1]
-        proc_ = get_runnable_for_func(call.func)
+        proc_ = get_runnable_for_task(call.func)
         name = proc_.name
         if name is None:
             raise ValueError("`call` functions must have a `__name__` attribute")

libs/langgraph/langgraph/pregel/call.py

@@ -1,12 +1,25 @@
 """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
+from typing import Any, Callable, Optional, TypeVar, Union
 
-from langgraph.constants import RETURN
+from typing_extensions import ParamSpec
+
+from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN, TAG_HIDDEN
 from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
-from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable
+from langgraph.types import RetryPolicy
+from langgraph.utils.config import get_config
+from langgraph.utils.runnable import (
+    RunnableCallable,
+    RunnableSeq,
+    is_async_callable,
+    run_in_executor,
+)
 
 ##
 # Utilities borrowed from cloudpickle.
@@ -107,18 +120,96 @@ def _lookup_module_and_qualname(
     return module, name
 
 
-def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq:
-    if func in CACHE:
-        return CACHE[func]
+def _explode_args_trace_inputs(
+    sig: inspect.Signature, input: tuple[tuple[Any, ...], dict[str, Any]]
+) -> dict[str, Any]:
+    args, kwargs = input
+    bound = sig.bind_partial(*args, **kwargs)
+    bound.apply_defaults()
+    arguments = dict(bound.arguments)
+    arguments.pop("self", None)
+    arguments.pop("cls", None)
+    for param_name, param in sig.parameters.items():
+        if param.kind == inspect.Parameter.VAR_KEYWORD:
+            # Update with the **kwargs, and remove the original entry
+            # This is to help flatten out keyword arguments
+            if param_name in arguments:
+                arguments.update(arguments.pop(param_name))
+    return arguments
+
+
+def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq:
+    key = (func, False)
+    if key in CACHE:
+        return CACHE[key]
     else:
+        if is_async_callable(func):
+            run = RunnableCallable(None, func, name=func.__name__, trace=False)
+        else:
+            afunc = functools.update_wrapper(
+                functools.partial(run_in_executor, None, func), func
+            )
+            run = RunnableCallable(
+                func,
+                afunc,
+                name=func.__name__,
+                trace=False,
+            )
         seq = RunnableSeq(
-            coerce_to_runnable(func, name=None, trace=False),
-            ChannelWrite([ChannelWriteEntry(RETURN)]),
+            run,
+            ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]),
             name=func.__name__,
         )
         if not _lookup_module_and_qualname(func):
             return seq
-        return CACHE.setdefault(func, seq)
+        return CACHE.setdefault(key, seq)
 
 
-CACHE: dict[Callable[..., Any], RunnableSeq] = {}
+def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq:
+    key = (func, True)
+    if key in CACHE:
+        return CACHE[key]
+    else:
+        if is_async_callable(func):
+            run = RunnableCallable(
+                None, func, explode_args=True, name=func.__name__, trace=False
+            )
+        else:
+            run = RunnableCallable(
+                func,
+                functools.wraps(func)(functools.partial(run_in_executor, None, func)),
+                explode_args=True,
+                name=func.__name__,
+                trace=False,
+            )
+        seq = RunnableSeq(
+            run,
+            ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]),
+            name=func.__name__,
+            trace_inputs=functools.partial(
+                _explode_args_trace_inputs, inspect.signature(func)
+            ),
+        )
+        if not _lookup_module_and_qualname(func):
+            return seq
+        return CACHE.setdefault(key, seq)
+
+
+CACHE: dict[tuple[Callable[..., Any], bool], RunnableSeq] = {}
+
+
+P = ParamSpec("P")
+P1 = TypeVar("P1")
+T = TypeVar("T")
+
+
+def call(
+    func: Callable[P, T],
+    *args: Any,
+    retry: Optional[RetryPolicy] = None,
+    **kwargs: Any,
+) -> Union[concurrent.futures.Future[T], asyncio.Future[T]]:
+    config = get_config()
+    impl = config[CONF][CONFIG_KEY_CALL]
+    fut = impl(func, (args, kwargs), retry=retry, callbacks=config["callbacks"])
+    return fut

libs/langgraph/langgraph/utils/runnable.py

@@ -120,6 +120,7 @@ class RunnableCallable(Runnable):
         tags: Optional[Sequence[str]] = None,
         trace: bool = True,
         recurse: bool = True,
+        explode_args: bool = False,
         **kwargs: Any,
     ) -> None:
         self.name = name
@@ -141,6 +142,7 @@ class RunnableCallable(Runnable):
         self.kwargs = kwargs
         self.trace = trace
         self.recurse = recurse
+        self.explode_args = explode_args
         # check signature
         if func is None and afunc is None:
             raise ValueError("At least one of func or afunc must be provided.")
@@ -176,7 +178,12 @@ class RunnableCallable(Runnable):
             )
         if config is None:
             config = ensure_config()
-        kwargs = {**self.kwargs, **kwargs}
+        if self.explode_args:
+            args, _kwargs = input
+            kwargs = {**self.kwargs, **_kwargs, **kwargs}
+        else:
+            args = (input,)
+            kwargs = {**self.kwargs, **kwargs}
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
@@ -208,7 +215,7 @@ class RunnableCallable(Runnable):
                 child_config = patch_config(config, callbacks=run_manager.get_child())
                 context = copy_context()
                 context.run(_set_config_context, child_config)
-                ret = context.run(self.func, input, **kwargs)
+                ret = context.run(self.func, *args, **kwargs)
             except BaseException as e:
                 run_manager.on_chain_error(e)
                 raise
@@ -216,7 +223,7 @@ class RunnableCallable(Runnable):
                 run_manager.on_chain_end(ret)
         else:
             context.run(_set_config_context, config)
-            ret = context.run(self.func, input, **kwargs)
+            ret = context.run(self.func, *args, **kwargs)
         if isinstance(ret, Runnable) and self.recurse:
             return ret.invoke(input, config)
         return ret
@@ -228,7 +235,12 @@ class RunnableCallable(Runnable):
             return self.invoke(input, config)
         if config is None:
             config = ensure_config()
-        kwargs = {**self.kwargs, **kwargs}
+        if self.explode_args:
+            args, _kwargs = input
+            kwargs = {**self.kwargs, **_kwargs, **kwargs}
+        else:
+            args = (input,)
+            kwargs = {**self.kwargs, **kwargs}
         if self.func_accepts_config:
             kwargs["config"] = config
         _conf = config[CONF]
@@ -258,7 +270,7 @@ class RunnableCallable(Runnable):
             try:
                 child_config = patch_config(config, callbacks=run_manager.get_child())
                 context.run(_set_config_context, child_config)
-                coro = cast(Coroutine[None, None, Any], self.afunc(input, **kwargs))
+                coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
                 if ASYNCIO_ACCEPTS_CONTEXT:
                     ret = await asyncio.create_task(coro, context=context)
                 else:
@@ -271,10 +283,10 @@ class RunnableCallable(Runnable):
         else:
             context.run(_set_config_context, config)
             if ASYNCIO_ACCEPTS_CONTEXT:
-                coro = cast(Coroutine[None, None, Any], self.afunc(input, **kwargs))
+                coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs))
                 ret = await asyncio.create_task(coro, context=context)
             else:
-                ret = await self.afunc(input, **kwargs)
+                ret = await self.afunc(*args, **kwargs)
         if isinstance(ret, Runnable) and self.recurse:
             return await ret.ainvoke(input, config)
         return ret
@@ -347,6 +359,7 @@ class RunnableSeq(Runnable):
         self,
         *steps: RunnableLike,
         name: Optional[str] = None,
+        trace_inputs: Optional[Callable[[Any], Any]] = None,
     ) -> None:
         """Create a new RunnableSeq.
 
@@ -371,6 +384,7 @@ class RunnableSeq(Runnable):
             )
         self.steps = steps_flat
         self.name = name
+        self.trace_inputs = trace_inputs
 
     def __or__(
         self,
@@ -432,7 +446,7 @@ class RunnableSeq(Runnable):
         # start the root run
         run_manager = callback_manager.on_chain_start(
             None,
-            input,
+            self.trace_inputs(input) if self.trace_inputs is not None else input,
             name=config.get("run_name") or self.get_name(),
             run_id=config.pop("run_id", None),
         )
@@ -442,7 +456,7 @@ class RunnableSeq(Runnable):
             for i, step in enumerate(self.steps):
                 # mark each step as a child run
                 config = patch_config(
-                    config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
+                    config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
                 )
                 if i == 0:
                     input = step.invoke(input, config, **kwargs)
@@ -469,7 +483,7 @@ class RunnableSeq(Runnable):
         # start the root run
         run_manager = await callback_manager.on_chain_start(
             None,
-            input,
+            self.trace_inputs(input) if self.trace_inputs is not None else input,
             name=config.get("run_name") or self.get_name(),
             run_id=config.pop("run_id", None),
         )
@@ -479,7 +493,7 @@ class RunnableSeq(Runnable):
             for i, step in enumerate(self.steps):
                 # mark each step as a child run
                 config = patch_config(
-                    config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
+                    config, callbacks=run_manager.get_child(f"seq:step:{i + 1}")
                 )
                 if i == 0:
                     input = await step.ainvoke(input, config, **kwargs)
@@ -506,7 +520,7 @@ class RunnableSeq(Runnable):
         # start the root run
         run_manager = callback_manager.on_chain_start(
             None,
-            input,
+            self.trace_inputs(input) if self.trace_inputs is not None else input,
             name=config.get("run_name") or self.get_name(),
             run_id=config.pop("run_id", None),
         )
@@ -519,7 +533,7 @@ class RunnableSeq(Runnable):
             for idx, step in enumerate(self.steps):
                 config = patch_config(
                     config,
-                    callbacks=run_manager.get_child(f"seq:step:{idx+1}"),
+                    callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
                 )
                 if idx == 0:
                     iterator = step.stream(input, config, **kwargs)
@@ -569,7 +583,7 @@ class RunnableSeq(Runnable):
         # start the root run
         run_manager = await callback_manager.on_chain_start(
             None,
-            input,
+            self.trace_inputs(input) if self.trace_inputs is not None else input,
             name=config.get("run_name") or self.get_name(),
             run_id=config.pop("run_id", None),
         )
@@ -583,7 +597,7 @@ class RunnableSeq(Runnable):
                 for idx, step in enumerate(self.steps):
                     config = patch_config(
                         config,
-                        callbacks=run_manager.get_child(f"seq:step:{idx+1}"),
+                        callbacks=run_manager.get_child(f"seq:step:{idx + 1}"),
                     )
                     if idx == 0:
                         aiterator = step.astream(input, config, **kwargs)

libs/langgraph/tests/test_pregel_async.py

@@ -2475,7 +2475,12 @@ async def test_imp_task(checkpointer_name: str) -> None:
         assert mapper_calls == 2
         assert len(tracer.runs) == 1
         assert len(tracer.runs[0].child_runs) == 1
-        assert tracer.runs[0].child_runs[0].name == "graph"
+        entrypoint_run = tracer.runs[0].child_runs[0]
+        assert entrypoint_run.name == "graph"
+        mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"]
+        assert len(mapper_runs) == 2
+        assert any(r.inputs == {"input": 0} for r in mapper_runs)
+        assert any(r.inputs == {"input": 1} for r in mapper_runs)
 
         assert await graph.ainvoke(Command(resume="answer"), thread1) == [
             "00answer",