diff --git a/libs/langgraph/langgraph/pregel/_call.py b/libs/langgraph/langgraph/pregel/_call.py index 8a07a9387..863c4d345 100644 --- a/libs/langgraph/langgraph/pregel/_call.py +++ b/libs/langgraph/langgraph/pregel/_call.py @@ -21,6 +21,7 @@ from langgraph._internal._runnable import ( run_in_executor, ) from langgraph.config import get_config +from langgraph.pregel._read import _validate_traceable_config from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry from langgraph.types import CachePolicy, RetryPolicy @@ -221,20 +222,15 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable: # Check for traceable config early to handle __unwrapped__ raw_config = getattr(func, "__traceable_config__", None) - traceable_config = None - if raw_config and isinstance(raw_config, dict): - traceable_config = { - "process_inputs": raw_config.get("process_inputs"), - "process_outputs": raw_config.get("process_outputs"), - "enabled": raw_config.get("enabled"), - "__unwrapped__": raw_config.get("__unwrapped__"), - } + traceable_config = _validate_traceable_config(raw_config) # Use unwrapped function if available to avoid double-tracing # when @traceable and @task are used together - func_to_run = func - if traceable_config and traceable_config.get("__unwrapped__"): - func_to_run = traceable_config["__unwrapped__"] + func_to_run: Callable[..., Any] = func + if traceable_config: + unwrapped = traceable_config.get("__unwrapped__") + if unwrapped is not None: + func_to_run = unwrapped if is_async_callable(func_to_run): run = RunnableCallable( @@ -262,10 +258,11 @@ def get_runnable_for_task(func: Callable[..., Any]) -> Runnable: base_trace_inputs = functools.partial( _explode_args_trace_inputs, inspect.signature(func) ) - if traceable_config and traceable_config.get("process_inputs"): - trace_inputs = _compose_trace_inputs( - base_trace_inputs, traceable_config["process_inputs"] - ) + process_inputs = ( + traceable_config.get("process_inputs") if traceable_config else None + ) + if process_inputs is not None: + trace_inputs = _compose_trace_inputs(base_trace_inputs, process_inputs) else: trace_inputs = base_trace_inputs diff --git a/libs/langgraph/langgraph/pregel/_read.py b/libs/langgraph/langgraph/pregel/_read.py index 9effaf391..74ae08811 100644 --- a/libs/langgraph/langgraph/pregel/_read.py +++ b/libs/langgraph/langgraph/pregel/_read.py @@ -41,11 +41,13 @@ def _validate_traceable_config(raw: Any) -> TraceableConfig | None: """ if not isinstance(raw, dict): return None + # Support both "__unwrapped__" (langsmith) and "wrapped" (legacy) keys + unwrapped = raw.get("__unwrapped__") or raw.get("wrapped") return { "process_inputs": raw.get("process_inputs"), "process_outputs": raw.get("process_outputs"), "enabled": raw.get("enabled"), # None means use external context - "wrapped": raw.get("wrapped"), + "__unwrapped__": unwrapped, } diff --git a/libs/langgraph/tests/test_traceable_integration.py b/libs/langgraph/tests/test_traceable_integration.py index c6207ccbd..e20ab74ec 100644 --- a/libs/langgraph/tests/test_traceable_integration.py +++ b/libs/langgraph/tests/test_traceable_integration.py @@ -52,7 +52,7 @@ def _set_traceable_config( "process_inputs": process_inputs, "process_outputs": process_outputs, "enabled": enabled, - "wrapped": func, + "__unwrapped__": func, }, ) return func