From 204c9c83f82078a3a6375f3a7054aa4e68e64f01 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 20 Jan 2025 10:42:29 -0800 Subject: [PATCH 1/5] Fix tracing of args for @task decorated functions - now using same logic as in langsmith sdk, treating as single args dict, based on function signature --- libs/langgraph/langgraph/func/__init__.py | 55 +++-------- libs/langgraph/langgraph/pregel/algo.py | 4 +- libs/langgraph/langgraph/pregel/call.py | 109 +++++++++++++++++++-- libs/langgraph/langgraph/utils/runnable.py | 44 ++++++--- 4 files changed, 143 insertions(+), 69 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 6bdc8a1ef..1ca426e8f 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index d1247dbc4..38db8ff67 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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") diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index d6b6e6d82..a7bbc1767 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/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 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), + run, ChannelWrite([ChannelWriteEntry(RETURN)]), 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)), # type: ignore[arg-type] + explode_args=True, + name=func.__name__, + trace=False, + ) + seq = RunnableSeq( + run, + ChannelWrite([ChannelWriteEntry(RETURN)]), + 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 diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 19a70e4cf..fbdb2502b 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/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) From 07695f5c5a8835f9226f66c53b7c4ebeea2ba255 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 20 Jan 2025 10:53:37 -0800 Subject: [PATCH 2/5] Lint --- libs/langgraph/langgraph/pregel/call.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index a7bbc1767..f8c6b22f3 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -177,7 +177,7 @@ def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: else: run = RunnableCallable( func, - functools.wraps(func)(functools.partial(run_in_executor, None, func)), # type: ignore[arg-type] + functools.wraps(func)(functools.partial(run_in_executor, None, func)), explode_args=True, name=func.__name__, trace=False, From b7c3ac4501e1a24cec06ad76deed7d1b8e19527e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 20 Jan 2025 13:55:13 -0800 Subject: [PATCH 3/5] Fix tracing output --- libs/langgraph/langgraph/pregel/call.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index f8c6b22f3..22feb46bb 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -10,7 +10,7 @@ from typing import Any, Callable, Optional, TypeVar, Union from typing_extensions import ParamSpec -from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN +from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN, TAG_HIDDEN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import RetryPolicy from langgraph.utils.config import get_config @@ -157,7 +157,7 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: ) seq = RunnableSeq( run, - ChannelWrite([ChannelWriteEntry(RETURN)]), + ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]), name=func.__name__, ) if not _lookup_module_and_qualname(func): @@ -184,7 +184,7 @@ def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: ) seq = RunnableSeq( run, - ChannelWrite([ChannelWriteEntry(RETURN)]), + ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]), name=func.__name__, trace_inputs=functools.partial( _explode_args_trace_inputs, inspect.signature(func) From 16c86c9de6acb6806bb93262bf907f0f18b1db11 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 21 Jan 2025 09:49:46 -0800 Subject: [PATCH 4/5] Add test --- libs/langgraph/tests/test_pregel_async.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 17847f4b8..e937b1949 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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 mapper_runs[0].inputs == {"input": 0} + assert mapper_runs[1].inputs == {"input": 1} assert await graph.ainvoke(Command(resume="answer"), thread1) == [ "00answer", From fe46576d98d5889e4d4663814785c7cfe8d0b996 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 21 Jan 2025 09:59:49 -0800 Subject: [PATCH 5/5] Fix --- libs/langgraph/tests/test_pregel_async.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index e937b1949..dab6285f4 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2479,8 +2479,8 @@ async def test_imp_task(checkpointer_name: str) -> None: 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 mapper_runs[0].inputs == {"input": 0} - assert mapper_runs[1].inputs == {"input": 1} + 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",