From 56fd3458cc2456b8cc59856f70ed2f1bdaccc380 Mon Sep 17 00:00:00 2001 From: Elior Nataf Lackritz Date: Tue, 6 Oct 2026 18:29:03 -0400 Subject: [PATCH] refactor(langgraph): check for put_writes' task_path in one place Both loop setups and bulk_update_state inspected the saver's signature themselves; they now share put_writes_accepts_task_path. --- libs/langgraph/langgraph/pregel/_checkpoint.py | 10 ++++++++++ libs/langgraph/langgraph/pregel/_loop.py | 8 +++----- libs/langgraph/langgraph/pregel/main.py | 5 +++-- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 2f86d60a0..359191a1c 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -3,6 +3,7 @@ from __future__ import annotations import uuid from collections.abc import Callable, Iterable, Mapping from datetime import datetime, timezone +from inspect import signature from typing import Any, Literal, cast from langchain_core.runnables import RunnableConfig @@ -49,6 +50,15 @@ def empty_checkpoint() -> Checkpoint: ) +def put_writes_accepts_task_path(put_writes: Callable[..., Any]) -> bool: + """Whether a saver's `put_writes` or `aput_writes` takes `task_path`. + + Savers written before the parameter existed don't, so it is passed only + when this is true. + """ + return signature(put_writes).parameters.get("task_path") is not None + + def exit_delta_task_id(step: int, task_id: str) -> str: """Synthetic task id for exit-mode DeltaChannel writes. diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 180591c6f..61b4280de 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -12,7 +12,6 @@ from contextlib import ( ExitStack, ) from datetime import datetime, timezone -from inspect import signature from types import TracebackType from typing import ( Any, @@ -107,6 +106,7 @@ from langgraph.pregel._checkpoint import ( delta_channels_with_pending_writes, empty_checkpoint, exit_delta_task_id, + put_writes_accepts_task_path, ) from langgraph.pregel._executor import ( AsyncBackgroundExecutor, @@ -1584,8 +1584,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.put_writes self.checkpointer_put_writes_accepts_task_path = ( - signature(checkpointer.put_writes).parameters.get("task_path") - is not None + put_writes_accepts_task_path(checkpointer.put_writes) ) else: self.checkpointer_get_next_version = increment @@ -1840,8 +1839,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.aput_writes self.checkpointer_put_writes_accepts_task_path = ( - signature(checkpointer.aput_writes).parameters.get("task_path") - is not None + put_writes_accepts_task_path(checkpointer.aput_writes) ) else: self.checkpointer_get_next_version = increment diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 16232d39a..56439ec0f 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -19,7 +19,7 @@ from collections.abc import ( from dataclasses import is_dataclass, replace from datetime import timedelta from functools import partial -from inspect import isclass, signature +from inspect import isclass from typing import ( Any, Generic, @@ -140,6 +140,7 @@ from langgraph.pregel._checkpoint import ( delta_channels_with_pending_writes, empty_checkpoint, get_updated_channels_from_tasks, + put_writes_accepts_task_path, versions_seen_without_bumps, ) from langgraph.pregel._draw import draw_graph @@ -4250,7 +4251,7 @@ def _task_path_kwarg(put_writes: Callable[..., Any], task: PregelTaskWrites) -> Savers that replay a checkpoint's writes in task path order then give back updates applied together in the order they were given. """ - if signature(put_writes).parameters.get("task_path") is None: + if not put_writes_accepts_task_path(put_writes): return {} return {"task_path": task_path_str(task.path)}