From 7bdbd62611ee8bb8572c3c30fc9ab37eb7bfbc87 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 Jan 2025 17:03:45 -0800 Subject: [PATCH 1/6] Fix tracing hierarchy for imperative api --- libs/langgraph/langgraph/func/__init__.py | 10 +++++----- libs/langgraph/langgraph/pregel/algo.py | 18 +++++++++++++----- libs/langgraph/langgraph/pregel/executor.py | 5 +++-- libs/langgraph/langgraph/pregel/runner.py | 12 ++++++++++-- libs/langgraph/langgraph/types.py | 5 +++-- libs/langgraph/langgraph/utils/config.py | 4 ++-- libs/langgraph/tests/test_pregel_async.py | 6 +++++- 7 files changed, 41 insertions(+), 19 deletions(-) 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", From 901273c0e28228f3c5d2280e1e266be105b4c5d6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 Jan 2025 17:09:30 -0800 Subject: [PATCH 2/6] Lint --- libs/langgraph/langgraph/pregel/algo.py | 3 +-- libs/langgraph/langgraph/pregel/executor.py | 5 ++++- libs/langgraph/langgraph/types.py | 3 +-- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 5be5e2311..c2d17d604 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,6 +1,5 @@ import sys from collections import defaultdict, deque -from contextvars import Context from functools import partial from hashlib import sha1 from typing import ( @@ -122,7 +121,7 @@ class Call: input: Any, *, retry: Optional[RetryPolicy], - callbacks: Optional[Context], + callbacks: Optional[Callbacks], ) -> None: self.func = func self.input = input diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 7e64a244a..80ce20e7c 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -65,7 +65,10 @@ class BackgroundExecutor(ContextManager): ) -> concurrent.futures.Future[T]: ctx = copy_context() if __next_tick__: - task = self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs) + task = cast( + concurrent.futures.Future[T], + self.executor.submit(next_tick, ctx.run, fn, *args, **kwargs), # type: ignore[arg-type] + ) else: task = self.executor.submit(ctx.run, fn, *args, **kwargs) self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index aa328bbb8..c0bd066ca 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -443,7 +443,6 @@ 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, @@ -456,7 +455,7 @@ def interrupt(value: Any) -> Any: from langgraph.errors import GraphInterrupt from langgraph.utils.config import get_config - conf = get_config()[CONF] + conf = get_config()["configurable"] # track interrupt index scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD] if "interrupt_counter" not in scratchpad: From 01331ef8586382b08fb5c82df44983438a67fee2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 08:33:09 -0800 Subject: [PATCH 3/6] Lint --- libs/langgraph/langgraph/pregel/algo.py | 4 ++-- libs/langgraph/langgraph/pregel/runner.py | 4 ++-- libs/langgraph/langgraph/utils/config.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c2d17d604..4031516c8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -113,7 +113,7 @@ class Call: func: Callable input: Any retry: Optional[RetryPolicy] - callbacks: Optional[Callbacks] + callbacks: Callbacks def __init__( self, @@ -121,7 +121,7 @@ class Call: input: Any, *, retry: Optional[RetryPolicy], - callbacks: Optional[Callbacks], + callbacks: Callbacks, ) -> None: self.func = func self.input = input diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index d4068aa5e..4a7114e46 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -150,7 +150,7 @@ class PregelRunner: input: Any, *, retry: Optional[RetryPolicy] = None, - callbacks: Optional[Callbacks] = None, + callbacks: Callbacks = None, ) -> concurrent.futures.Future[Any]: (fut,) = writer( task, @@ -342,7 +342,7 @@ class PregelRunner: input: Any, *, retry: Optional[RetryPolicy] = None, - callbacks: Optional[Callbacks] = None, + callbacks: Callbacks = None, ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: (fut,) = writer( task, diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index f87d12e89..372ea9616 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -132,7 +132,7 @@ def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig: def patch_config( config: Optional[RunnableConfig], *, - callbacks: Optional[Callbacks] = None, + callbacks: Callbacks = None, recursion_limit: Optional[int] = None, max_concurrency: Optional[int] = None, run_name: Optional[str] = None, From c767d86c9c1ee95190ab6d55d95d38adce6ed5c3 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 08:37:24 -0800 Subject: [PATCH 4/6] Remove ci job for removed flag --- .github/workflows/_test_langgraph.yml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index 2708d0f23..c9d1035b8 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -19,14 +19,9 @@ jobs: - "3.13" core-version: - "latest" - ff-send-v2: - - "false" include: - python-version: "3.11" core-version: ">=0.2.42,<0.3.0" - - python-version: "3.11" - core-version: "latest" - ff-send-v2: "true" defaults: run: @@ -57,8 +52,6 @@ jobs: - name: Run tests shell: bash - env: - LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }} run: | make test_parallel From 65a41942ef3718e6a2afc6736d62ef1a6592d263 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 08:39:09 -0800 Subject: [PATCH 5/6] Remove CI job to test against core 0.2.x --- .github/workflows/_test_langgraph.yml | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index c9d1035b8..ba1cbd08e 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -17,16 +17,11 @@ jobs: - "3.11" - "3.12" - "3.13" - core-version: - - "latest" - include: - - python-version: "3.11" - core-version: ">=0.2.42,<0.3.0" defaults: run: working-directory: libs/langgraph - name: "test #${{ matrix.python-version }} (langchain-core: ${{ matrix.core-version }}, ff-send-v2: ${{ matrix.ff-send-v2 }})" + name: "test #${{ matrix.python-version }}" steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} + Poetry ${{ env.POETRY_VERSION }} @@ -46,9 +41,6 @@ jobs: shell: bash run: | poetry install --with dev - if [ "${{ matrix.core-version }}" != "latest" ]; then - poetry run pip install "langchain-core${{ matrix.core-version }}" - fi - name: Run tests shell: bash From b7e4656c91d92efddef84683ab8415981c59ffa1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 08:45:35 -0800 Subject: [PATCH 6/6] Fix flaky test output order --- libs/langgraph/tests/test_pregel_async.py | 1 + 1 file changed, 1 insertion(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3945918b0..26352e545 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2456,6 +2456,7 @@ async def test_imp_task(checkpointer_name: str) -> None: async def mapper(input: int) -> str: nonlocal mapper_calls mapper_calls += 1 + await asyncio.sleep(0.1 * input) return str(input) * 2 @entrypoint(checkpointer=checkpointer)