diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 26d583ed2..6209896b4 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -19,7 +19,7 @@ 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 END, START, TAG_HIDDEN +from langgraph.constants import CONF, END, START, TAG_HIDDEN from langgraph.pregel import Pregel from langgraph.pregel.call import get_runnable_for_func from langgraph.pregel.read import PregelNode @@ -39,11 +39,11 @@ def call( retry: Optional[RetryPolicy] = None, ) -> concurrent.futures.Future[T]: from langgraph.constants import CONFIG_KEY_CALL - from langgraph.utils.config import get_configurable + from langgraph.utils.config import get_config - conf = get_configurable() - impl = conf[CONFIG_KEY_CALL] - fut = impl(func, input, retry=retry) + config = get_config() + impl = config[CONF][CONFIG_KEY_CALL] + fut = impl(func, input, retry=retry, callbacks=config["callbacks"]) return fut diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 130d7e983..5be5e2311 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,5 +1,6 @@ import sys from collections import defaultdict, deque +from contextvars import Context from functools import partial from hashlib import sha1 from typing import ( @@ -19,6 +20,7 @@ from typing import ( ) from uuid import UUID +from langchain_core.callbacks import Callbacks from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager from langchain_core.runnables.config import RunnableConfig @@ -107,18 +109,25 @@ class PregelTaskWrites(NamedTuple): class Call: - __slots__ = ("func", "input", "retry") + __slots__ = ("func", "input", "retry", "callbacks") func: Callable input: Any retry: Optional[RetryPolicy] + callbacks: Optional[Callbacks] def __init__( - self, func: Callable, input: Any, *, retry: Optional[RetryPolicy] + self, + func: Callable, + input: Any, + *, + retry: Optional[RetryPolicy], + callbacks: Optional[Context], ) -> None: self.func = func self.input = input self.retry = retry + self.callbacks = callbacks def should_interrupt( @@ -465,9 +474,8 @@ def prepare_single_task( patch_config( merge_configs(config, {"metadata": metadata}), run_name=name, - callbacks=( - manager.get_child(f"graph:step:{step}") if manager else None - ), + callbacks=call.callbacks + or (manager.get_child(f"graph:step:{step}") if manager else None), configurable={ CONFIG_KEY_TASK_ID: task_id, # deque.extend is thread-safe diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 46a4e6036..7e64a244a 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -63,10 +63,11 @@ class BackgroundExecutor(ContextManager): __next_tick__: bool = False, **kwargs: P.kwargs, ) -> concurrent.futures.Future[T]: + ctx = copy_context() if __next_tick__: - task = self.executor.submit(next_tick, fn, *args, **kwargs) + task = self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs) else: - task = self.executor.submit(fn, *args, **kwargs) + task = self.executor.submit(ctx.run, fn, *args, **kwargs) self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__) task.add_done_callback(self.done) return task diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index e680518a5..d4068aa5e 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -17,6 +17,8 @@ from typing import ( cast, ) +from langchain_core.callbacks import Callbacks + from langgraph.constants import ( CONF, CONFIG_KEY_CALL, @@ -148,9 +150,12 @@ class PregelRunner: input: Any, *, retry: Optional[RetryPolicy] = None, + callbacks: Optional[Callbacks] = None, ) -> concurrent.futures.Future[Any]: (fut,) = writer( - task, [(PUSH, None)], calls=[Call(func, input, retry=retry)] + task, + [(PUSH, None)], + calls=[Call(func, input, retry=retry, callbacks=callbacks)], ) assert fut is not None, "writer did not return a future for call" return fut @@ -337,9 +342,12 @@ class PregelRunner: input: Any, *, retry: Optional[RetryPolicy] = None, + callbacks: Optional[Callbacks] = None, ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: (fut,) = writer( - task, [(PUSH, None)], calls=[Call(func, input, retry=retry)] + task, + [(PUSH, None)], + calls=[Call(func, input, retry=retry, callbacks=callbacks)], ) assert fut is not None, "writer did not return a future for call" if asyncio.iscoroutinefunction(func): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 53dc6bd57..aa328bbb8 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -443,6 +443,7 @@ def interrupt(value: Any) -> Any: GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client. """ from langgraph.constants import ( + CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, @@ -453,9 +454,9 @@ def interrupt(value: Any) -> Any: RESUME, ) from langgraph.errors import GraphInterrupt - from langgraph.utils.config import get_configurable + from langgraph.utils.config import get_config - conf = get_configurable() + conf = get_config()[CONF] # track interrupt index scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD] if "interrupt_counter" not in scratchpad: diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index 5bff9e848..f87d12e89 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -297,7 +297,7 @@ def ensure_config(*configs: Optional[RunnableConfig]) -> RunnableConfig: return empty -def get_configurable() -> dict[str, Any]: +def get_config() -> RunnableConfig: if sys.version_info < (3, 11): try: if asyncio.current_task(): @@ -307,6 +307,6 @@ def get_configurable() -> dict[str, Any]: except RuntimeError: pass if var_config := var_child_runnable_config.get(): - return var_config[CONF] + return var_config else: raise RuntimeError("Called get_configurable outside of a runnable context") diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 347e31574..329b22c28 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2465,7 +2465,8 @@ async def test_imp_task(checkpointer_name: str) -> None: answer = interrupt("question") return [m + answer for m in mapped] - thread1 = {"configurable": {"thread_id": "1"}} + tracer = FakeTracer() + thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} assert [c async for c in graph.astream([0, 1], thread1)] == [ {"mapper": "00"}, {"mapper": "11"}, @@ -2481,6 +2482,9 @@ 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" assert await graph.ainvoke(Command(resume="answer"), thread1) == [ "00answer",