mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-13 03:35:21 +02:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6b84ae69a0 | ||
|
|
48b2346e33 |
@@ -14,8 +14,8 @@ async def memory_checkpointer():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_memory():
|
||||
"""InMemorySaver passes the tests of every capability it implements."""
|
||||
async def test_validate_memory_base():
|
||||
"""InMemorySaver passes all base capability tests."""
|
||||
report = await validate(memory_checkpointer)
|
||||
report.print_report()
|
||||
assert report.passed_all(), f"Capability tests failed: {report.to_dict()}"
|
||||
assert report.passed_all_base(), f"Base tests failed: {report.to_dict()}"
|
||||
|
||||
@@ -451,8 +451,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
* Stage 1 (paged): dynamic SELECT over `checkpoints` with three
|
||||
columns per channel: its version, an `EXISTS` probe for a stored
|
||||
blob at that version, and its inline value. Pages newest-first by
|
||||
`checkpoint_id`, starting at the target; page size is
|
||||
`_DELTA_PAGE_SIZE`.
|
||||
`checkpoint_id` with a cursor; page size is `_DELTA_PAGE_SIZE`.
|
||||
Stops paging when every channel has found its seed or a page comes
|
||||
back short.
|
||||
|
||||
@@ -469,7 +468,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if not checkpoint_id:
|
||||
if checkpoint_id is None:
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
@@ -477,7 +476,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
|
||||
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
|
||||
# Python after each page. Stops as soon as every channel has its seed.
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
@@ -487,7 +486,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = checkpoint_id
|
||||
cursor: str | None = None
|
||||
|
||||
with self._cursor() as cur:
|
||||
while True:
|
||||
@@ -496,7 +495,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
cur.execute(stage1_sql, stage1_params)
|
||||
page = cur.fetchall()
|
||||
@@ -528,7 +527,6 @@ class PostgresSaver(BasePostgresSaver):
|
||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
||||
break
|
||||
cursor = oldest
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
|
||||
# Stage 2: per-channel UNION ALL — one writes branch per channel
|
||||
# with non-empty chain, plus one blob branch per seeded channel.
|
||||
|
||||
@@ -417,13 +417,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if not checkpoint_id:
|
||||
if checkpoint_id is None:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
@@ -433,7 +433,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = checkpoint_id
|
||||
cursor: str | None = None
|
||||
|
||||
async with self._cursor() as cur:
|
||||
while True:
|
||||
@@ -442,7 +442,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
await cur.execute(stage1_sql, stage1_params)
|
||||
page = await cur.fetchall()
|
||||
@@ -472,7 +472,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
||||
break
|
||||
cursor = oldest
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
|
||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
||||
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
||||
|
||||
@@ -178,13 +178,10 @@ class _DeltaStage2Row(TypedDict, total=False):
|
||||
# `_build_delta_stage2_sql` document their shapes.
|
||||
|
||||
|
||||
def _build_delta_stage1_sql(
|
||||
channels: Sequence[str], *, paged: bool, include_cursor: bool = False
|
||||
) -> str:
|
||||
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
"""Build stage 1 SQL with K parallel version lookups + seed probes.
|
||||
|
||||
For channels=["messages", "files"] (with `paged=True, include_cursor=True`)
|
||||
the result is::
|
||||
For channels=["messages", "files"] (with `paged=True`) the result is::
|
||||
|
||||
SELECT checkpoint_id, parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
||||
@@ -200,7 +197,7 @@ def _build_delta_stage1_sql(
|
||||
checkpoint -> 'channel_values' -> %s AS inline_1
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
AND checkpoint_id <= %s
|
||||
AND (%s::text IS NULL OR checkpoint_id < %s)
|
||||
ORDER BY checkpoint_id DESC
|
||||
LIMIT %s
|
||||
|
||||
@@ -241,16 +238,9 @@ def _build_delta_stage1_sql(
|
||||
and uses safe identifiers).
|
||||
|
||||
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
|
||||
cursor, page_size]` when `paged=True` — four per channel: the version
|
||||
lookup, the blob's channel, the version the blob must match, and the inline
|
||||
lookup.
|
||||
|
||||
Pages run newest-first from the target down. The first page passes the
|
||||
target as the cursor with `include_cursor=True`, so it opens with the
|
||||
target's own row, whose parent starts the walk; each later page continues
|
||||
below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6
|
||||
is time-ordered, which `get_tuple` also relies on to find the latest
|
||||
checkpoint), so no row newer than the target is part of its chain.
|
||||
cursor, cursor, page_size]` when `paged=True` — four per channel: the
|
||||
version lookup, the blob's channel, the version the blob must match, and the
|
||||
inline lookup.
|
||||
|
||||
When `paged=False`, the WHERE has no cursor predicate and there's no
|
||||
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
||||
@@ -274,7 +264,7 @@ def _build_delta_stage1_sql(
|
||||
)
|
||||
if paged:
|
||||
sql += (
|
||||
f" AND checkpoint_id {'<=' if include_cursor else '<'} %s"
|
||||
" AND (%s::text IS NULL OR checkpoint_id < %s)"
|
||||
" ORDER BY checkpoint_id DESC LIMIT %s"
|
||||
)
|
||||
return sql
|
||||
@@ -423,8 +413,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
materialized at this point),
|
||||
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
||||
a later page; the cursor stays put), or
|
||||
(d) the target's own row isn't in `parent_of` (the target doesn't
|
||||
exist, so the walk never starts).
|
||||
(d) the target's own row isn't in `parent_of` yet (the walk has
|
||||
not started; no cursor is set, so a later page retries).
|
||||
|
||||
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
||||
`walk_cursor_by_ch`, and `seeded` in place.
|
||||
@@ -432,9 +422,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
for i, ch in enumerate(channels):
|
||||
if ch in seeded:
|
||||
continue
|
||||
# The first page opens with the target's row, so it's missing only
|
||||
# when the target doesn't exist; a `None` cursor would read as
|
||||
# "target is a root".
|
||||
# Pages start at the thread head, so the target may not have
|
||||
# loaded yet; a `None` cursor would read as "target is a root".
|
||||
if ch not in walk_cursor_by_ch:
|
||||
if target_id not in parent_of:
|
||||
continue
|
||||
|
||||
@@ -36,7 +36,6 @@ from langgraph.store.base import (
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
validate_op_namespace,
|
||||
)
|
||||
from psycopg import Capabilities, Connection, Cursor, Pipeline
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
@@ -1387,7 +1386,6 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
|
||||
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
|
||||
tot = 0
|
||||
for idx, op in enumerate(ops):
|
||||
validate_op_namespace(op)
|
||||
grouped_ops[type(op)].append((idx, op))
|
||||
tot += 1
|
||||
return grouped_ops, tot
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.3"
|
||||
version = "3.1.2"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -13,10 +13,8 @@ import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
@@ -873,65 +871,3 @@ async def test_omit_expired_search_pagination(store: AsyncPostgresStore) -> None
|
||||
page2 = await store.asearch(ns, limit=2, offset=2)
|
||||
assert [i.key for i in page1] == ["a", "b"]
|
||||
assert [i.key for i in page2] == ["c"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
||||
async def test_abatch_rejects_invalid_namespace_labels(
|
||||
store: AsyncPostgresStore, namespace: tuple
|
||||
) -> None:
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
for op in (
|
||||
GetOp(namespace, "key"),
|
||||
GetOp(namespace, "key", refresh_ttl=True),
|
||||
PutOp(namespace, "key", {"changed": True}),
|
||||
PutOp(namespace, "key", None),
|
||||
SearchOp(namespace),
|
||||
ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await store.abatch([op])
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
|
||||
async def test_invalid_namespace_only_fails_its_own_call(
|
||||
store: AsyncPostgresStore,
|
||||
) -> None:
|
||||
"""Concurrent calls share one `abatch`, which fails every op if it raises.
|
||||
|
||||
Labels are checked before an op is queued, so one caller's bad label cannot
|
||||
fail another caller's request.
|
||||
"""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
valid, invalid = await asyncio.gather(
|
||||
store.aget(("foo", "bar"), "key"),
|
||||
store.aget(("foo.bar",), "key"),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert isinstance(valid, Item) and valid.value == {"original": True}
|
||||
assert isinstance(invalid, InvalidNamespaceError)
|
||||
|
||||
|
||||
async def test_sync_methods_reject_invalid_namespace_labels(
|
||||
store: AsyncPostgresStore,
|
||||
) -> None:
|
||||
"""The sync wrappers run off the event loop thread and must validate too."""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
for call in (
|
||||
lambda: store.get(("foo.bar",), "key"),
|
||||
lambda: store.search(("foo.bar",)),
|
||||
lambda: store.delete(("foo.bar",), "key"),
|
||||
lambda: store.list_namespaces(prefix=("foo.bar",)),
|
||||
lambda: store.batch([GetOp(("foo.bar",), "key")]),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await asyncio.to_thread(call)
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
@@ -15,7 +14,7 @@ from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE, BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
CHANNEL = "items"
|
||||
@@ -24,8 +23,8 @@ SEED_STEP = 1
|
||||
SEED_VALUE = [10, 20]
|
||||
TARGET_STEP = 4
|
||||
|
||||
# The real page size is the control; the rest split the walk from the target
|
||||
# to its seed across pages.
|
||||
# The real page size is the control; the rest leave the target off the first
|
||||
# page (three checkpoints are newer than it).
|
||||
PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1]
|
||||
|
||||
|
||||
@@ -93,7 +92,7 @@ def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None:
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
async def test_async_walk_continues_across_pages(
|
||||
async def test_async_target_older_than_the_first_page(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size)
|
||||
@@ -107,7 +106,7 @@ async def test_async_walk_continues_across_pages(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
def test_sync_walk_continues_across_pages(
|
||||
def test_sync_target_older_than_the_first_page(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size)
|
||||
@@ -120,84 +119,6 @@ def test_sync_walk_continues_across_pages(
|
||||
_assert_history(result[CHANNEL], page_size)
|
||||
|
||||
|
||||
def _record_rows_read(monkeypatch: pytest.MonkeyPatch) -> list[str]:
|
||||
read: list[str] = []
|
||||
ingest = BasePostgresSaver._ingest_stage1_page
|
||||
|
||||
def record(rows: Sequence[Mapping[str, Any]], *args: Any) -> str | None:
|
||||
read.extend(row["checkpoint_id"] for row in rows)
|
||||
return ingest(rows, *args)
|
||||
|
||||
monkeypatch.setattr(BasePostgresSaver, "_ingest_stage1_page", staticmethod(record))
|
||||
return read
|
||||
|
||||
|
||||
def _ids_from_target_down(configs: list[dict]) -> list[str]:
|
||||
return [c["configurable"]["checkpoint_id"] for c in configs[TARGET_STEP::-1]]
|
||||
|
||||
|
||||
async def test_async_walk_reads_nothing_newer_than_the_target(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
read = _record_rows_read(monkeypatch)
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], _DELTA_PAGE_SIZE)
|
||||
assert read == _ids_from_target_down(configs)
|
||||
|
||||
|
||||
def test_sync_walk_reads_nothing_newer_than_the_target(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
read = _record_rows_read(monkeypatch)
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
configs = _build_chain(saver)
|
||||
result = saver.get_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], _DELTA_PAGE_SIZE)
|
||||
assert read == _ids_from_target_down(configs)
|
||||
|
||||
|
||||
def _empty_checkpoint_id(config: dict) -> dict:
|
||||
return {"configurable": {**config["configurable"], "checkpoint_id": ""}}
|
||||
|
||||
|
||||
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
latest = await saver.aget_delta_channel_history(
|
||||
config=configs[-1], channels=[CHANNEL]
|
||||
)
|
||||
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert latest[CHANNEL]["writes"]
|
||||
assert result == latest
|
||||
|
||||
|
||||
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
configs = _build_chain(saver)
|
||||
latest = saver.get_delta_channel_history(config=configs[-1], channels=[CHANNEL])
|
||||
|
||||
result = saver.get_delta_channel_history(
|
||||
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert latest[CHANNEL]["writes"]
|
||||
assert result == latest
|
||||
|
||||
|
||||
async def test_root_target_has_no_history_and_still_terminates(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
@@ -11,7 +11,6 @@ import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
@@ -1165,51 +1164,3 @@ def test_namespace_labels_with_trailing_newline(store) -> None:
|
||||
assert set(store.list_namespaces(prefix=["users", "alice"], limit=100)) == {
|
||||
("users", "alice"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
||||
@pytest.mark.parametrize(
|
||||
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
|
||||
)
|
||||
def test_batch_rejects_invalid_namespace_labels(
|
||||
store, namespace: tuple, kind: str
|
||||
) -> None:
|
||||
"""Ops passed straight to `batch` must not reach another namespace.
|
||||
|
||||
Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same
|
||||
text as `("foo", "bar")`. `BaseStore` methods validate labels themselves,
|
||||
but `batch` takes ops as given.
|
||||
"""
|
||||
op = {
|
||||
"get": GetOp(namespace, "key"),
|
||||
"put": PutOp(namespace, "key", {"changed": True}),
|
||||
"delete": PutOp(namespace, "key", None),
|
||||
"search": SearchOp(namespace),
|
||||
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
}[kind]
|
||||
store.put(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
store.batch([PutOp(("valid",), "key", {}), op])
|
||||
|
||||
item = store.get(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
# The whole batch is rejected before any SQL runs.
|
||||
assert store.get(("valid",), "key") is None
|
||||
|
||||
|
||||
def test_batch_allows_empty_search_prefix_and_listing_wildcards(
|
||||
store,
|
||||
) -> None:
|
||||
store.put(("foo", "bar"), "key", {"v": 1})
|
||||
|
||||
found, listed = store.batch(
|
||||
[
|
||||
SearchOp(()),
|
||||
ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)),
|
||||
]
|
||||
)
|
||||
|
||||
assert [item.namespace for item in found] == [("foo", "bar")]
|
||||
assert listed == [("foo", "bar")]
|
||||
|
||||
Generated
+1
-1
@@ -353,7 +353,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.3"
|
||||
version = "3.1.2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
|
||||
@@ -539,7 +539,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if not checkpoint_id:
|
||||
if checkpoint_id is None:
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -653,7 +653,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if not checkpoint_id:
|
||||
if checkpoint_id is None:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -28,7 +28,6 @@ from langgraph.store.base import (
|
||||
ensure_embeddings,
|
||||
get_text_at_path,
|
||||
tokenize_path,
|
||||
validate_op_namespace,
|
||||
)
|
||||
|
||||
_AIO_ERROR_MSG = (
|
||||
@@ -258,7 +257,6 @@ def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int
|
||||
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
|
||||
tot = 0
|
||||
for idx, op in enumerate(ops):
|
||||
validate_op_namespace(op)
|
||||
grouped_ops[type(op)].append((idx, op))
|
||||
tot += 1
|
||||
return grouped_ops, tot
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.1"
|
||||
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -9,10 +9,8 @@ from typing import cast
|
||||
import pytest
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
PutOp,
|
||||
SearchOp,
|
||||
)
|
||||
@@ -747,65 +745,3 @@ async def test_async_namespace_segment_boundary(store: AsyncSqliteStore) -> None
|
||||
assert set(await store.alist_namespaces(suffix=["alice"], limit=100)) == {
|
||||
("uid", "users", "alice"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
||||
async def test_abatch_rejects_invalid_namespace_labels(
|
||||
store: AsyncSqliteStore, namespace: tuple
|
||||
) -> None:
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
for op in (
|
||||
GetOp(namespace, "key"),
|
||||
GetOp(namespace, "key", refresh_ttl=True),
|
||||
PutOp(namespace, "key", {"changed": True}),
|
||||
PutOp(namespace, "key", None),
|
||||
SearchOp(namespace),
|
||||
ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await store.abatch([op])
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
|
||||
async def test_invalid_namespace_only_fails_its_own_call(
|
||||
store: AsyncSqliteStore,
|
||||
) -> None:
|
||||
"""Concurrent calls share one `abatch`, which fails every op if it raises.
|
||||
|
||||
Labels are checked before an op is queued, so one caller's bad label cannot
|
||||
fail another caller's request.
|
||||
"""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
valid, invalid = await asyncio.gather(
|
||||
store.aget(("foo", "bar"), "key"),
|
||||
store.aget(("foo.bar",), "key"),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert isinstance(valid, Item) and valid.value == {"original": True}
|
||||
assert isinstance(invalid, InvalidNamespaceError)
|
||||
|
||||
|
||||
async def test_sync_methods_reject_invalid_namespace_labels(
|
||||
store: AsyncSqliteStore,
|
||||
) -> None:
|
||||
"""The sync wrappers run off the event loop thread and must validate too."""
|
||||
await store.aput(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
for call in (
|
||||
lambda: store.get(("foo.bar",), "key"),
|
||||
lambda: store.search(("foo.bar",)),
|
||||
lambda: store.delete(("foo.bar",), "key"),
|
||||
lambda: store.list_namespaces(prefix=("foo.bar",)),
|
||||
lambda: store.batch([GetOp(("foo.bar",), "key")]),
|
||||
):
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await asyncio.to_thread(call)
|
||||
|
||||
item = await store.aget(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
|
||||
@@ -65,39 +65,6 @@ async def test_async_walk_reaches_parent_whatever_the_id_order(
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
EMPTY_CHECKPOINT_ID: dict[str, Any] = {
|
||||
"configurable": {**CONFIG["configurable"], "checkpoint_id": ""}
|
||||
}
|
||||
|
||||
|
||||
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = saver.put(CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {})
|
||||
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
saver.put(root, _checkpoint("z-newer", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(
|
||||
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = await saver.aput(
|
||||
CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {}
|
||||
)
|
||||
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
await saver.aput(root, _checkpoint("z-newer", {}), {}, {})
|
||||
|
||||
got = await saver.aget_delta_channel_history(
|
||||
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_walk_reaches_root_of_long_chain_with_descending_ids() -> None:
|
||||
steps = 40
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
|
||||
@@ -14,7 +14,6 @@ import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
@@ -1436,51 +1435,3 @@ def test_list_namespaces_metacharacter_labels(store: SqliteStore) -> None:
|
||||
assert set(store.list_namespaces(prefix=[label, "child"], limit=100)) == {
|
||||
(label, "child"),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [("foo.bar",), ("foo", ""), ("foo", 1)])
|
||||
@pytest.mark.parametrize(
|
||||
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
|
||||
)
|
||||
def test_batch_rejects_invalid_namespace_labels(
|
||||
store: SqliteStore, namespace: tuple, kind: str
|
||||
) -> None:
|
||||
"""Ops passed straight to `batch` must not reach another namespace.
|
||||
|
||||
Namespaces are stored dot-joined, so `("foo.bar",)` flattens to the same
|
||||
text as `("foo", "bar")`. `BaseStore` methods validate labels themselves,
|
||||
but `batch` takes ops as given.
|
||||
"""
|
||||
op = {
|
||||
"get": GetOp(namespace, "key"),
|
||||
"put": PutOp(namespace, "key", {"changed": True}),
|
||||
"delete": PutOp(namespace, "key", None),
|
||||
"search": SearchOp(namespace),
|
||||
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
}[kind]
|
||||
store.put(("foo", "bar"), "key", {"original": True})
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
store.batch([PutOp(("valid",), "key", {}), op])
|
||||
|
||||
item = store.get(("foo", "bar"), "key")
|
||||
assert item is not None and item.value == {"original": True}
|
||||
# The whole batch is rejected before any SQL runs.
|
||||
assert store.get(("valid",), "key") is None
|
||||
|
||||
|
||||
def test_batch_allows_empty_search_prefix_and_listing_wildcards(
|
||||
store: SqliteStore,
|
||||
) -> None:
|
||||
store.put(("foo", "bar"), "key", {"v": 1})
|
||||
|
||||
found, listed = store.batch(
|
||||
[
|
||||
SearchOp(()),
|
||||
ListNamespacesOp((MatchCondition("prefix", ("foo", "*")),)),
|
||||
]
|
||||
)
|
||||
|
||||
assert [item.namespace for item in found] == [("foo", "bar")]
|
||||
assert listed == [("foo", "bar")]
|
||||
|
||||
Generated
+1
-1
@@ -360,7 +360,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -170,8 +170,8 @@ class InMemorySaver(
|
||||
return {}
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"].get("checkpoint_id", "")
|
||||
ns_storage = self.storage.get(thread_id, {}).get(checkpoint_ns, {})
|
||||
checkpoint_id = get_checkpoint_id(config) or max(ns_storage, default="")
|
||||
|
||||
chain: list[str] = []
|
||||
target_entry = ns_storage.get(checkpoint_id)
|
||||
|
||||
@@ -55,7 +55,6 @@ logger = logging.getLogger(__name__)
|
||||
_MAX_WARNED_TYPES = 1000
|
||||
_warned_unregistered_types: set[tuple[str, str]] = set()
|
||||
_warned_blocked_types: set[tuple[str, str]] = set()
|
||||
_warned_unreconstructable_types: set[tuple[str, str]] = set()
|
||||
|
||||
|
||||
def _is_safe_json_type(id_list: list[str]) -> bool:
|
||||
@@ -80,27 +79,6 @@ def _warn_once(
|
||||
logger.warning(msg, *args)
|
||||
|
||||
|
||||
def _reconstruction_fallback(tup: Any, exc: Exception) -> Any:
|
||||
"""Return the serialized payload of an object that could not be rebuilt.
|
||||
|
||||
Returning `None` here would silently erase the value from restored state.
|
||||
"""
|
||||
try:
|
||||
module, name, payload = tup[0], tup[1], tup[2]
|
||||
except Exception:
|
||||
return None
|
||||
_warn_once(
|
||||
_warned_unreconstructable_types,
|
||||
(str(module), str(name)),
|
||||
"Could not reconstruct %s.%s from checkpoint (%s); "
|
||||
"returning its serialized data instead.",
|
||||
module,
|
||||
name,
|
||||
type(exc).__name__,
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
class JsonPlusSerializer(SerializerProtocol):
|
||||
"""Serializer that uses ormsgpack, with optional fallbacks.
|
||||
|
||||
@@ -660,7 +638,6 @@ def _create_msgpack_ext_hook(
|
||||
)
|
||||
)
|
||||
elif code == EXT_CONSTRUCTOR_SINGLE_ARG:
|
||||
tup = None
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
@@ -672,10 +649,9 @@ def _create_msgpack_ext_hook(
|
||||
return tup[2]
|
||||
# module, name, arg
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(tup[2])
|
||||
except Exception as exc:
|
||||
return _reconstruction_fallback(tup, exc)
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_CONSTRUCTOR_POS_ARGS:
|
||||
tup = None
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
@@ -686,10 +662,9 @@ def _create_msgpack_ext_hook(
|
||||
return _send_from_args(tup[2])
|
||||
# module, name, args
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(*tup[2])
|
||||
except Exception as exc:
|
||||
return _reconstruction_fallback(tup, exc)
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_CONSTRUCTOR_KW_ARGS:
|
||||
tup = None
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
@@ -698,10 +673,9 @@ def _create_msgpack_ext_hook(
|
||||
return tup[2]
|
||||
# module, name, kwargs
|
||||
return getattr(importlib.import_module(tup[0]), tup[1])(**tup[2])
|
||||
except Exception as exc:
|
||||
return _reconstruction_fallback(tup, exc)
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_METHOD_SINGLE_ARG:
|
||||
tup = None
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
data, ext_hook=ext_hook, option=ormsgpack.OPT_NON_STR_KEYS
|
||||
@@ -712,8 +686,8 @@ def _create_msgpack_ext_hook(
|
||||
return getattr(
|
||||
getattr(importlib.import_module(tup[0]), tup[1]), tup[3]
|
||||
)(tup[2])
|
||||
except Exception as exc:
|
||||
return _reconstruction_fallback(tup, exc)
|
||||
except Exception:
|
||||
return None
|
||||
elif code == EXT_PYDANTIC_V1:
|
||||
try:
|
||||
tup = ormsgpack.unpackb(
|
||||
|
||||
@@ -771,12 +771,7 @@ class BaseStore(ABC):
|
||||
|
||||
Returns:
|
||||
The retrieved item or `None` if not found.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a namespace label is empty, is not a string,
|
||||
or contains a period (`.`).
|
||||
"""
|
||||
_validate_namespace_labels(namespace)
|
||||
return self.batch(
|
||||
[GetOp(namespace, str(key), _ensure_refresh(self.ttl_config, refresh_ttl))]
|
||||
)[0]
|
||||
@@ -806,10 +801,6 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a
|
||||
string, or contains a period (`.`).
|
||||
|
||||
???+ example "Examples"
|
||||
|
||||
Basic filtering:
|
||||
@@ -849,7 +840,6 @@ class BaseStore(ABC):
|
||||
Natural language search support depends on your store implementation
|
||||
and requires proper embedding configuration.
|
||||
"""
|
||||
_validate_namespace_labels(namespace_prefix)
|
||||
return self.batch(
|
||||
[
|
||||
SearchOp(
|
||||
@@ -897,11 +887,6 @@ class BaseStore(ABC):
|
||||
By default, the expiration timer refreshes on both read operations (get/search)
|
||||
and write operations (put/update), whenever the item is included in the operation.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If the namespace is empty, its root label is
|
||||
`"langgraph"`, or a label is empty, is not a string, or contains a
|
||||
period (`.`).
|
||||
|
||||
Note:
|
||||
Indexing support depends on your store implementation.
|
||||
If you do not initialize the store with indexing capabilities,
|
||||
@@ -955,12 +940,7 @@ class BaseStore(ABC):
|
||||
Args:
|
||||
namespace: Hierarchical path for the item.
|
||||
key: Unique identifier within the namespace.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a namespace label is empty, is not a string,
|
||||
or contains a period (`.`).
|
||||
"""
|
||||
_validate_namespace_labels(namespace)
|
||||
self.batch([PutOp(namespace, str(key), None, ttl=None)])
|
||||
|
||||
def list_namespaces(
|
||||
@@ -989,10 +969,6 @@ class BaseStore(ABC):
|
||||
A list of namespace tuples that match the criteria. Each tuple represents a
|
||||
full namespace path up to `max_depth`.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a
|
||||
string, or contains a period (`.`).
|
||||
|
||||
???+ example "Examples":
|
||||
|
||||
Setting `max_depth=3`. Given the namespaces:
|
||||
@@ -1008,8 +984,6 @@ class BaseStore(ABC):
|
||||
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
_validate_namespace_labels(prefix or ())
|
||||
_validate_namespace_labels(suffix or ())
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
@@ -1039,12 +1013,7 @@ class BaseStore(ABC):
|
||||
|
||||
Returns:
|
||||
The retrieved item or `None` if not found.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a namespace label is empty, is not a string,
|
||||
or contains a period (`.`).
|
||||
"""
|
||||
_validate_namespace_labels(namespace)
|
||||
return (
|
||||
await self.abatch(
|
||||
[
|
||||
@@ -1083,10 +1052,6 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a `namespace_prefix` label is empty, is not a
|
||||
string, or contains a period (`.`).
|
||||
|
||||
???+ example "Examples"
|
||||
|
||||
Basic filtering:
|
||||
@@ -1126,7 +1091,6 @@ class BaseStore(ABC):
|
||||
Natural language search support depends on your store implementation
|
||||
and requires proper embedding configuration.
|
||||
"""
|
||||
_validate_namespace_labels(namespace_prefix)
|
||||
return (
|
||||
await self.abatch(
|
||||
[
|
||||
@@ -1176,11 +1140,6 @@ class BaseStore(ABC):
|
||||
By default, the expiration timer refreshes on both read operations (get/search)
|
||||
and write operations (put/update), whenever the item is included in the operation.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If the namespace is empty, its root label is
|
||||
`"langgraph"`, or a label is empty, is not a string, or contains a
|
||||
period (`.`).
|
||||
|
||||
Note:
|
||||
Indexing support depends on your store implementation.
|
||||
If you do not initialize the store with indexing capabilities,
|
||||
@@ -1242,12 +1201,7 @@ class BaseStore(ABC):
|
||||
Args:
|
||||
namespace: Hierarchical path for the item.
|
||||
key: Unique identifier within the namespace.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a namespace label is empty, is not a string,
|
||||
or contains a period (`.`).
|
||||
"""
|
||||
_validate_namespace_labels(namespace)
|
||||
await self.abatch([PutOp(namespace, str(key), None)])
|
||||
|
||||
async def alist_namespaces(
|
||||
@@ -1276,10 +1230,6 @@ class BaseStore(ABC):
|
||||
A list of namespace tuples that match the criteria. Each tuple represents a
|
||||
full namespace path up to `max_depth`.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a `prefix` or `suffix` label is empty, is not a
|
||||
string, or contains a period (`.`).
|
||||
|
||||
???+ example "Examples"
|
||||
|
||||
Setting `max_depth=3` with existing namespaces:
|
||||
@@ -1295,8 +1245,6 @@ class BaseStore(ABC):
|
||||
# Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
|
||||
```
|
||||
"""
|
||||
_validate_namespace_labels(prefix or ())
|
||||
_validate_namespace_labels(suffix or ())
|
||||
match_conditions = []
|
||||
if prefix:
|
||||
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
||||
@@ -1315,14 +1263,6 @@ class BaseStore(ABC):
|
||||
def _validate_namespace(namespace: tuple[str, ...]) -> None:
|
||||
if not namespace:
|
||||
raise InvalidNamespaceError("Namespace cannot be empty.")
|
||||
_validate_namespace_labels(namespace)
|
||||
if namespace[0] == "langgraph":
|
||||
raise InvalidNamespaceError(
|
||||
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
|
||||
)
|
||||
|
||||
|
||||
def _validate_namespace_labels(namespace: tuple[str, ...]) -> None:
|
||||
for label in namespace:
|
||||
if not isinstance(label, str):
|
||||
raise InvalidNamespaceError(
|
||||
@@ -1337,27 +1277,10 @@ def _validate_namespace_labels(namespace: tuple[str, ...]) -> None:
|
||||
raise InvalidNamespaceError(
|
||||
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
|
||||
)
|
||||
|
||||
|
||||
def validate_op_namespace(op: Op) -> None:
|
||||
"""Validate the namespace labels an op carries before a store executes it.
|
||||
|
||||
`BaseStore` methods check labels before batching, but ops passed directly to
|
||||
`batch`/`abatch` skip those methods. Stores that serialize namespaces as
|
||||
delimited text should call this for every op they execute, so a label such
|
||||
as `"foo.bar"` cannot address the namespace `("foo", "bar")`.
|
||||
|
||||
Raises:
|
||||
InvalidNamespaceError: If a label is empty, is not a string, or contains
|
||||
a period (`.`).
|
||||
"""
|
||||
if isinstance(op, (GetOp, PutOp)):
|
||||
_validate_namespace_labels(op.namespace)
|
||||
elif isinstance(op, SearchOp):
|
||||
_validate_namespace_labels(op.namespace_prefix)
|
||||
elif isinstance(op, ListNamespacesOp):
|
||||
for condition in op.match_conditions or ():
|
||||
_validate_namespace_labels(condition.path)
|
||||
if namespace[0] == "langgraph":
|
||||
raise InvalidNamespaceError(
|
||||
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
|
||||
)
|
||||
|
||||
|
||||
def _ensure_refresh(
|
||||
@@ -1396,5 +1319,4 @@ __all__ = [
|
||||
"ensure_embeddings",
|
||||
"tokenize_path",
|
||||
"get_text_at_path",
|
||||
"validate_op_namespace",
|
||||
]
|
||||
|
||||
@@ -25,7 +25,6 @@ from langgraph.store.base import (
|
||||
_ensure_refresh,
|
||||
_ensure_ttl,
|
||||
_validate_namespace,
|
||||
_validate_namespace_labels,
|
||||
)
|
||||
|
||||
F = TypeVar("F", bound=Callable)
|
||||
@@ -87,7 +86,6 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
*,
|
||||
refresh_ttl: bool | None = None,
|
||||
) -> Item | None:
|
||||
_validate_namespace_labels(namespace)
|
||||
self._ensure_task()
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue.put_nowait(
|
||||
@@ -113,7 +111,6 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
offset: int = 0,
|
||||
refresh_ttl: bool | None = None,
|
||||
) -> list[SearchItem]:
|
||||
_validate_namespace_labels(namespace_prefix)
|
||||
self._ensure_task()
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue.put_nowait(
|
||||
@@ -158,7 +155,6 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
namespace: tuple[str, ...],
|
||||
key: str,
|
||||
) -> None:
|
||||
_validate_namespace_labels(namespace)
|
||||
self._ensure_task()
|
||||
fut = self._loop.create_future()
|
||||
self._aqueue.put_nowait((fut, PutOp(namespace, key, None)))
|
||||
@@ -173,8 +169,6 @@ class AsyncBatchedBaseStore(BaseStore):
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
) -> list[tuple[str, ...]]:
|
||||
_validate_namespace_labels(prefix or ())
|
||||
_validate_namespace_labels(suffix or ())
|
||||
self._ensure_task()
|
||||
fut = self._loop.create_future()
|
||||
match_conditions = []
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Run delta-channel conformance capabilities against InMemorySaver."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
conformance = pytest.importorskip(
|
||||
"langgraph.checkpoint.conformance",
|
||||
reason="langgraph-checkpoint-conformance not installed",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_channel_conformance():
|
||||
# Imported inside the test: the module-level importorskip above is what
|
||||
# makes these safe, so they cannot move to the top of the file.
|
||||
from langgraph.checkpoint.conformance import validate # noqa: PLC0415
|
||||
from langgraph.checkpoint.conformance.initializer import ( # noqa: PLC0415
|
||||
checkpointer_test,
|
||||
)
|
||||
|
||||
from langgraph.checkpoint.memory import InMemorySaver # noqa: PLC0415
|
||||
|
||||
@checkpointer_test(name="InMemorySaver")
|
||||
async def mem_saver():
|
||||
yield InMemorySaver()
|
||||
|
||||
report = await validate(
|
||||
mem_saver,
|
||||
capabilities={
|
||||
"delta_channel_history",
|
||||
},
|
||||
)
|
||||
for cap, result in report.results.items():
|
||||
if result.passed is False:
|
||||
details = "\n".join(result.failures or [])
|
||||
pytest.fail(f"Capability {cap} failed:\n{details}")
|
||||
@@ -7,7 +7,6 @@ import pickle
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import uuid
|
||||
from collections import deque
|
||||
from datetime import date, datetime, time, timezone
|
||||
@@ -34,16 +33,12 @@ from langgraph.checkpoint.serde.event_hooks import (
|
||||
register_serde_event_listener,
|
||||
)
|
||||
from langgraph.checkpoint.serde.jsonplus import (
|
||||
EXT_CONSTRUCTOR_KW_ARGS,
|
||||
EXT_CONSTRUCTOR_POS_ARGS,
|
||||
EXT_CONSTRUCTOR_SINGLE_ARG,
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
InvalidModuleError,
|
||||
JsonPlusSerializer,
|
||||
_msgpack_enc,
|
||||
_msgpack_ext_hook_to_json,
|
||||
_warned_blocked_types,
|
||||
_warned_unreconstructable_types,
|
||||
_warned_unregistered_types,
|
||||
)
|
||||
from langgraph.store.base import Item
|
||||
@@ -826,7 +821,6 @@ def _reset_warned_types() -> None:
|
||||
# a fresh slate and assertions about warning emission are stable.
|
||||
_warned_unregistered_types.clear()
|
||||
_warned_blocked_types.clear()
|
||||
_warned_unreconstructable_types.clear()
|
||||
|
||||
|
||||
def test_msgpack_pydantic_warns_by_default(caplog: pytest.LogCaptureFixture) -> None:
|
||||
@@ -1236,58 +1230,3 @@ def test_msgpack_nested_pydantic_serializes_as_dict(
|
||||
# No blocking should occur - inner is serialized as dict, not ext
|
||||
assert "blocked" not in caplog.text.lower()
|
||||
assert result == obj
|
||||
|
||||
|
||||
def test_msgpack_dataclass_from_removed_module_restores_payload(
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
@dataclasses.dataclass
|
||||
class SavedObject:
|
||||
value: int
|
||||
|
||||
SavedObject.__module__ = "removed_module"
|
||||
module = types.ModuleType("removed_module")
|
||||
module.SavedObject = SavedObject
|
||||
monkeypatch.setitem(sys.modules, "removed_module", module)
|
||||
serde = JsonPlusSerializer(
|
||||
allowed_msgpack_modules=[("removed_module", "SavedObject")]
|
||||
)
|
||||
dumped = serde.dumps_typed({"state": SavedObject(123)})
|
||||
assert serde.loads_typed(dumped) == {"state": SavedObject(123)}
|
||||
|
||||
monkeypatch.delitem(sys.modules, "removed_module")
|
||||
caplog.set_level(logging.WARNING, logger="langgraph.checkpoint.serde.jsonplus")
|
||||
|
||||
assert serde.loads_typed(dumped) == {"state": {"value": 123}}
|
||||
assert (
|
||||
"could not reconstruct removed_module.savedobject from checkpoint "
|
||||
"(modulenotfounderror)" in caplog.text.lower()
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("code", "tup", "expected"),
|
||||
[
|
||||
(EXT_CONSTRUCTOR_SINGLE_ARG, ("missing_module", "Thing", "x"), "x"),
|
||||
(EXT_CONSTRUCTOR_POS_ARGS, ("missing_module", "Thing", [1, 2]), [1, 2]),
|
||||
(
|
||||
EXT_CONSTRUCTOR_KW_ARGS,
|
||||
("missing_module", "Thing", {"value": 123}),
|
||||
{"value": 123},
|
||||
),
|
||||
(
|
||||
EXT_METHOD_SINGLE_ARG,
|
||||
("datetime", "datetime", "not-a-date", "fromisoformat"),
|
||||
"not-a-date",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_msgpack_failed_reconstruction_returns_payload(
|
||||
code: int, tup: tuple, expected: object
|
||||
) -> None:
|
||||
serde = JsonPlusSerializer(allowed_msgpack_modules=True)
|
||||
payload = ormsgpack.packb(
|
||||
ormsgpack.Ext(code, _msgpack_enc(tup)), option=ormsgpack.OPT_NON_STR_KEYS
|
||||
)
|
||||
|
||||
assert serde.loads_typed(("msgpack", payload)) == expected
|
||||
|
||||
@@ -420,19 +420,6 @@ class TestInMemorySaverDeltaChannel:
|
||||
assert "seed" not in result
|
||||
assert result["writes"] == []
|
||||
|
||||
def test_get_channel_writes_without_checkpoint_id_reads_the_latest(self) -> None:
|
||||
saver = InMemorySaver()
|
||||
thread: RunnableConfig = {
|
||||
"configurable": {"thread_id": "t1", "checkpoint_ns": ""}
|
||||
}
|
||||
parent = saver.put(thread, empty_checkpoint(), {}, {})
|
||||
saver.put_writes(parent, [("messages", "hi")], "task1")
|
||||
saver.put(parent, empty_checkpoint(), {}, {})
|
||||
|
||||
result = saver.get_delta_channel_history(config=thread, channels=["messages"])
|
||||
|
||||
assert result == {"messages": {"writes": [("task1", "messages", "hi")]}}
|
||||
|
||||
|
||||
class TestBaseFallbackGetChannelWrites:
|
||||
"""Exercises the `BaseCheckpointSaver.get_delta_channel_history` default
|
||||
|
||||
@@ -13,14 +13,10 @@ from langgraph.store.base import (
|
||||
GetOp,
|
||||
InvalidNamespaceError,
|
||||
Item,
|
||||
ListNamespacesOp,
|
||||
MatchCondition,
|
||||
Op,
|
||||
PutOp,
|
||||
Result,
|
||||
SearchOp,
|
||||
get_text_at_path,
|
||||
validate_op_namespace,
|
||||
)
|
||||
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
@@ -532,127 +528,6 @@ async def test_cannot_put_empty_namespace() -> None:
|
||||
assert (await async_store.aget(("valid", "namespace"), "key")) is None
|
||||
|
||||
|
||||
INVALID_NAMESPACES = [("foo.bar",), ("foo", ""), (123,)]
|
||||
NAMESPACE_METHODS = ["get", "delete", "search", "prefix", "suffix"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
|
||||
@pytest.mark.parametrize("method", NAMESPACE_METHODS)
|
||||
def test_rejects_invalid_namespace_labels(
|
||||
mocker: MockerFixture, namespace: tuple, method: str
|
||||
) -> None:
|
||||
store = InMemoryStore()
|
||||
batch = mocker.spy(InMemoryStore, "batch")
|
||||
call = {
|
||||
"get": lambda: store.get(namespace, "key"),
|
||||
"delete": lambda: store.delete(namespace, "key"),
|
||||
"search": lambda: store.search(namespace),
|
||||
"prefix": lambda: store.list_namespaces(prefix=namespace),
|
||||
"suffix": lambda: store.list_namespaces(suffix=namespace),
|
||||
}[method]
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
call()
|
||||
|
||||
batch.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batched", [False, True])
|
||||
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
|
||||
@pytest.mark.parametrize("method", NAMESPACE_METHODS)
|
||||
async def test_async_rejects_invalid_namespace_labels(
|
||||
mocker: MockerFixture, batched: bool, namespace: tuple, method: str
|
||||
) -> None:
|
||||
# The batched store must reject before queueing: a failure inside the
|
||||
# shared `abatch` would fail every op queued alongside this one.
|
||||
store = MockAsyncBatchedStore() if batched else InMemoryStore()
|
||||
# `MockAsyncBatchedStore` dispatches through `InMemoryStore.batch`.
|
||||
batch = mocker.spy(InMemoryStore, "batch")
|
||||
abatch = mocker.spy(InMemoryStore, "abatch")
|
||||
call = {
|
||||
"get": lambda: store.aget(namespace, "key"),
|
||||
"delete": lambda: store.adelete(namespace, "key"),
|
||||
"search": lambda: store.asearch(namespace),
|
||||
"prefix": lambda: store.alist_namespaces(prefix=namespace),
|
||||
"suffix": lambda: store.alist_namespaces(suffix=namespace),
|
||||
}[method]
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
await call()
|
||||
|
||||
batch.assert_not_called()
|
||||
abatch.assert_not_called()
|
||||
|
||||
|
||||
def test_search_and_listing_keep_empty_prefixes_and_wildcards() -> None:
|
||||
store = InMemoryStore()
|
||||
store.put(("tenant", "a_%"), "key", {"v": 1})
|
||||
store.put(("tenant", "b", "child"), "key", {"v": 1})
|
||||
|
||||
assert len(store.search(())) == 2
|
||||
assert [item.namespace for item in store.search(("tenant", "a_%"))] == [
|
||||
("tenant", "a_%")
|
||||
]
|
||||
assert sorted(store.list_namespaces(prefix=("tenant", "*"), suffix=("*",))) == [
|
||||
("tenant", "a_%"),
|
||||
("tenant", "b", "child"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batched", [False, True])
|
||||
async def test_async_search_and_listing_keep_empty_prefixes_and_wildcards(
|
||||
batched: bool,
|
||||
) -> None:
|
||||
store = MockAsyncBatchedStore() if batched else InMemoryStore()
|
||||
await store.aput(("tenant", "a_%"), "key", {"v": 1})
|
||||
await store.aput(("tenant", "b", "child"), "key", {"v": 1})
|
||||
|
||||
assert len(await store.asearch(())) == 2
|
||||
assert [item.namespace for item in await store.asearch(("tenant", "a_%"))] == [
|
||||
("tenant", "a_%")
|
||||
]
|
||||
assert sorted(
|
||||
await store.alist_namespaces(prefix=("tenant", "*"), suffix=("*",))
|
||||
) == [("tenant", "a_%"), ("tenant", "b", "child")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", INVALID_NAMESPACES)
|
||||
@pytest.mark.parametrize(
|
||||
"kind", ["get", "put", "delete", "search", "list_prefix", "list_suffix"]
|
||||
)
|
||||
def test_validate_op_namespace_rejects_invalid_labels(
|
||||
namespace: tuple, kind: str
|
||||
) -> None:
|
||||
op = {
|
||||
"get": GetOp(namespace, "key"),
|
||||
"put": PutOp(namespace, "key", {"v": 1}),
|
||||
"delete": PutOp(namespace, "key", None),
|
||||
"search": SearchOp(namespace),
|
||||
"list_prefix": ListNamespacesOp((MatchCondition("prefix", namespace),)),
|
||||
"list_suffix": ListNamespacesOp((MatchCondition("suffix", namespace),)),
|
||||
}[kind]
|
||||
|
||||
with pytest.raises(InvalidNamespaceError):
|
||||
validate_op_namespace(op)
|
||||
|
||||
|
||||
def test_validate_op_namespace_allows_empty_prefix_and_wildcards() -> None:
|
||||
for op in (
|
||||
SearchOp(()),
|
||||
ListNamespacesOp(),
|
||||
ListNamespacesOp(
|
||||
(
|
||||
MatchCondition("prefix", ("tenant", "*")),
|
||||
MatchCondition("suffix", ("*",)),
|
||||
)
|
||||
),
|
||||
GetOp(("tenant", "a_%"), "key"),
|
||||
# Write-only rules belong to `put`, not to op validation.
|
||||
PutOp(("langgraph", "x"), "key", {"v": 1}),
|
||||
):
|
||||
validate_op_namespace(op)
|
||||
|
||||
|
||||
async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None:
|
||||
abatch = mocker.spy(InMemoryStore, "batch")
|
||||
store = MockAsyncBatchedStore()
|
||||
|
||||
@@ -2518,6 +2518,13 @@ def deploy_listeners() -> None:
|
||||
|
||||
@OPT_HOST_API_KEY
|
||||
@OPT_HOST_URL
|
||||
@click.option(
|
||||
"--json",
|
||||
"json_output",
|
||||
is_flag=True,
|
||||
default=False,
|
||||
help="Emit the listeners as a JSON array instead of a table.",
|
||||
)
|
||||
@deploy_listeners.command(
|
||||
"list",
|
||||
help=(
|
||||
@@ -2526,11 +2533,18 @@ def deploy_listeners() -> None:
|
||||
"--listener-id <id>` to deploy through it."
|
||||
),
|
||||
)
|
||||
def deploy_listeners_list(api_key: str | None, host_url: str | None) -> None:
|
||||
def deploy_listeners_list(
|
||||
api_key: str | None, host_url: str | None, json_output: bool
|
||||
) -> None:
|
||||
global _no_input
|
||||
_no_input = json_output
|
||||
client = _create_host_backend_client(host_url, api_key)
|
||||
listeners = _call_host_backend_with_optional_tenant(
|
||||
client, lambda c: c.list_listeners()
|
||||
)
|
||||
if json_output:
|
||||
click.echo(json_mod.dumps(listeners))
|
||||
return
|
||||
if not listeners:
|
||||
click.echo("No listeners found for this workspace.")
|
||||
return
|
||||
|
||||
@@ -14,6 +14,7 @@ import langgraph_cli.deploy as deploy_module
|
||||
from langgraph_cli.cli import cli, prepare_args_and_stdin
|
||||
from langgraph_cli.config import Config, _get_pip_cleanup_lines, validate_config
|
||||
from langgraph_cli.docker import DEFAULT_POSTGRES_URI, DockerCapabilities, Version
|
||||
from langgraph_cli.host_backend import HostBackendError
|
||||
from langgraph_cli.util import clean_empty_lines
|
||||
|
||||
FORMATTED_CLEANUP_LINES = _get_pip_cleanup_lines(
|
||||
@@ -535,6 +536,131 @@ def test_deploy_listeners_list_command_no_results(monkeypatch) -> None:
|
||||
assert result.output.strip() == "No listeners found for this workspace."
|
||||
|
||||
|
||||
def test_deploy_listeners_list_command_json(monkeypatch) -> None:
|
||||
# --json sets the module-level _no_input global as a side effect; pin
|
||||
# the "before" value so monkeypatch restores it after this test.
|
||||
monkeypatch.setattr(deploy_module, "_no_input", False)
|
||||
runner = CliRunner()
|
||||
|
||||
listeners = [
|
||||
{
|
||||
"id": "listener-1",
|
||||
"compute_id": "prod-cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents"]},
|
||||
},
|
||||
{
|
||||
"id": "listener-2",
|
||||
"compute_id": "multi-cluster",
|
||||
"compute_config": {"k8s_namespaces": ["agents", "agents-staging"]},
|
||||
},
|
||||
]
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_listeners(self):
|
||||
return listeners
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"listeners",
|
||||
"list",
|
||||
"--json",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.output) == listeners
|
||||
|
||||
|
||||
def test_deploy_listeners_list_command_json_no_results(monkeypatch) -> None:
|
||||
# --json sets the module-level _no_input global as a side effect; pin
|
||||
# the "before" value so monkeypatch restores it after this test.
|
||||
monkeypatch.setattr(deploy_module, "_no_input", False)
|
||||
runner = CliRunner()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_listeners(self):
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"listeners",
|
||||
"list",
|
||||
"--json",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.output) == []
|
||||
|
||||
|
||||
def test_deploy_listeners_list_command_json_org_scoped_403(monkeypatch) -> None:
|
||||
"""With --json, an org-scoped 403 must raise cleanly instead of
|
||||
prompting for a workspace ID, which would write to stdout and hang
|
||||
waiting for input that a scripted invocation never provides."""
|
||||
# deploy_listeners_list sets the module-level _no_input global as a
|
||||
# side effect; pin its "before" value so monkeypatch restores it after
|
||||
# this test instead of leaking True into whichever test runs next.
|
||||
monkeypatch.setattr(deploy_module, "_no_input", False)
|
||||
runner = CliRunner()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, host_url: str, api_key: str, tenant_id: str | None = None):
|
||||
pass
|
||||
|
||||
def list_listeners(self):
|
||||
raise HostBackendError("requires workspace specification", status_code=403)
|
||||
|
||||
def set_tenant(self, tenant_id: str) -> None:
|
||||
raise AssertionError("must not retry with a prompted tenant with --json")
|
||||
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", FakeClient)
|
||||
|
||||
prompt_calls = []
|
||||
monkeypatch.setattr(
|
||||
click, "prompt", lambda *a, **k: prompt_calls.append((a, k)) or "unreachable"
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
cli,
|
||||
[
|
||||
"deploy",
|
||||
"listeners",
|
||||
"list",
|
||||
"--json",
|
||||
"--api-key",
|
||||
"test-key",
|
||||
"--host-url",
|
||||
"https://api.example.com",
|
||||
],
|
||||
)
|
||||
|
||||
assert prompt_calls == [], "must not prompt interactively with --json"
|
||||
assert result.exit_code != 0
|
||||
assert "LANGSMITH_TENANT_ID" in result.output
|
||||
|
||||
|
||||
def test_deploy_revisions_list_command(monkeypatch) -> None:
|
||||
runner = CliRunner()
|
||||
captured: dict[str, str] = {}
|
||||
|
||||
@@ -28,7 +28,6 @@ __all__ = (
|
||||
"ParentCommand",
|
||||
"EmptyInputError",
|
||||
"TaskNotFound",
|
||||
"is_invalid_resume",
|
||||
)
|
||||
|
||||
|
||||
@@ -240,21 +239,3 @@ class NodeTimeoutError(Exception):
|
||||
self.kind = kind
|
||||
self.idle_timeout = idle_timeout
|
||||
self.run_timeout = run_timeout
|
||||
|
||||
|
||||
_INVALID_RESUME = "_langgraph_invalid_resume"
|
||||
|
||||
|
||||
def _mark_invalid_resume(error: BaseException) -> None:
|
||||
setattr(error, _INVALID_RESUME, True)
|
||||
|
||||
|
||||
def is_invalid_resume(error: BaseException) -> bool:
|
||||
"""Whether `error` was raised because a resume value didn't match `response_schema`.
|
||||
|
||||
`interrupt()` raises a `pydantic.ValidationError` in that case. `ToolNode` uses
|
||||
this to tell it apart from invalid tool arguments when a tool calls `interrupt()`
|
||||
or runs a graph that does, so the resume fails and the interrupt can be answered
|
||||
again.
|
||||
"""
|
||||
return getattr(error, _INVALID_RESUME, False) is True
|
||||
|
||||
@@ -24,9 +24,6 @@ from typing_extensions import TypedDict, deprecated
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel._messages import (
|
||||
_messages_delta_reducer as _messages_delta_reducer,
|
||||
)
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
|
||||
__all__ = (
|
||||
@@ -247,6 +244,71 @@ def add_messages(
|
||||
return merged
|
||||
|
||||
|
||||
def _messages_delta_reducer(
|
||||
state: list[AnyMessage], writes: list[list[AnyMessage]]
|
||||
) -> list[AnyMessage]:
|
||||
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
|
||||
|
||||
Processes all writes in one pass — dedup by ID, `RemoveMessage`
|
||||
tombstoning — without calling `add_messages`.
|
||||
|
||||
This reducer is batching-invariant, as required by `DeltaChannel`:
|
||||
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
|
||||
|
||||
Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
|
||||
objects so that HTTP-driven graphs work without a separate coercion
|
||||
step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
|
||||
unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
|
||||
`BaseMessageChunk` conversion are not handled here.
|
||||
|
||||
Example::
|
||||
|
||||
from typing import Annotated
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
"""
|
||||
|
||||
# Each write is either a list of message-likes or a single message-like
|
||||
# (BaseMessage / dict / str / tuple). Only lists flatten; everything
|
||||
# else is one message.
|
||||
flat: list[Any] = []
|
||||
for w in writes:
|
||||
if isinstance(w, list):
|
||||
flat.extend(w)
|
||||
else:
|
||||
flat.append(w)
|
||||
# Steady state: the reducer's own output is already typed, so skip
|
||||
# `convert_to_messages` on state when the first element is a BaseMessage.
|
||||
# Only raw input (initial dicts, deserialized blobs) hits the slow path.
|
||||
if state and isinstance(state[0], BaseMessage):
|
||||
state_msgs = state
|
||||
else:
|
||||
state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
|
||||
msgs = cast("list[AnyMessage]", convert_to_messages(flat))
|
||||
|
||||
index: dict[str, int] = {
|
||||
m.id: i for i, m in enumerate(state_msgs) if m.id is not None
|
||||
}
|
||||
result: list[AnyMessage | None] = list(state_msgs)
|
||||
for msg in msgs:
|
||||
mid = msg.id
|
||||
if mid is None:
|
||||
result.append(msg)
|
||||
elif isinstance(msg, RemoveMessage):
|
||||
if mid in index:
|
||||
result[index[mid]] = None
|
||||
del index[mid]
|
||||
elif mid in index:
|
||||
result[index[mid]] = msg
|
||||
else:
|
||||
index[mid] = len(result)
|
||||
result.append(msg)
|
||||
return [m for m in result if m is not None]
|
||||
|
||||
|
||||
@deprecated(
|
||||
"MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.",
|
||||
category=None,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from inspect import signature
|
||||
from typing import Any, Literal, cast
|
||||
@@ -32,7 +32,6 @@ from langgraph._internal._constants import (
|
||||
)
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
@@ -153,23 +152,6 @@ def delta_channels_with_pending_writes(
|
||||
}
|
||||
|
||||
|
||||
def delta_channels_overwritten(
|
||||
specs: Mapping[str, Any], writes: Iterable[tuple[str, Any]]
|
||||
) -> set[str]:
|
||||
"""Return the names of the DeltaChannels that `writes` set with an `Overwrite`.
|
||||
|
||||
`update_state` saves a full snapshot of these channels in the checkpoint it
|
||||
creates, like the loop does when a node returns an `Overwrite`. Otherwise,
|
||||
reading the channel later starts from an older snapshot and replays the
|
||||
writes the `Overwrite` threw away.
|
||||
"""
|
||||
return {
|
||||
ch
|
||||
for ch, value in writes
|
||||
if isinstance(specs.get(ch), DeltaChannel) and _get_overwrite(value)[0]
|
||||
}
|
||||
|
||||
|
||||
def checkpoint_superseded(
|
||||
saver: BaseCheckpointSaver, config: RunnableConfig, saved: CheckpointTuple
|
||||
) -> bool:
|
||||
@@ -274,7 +256,6 @@ def create_checkpoint(
|
||||
get_next_version: GetNextVersion | None = None,
|
||||
channels_to_snapshot: set[str] | None = None,
|
||||
stored_versions: ChannelVersions | None = None,
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Build a new Checkpoint from the previous one and live channel state.
|
||||
|
||||
@@ -333,9 +314,7 @@ def create_checkpoint(
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=channel_versions,
|
||||
versions_seen=_mark_bumps_seen(
|
||||
checkpoint["versions_seen"], bumped, trigger_to_nodes or {}
|
||||
),
|
||||
versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped),
|
||||
updated_channels=None if updated_channels is None else sorted(updated_channels),
|
||||
)
|
||||
|
||||
@@ -343,27 +322,19 @@ def create_checkpoint(
|
||||
def _mark_bumps_seen(
|
||||
versions_seen: dict[str, ChannelVersions],
|
||||
bumped: Mapping[str, tuple[Any, Any]],
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]],
|
||||
) -> dict[str, ChannelVersions]:
|
||||
"""Advance whoever had seen a bumped channel's old version to the new one.
|
||||
|
||||
A bump that only stores a snapshot is not a write. Left unseen, it would
|
||||
re-fire `interrupt_before` and rerun the channel's subscribers. A channel
|
||||
bumped from no version was never written, so it also goes to the
|
||||
subscribers that never ran: they have no entry, and would start on the bump.
|
||||
For each entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the
|
||||
one the node really read, so `versions_seen_without_bumps` can put the read
|
||||
back.
|
||||
re-fire `interrupt_before` and rerun the channel's subscribers. For each
|
||||
entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the
|
||||
node really read, so `versions_seen_without_bumps` can put the read back.
|
||||
"""
|
||||
if not bumped:
|
||||
return versions_seen
|
||||
out = dict(versions_seen)
|
||||
for k, (old, _) in bumped.items():
|
||||
if old is None:
|
||||
for node in trigger_to_nodes.get(k, ()):
|
||||
out.setdefault(node, {})
|
||||
marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {}))
|
||||
for node, seen in out.items():
|
||||
for node, seen in versions_seen.items():
|
||||
if node == SNAPSHOT_BUMPS:
|
||||
continue
|
||||
for k, (old, new) in bumped.items():
|
||||
|
||||
@@ -19,7 +19,6 @@ from typing import (
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
from uuid import UUID, uuid5
|
||||
|
||||
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
@@ -161,12 +160,6 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol:
|
||||
return StreamProtocol(__call__, {mode for s in streams for mode in s.modes})
|
||||
|
||||
|
||||
def _cacheable(writes: WritesT) -> bool:
|
||||
# A cache hit skips the node, so a cached error would skip it without
|
||||
# raising: only a node that finished goes to the cache.
|
||||
return not any(c in (INTERRUPT, ERROR) for c, _ in writes)
|
||||
|
||||
|
||||
class PregelLoop:
|
||||
config: RunnableConfig
|
||||
store: BaseStore | None
|
||||
@@ -198,7 +191,6 @@ class PregelLoop:
|
||||
Callable[
|
||||
[
|
||||
concurrent.futures.Future | None,
|
||||
Sequence[Any],
|
||||
RunnableConfig,
|
||||
Checkpoint,
|
||||
str,
|
||||
@@ -212,13 +204,11 @@ class PregelLoop:
|
||||
submit: Submit
|
||||
channels: Mapping[str, BaseChannel]
|
||||
# Futures from `checkpointer.put_writes` calls that produced delta-channel
|
||||
# writes. `_put_checkpoint` hands this list to the save it submits, which
|
||||
# waits for them first, so a checkpoint never becomes durable before the
|
||||
# writes that produced it. If a write or the previous save failed, the
|
||||
# save fails too: a DeltaChannel is rebuilt from its writes along the
|
||||
# parent chain, so a checkpoint saved past either gap reads back short
|
||||
# for good. Initialised to `[]` in both sync and async `__enter__`;
|
||||
# stays `None` only when no checkpointer.
|
||||
# writes. `_checkpointer_put_after_previous` drains this list (swap to a
|
||||
# local `futs` then reset to `[]` and wait/gather) before putting the
|
||||
# next checkpoint, so a checkpoint never becomes durable before the
|
||||
# writes that produced it. Initialised to `[]` in both sync and async
|
||||
# `__enter__`; stays `None` only when no checkpointer.
|
||||
_delta_write_futs: list[Any] | None = None
|
||||
|
||||
# Same pattern as `_delta_write_futs` but for error-handler writes.
|
||||
@@ -276,9 +266,6 @@ class PregelLoop:
|
||||
# `_put_exit_delta_writes` uses this to decide between anchoring on
|
||||
# the existing parent (True) or creating a lazy stub (False).
|
||||
_has_persisted_parent: bool = False
|
||||
# True iff `__enter__` loaded the thread's latest checkpoint, not one a
|
||||
# `checkpoint_id` addressed, so nothing has been built on it yet.
|
||||
_loaded_latest: bool = False
|
||||
|
||||
managed: ManagedValueMapping
|
||||
checkpoint: Checkpoint
|
||||
@@ -449,18 +436,13 @@ class PregelLoop:
|
||||
return None
|
||||
return self._graph_lifecycle_events.popleft()
|
||||
|
||||
def put_writes(
|
||||
self, task_id: str, writes: WritesT, *, cached: bool = False
|
||||
) -> None:
|
||||
def put_writes(self, task_id: str, writes: WritesT) -> None:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
if not writes:
|
||||
return
|
||||
# deduplicate writes to special channels, last write wins
|
||||
if all(w[0] in WRITES_IDX_MAP for w in writes):
|
||||
writes = list({w[0]: w for w in writes}.values())
|
||||
# Before the writes are saved, kept in `checkpoint_pending_writes` or
|
||||
# applied, so all of them hold the same message ids.
|
||||
ensure_message_ids(self.specs, writes)
|
||||
if task_id == NULL_TASK_ID:
|
||||
# writes for the null task are accumulated
|
||||
self.checkpoint_pending_writes = [
|
||||
@@ -497,6 +479,14 @@ class PregelLoop:
|
||||
|
||||
# save writes
|
||||
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
|
||||
# Assign stable IDs to any id=None BaseMessages in DeltaChannel writes
|
||||
# before the background thread serialises them. Without this, reducers
|
||||
# that assign IDs inside apply_writes() race with serialisation and
|
||||
# store id=None, causing get_state() replays to produce a different UUID
|
||||
# on every call.
|
||||
for c, v in writes_to_save:
|
||||
if isinstance(self.specs.get(c), DeltaChannel):
|
||||
ensure_message_ids(v)
|
||||
if self.durability != "exit" and self.checkpointer_put_writes is not None:
|
||||
config = patch_configurable(
|
||||
self.checkpoint_config,
|
||||
@@ -539,7 +529,7 @@ class PregelLoop:
|
||||
self._error_handler_write_futs.append(fut)
|
||||
# output writes
|
||||
if hasattr(self, "tasks"):
|
||||
self.output_writes(task_id, writes, cached=cached)
|
||||
self.output_writes(task_id, writes)
|
||||
|
||||
def _put_pending_writes(self) -> None:
|
||||
if self.checkpointer_put_writes is None:
|
||||
@@ -688,6 +678,10 @@ class PregelLoop:
|
||||
self.status = "done"
|
||||
return False
|
||||
|
||||
if self.control is not None and self.control.drain_requested:
|
||||
self.status = "draining"
|
||||
return False
|
||||
|
||||
# if there are pending writes from a previous loop, apply them
|
||||
if self._reapplies_pending_writes and self.checkpoint_pending_writes:
|
||||
self._reapply_writes_to_succeeded_nodes(self.tasks)
|
||||
@@ -700,13 +694,6 @@ class PregelLoop:
|
||||
self.status = "interrupt_before"
|
||||
raise GraphInterrupt()
|
||||
|
||||
# stop before running the next tasks if a drain was requested. after
|
||||
# the breakpoint check: a resume passes the next step's breakpoints,
|
||||
# so draining in front of one would skip it
|
||||
if self.control is not None and self.control.drain_requested:
|
||||
self.status = "draining"
|
||||
return False
|
||||
|
||||
# produce debug output
|
||||
self._emit("tasks", map_debug_tasks, self.tasks.values())
|
||||
|
||||
@@ -1018,8 +1005,6 @@ class PregelLoop:
|
||||
# save writes
|
||||
for tid, ws in writes.items():
|
||||
if forks:
|
||||
# These skip `put_writes`, which gives messages their ids.
|
||||
ensure_message_ids(self.specs, ws)
|
||||
carried.extend((tid, c, v) for c, v in ws)
|
||||
else:
|
||||
self.put_writes(tid, ws)
|
||||
@@ -1084,9 +1069,6 @@ class PregelLoop:
|
||||
)
|
||||
# map inputs to channel updates
|
||||
elif input_writes := deque(map_input(input_keys, self.input)):
|
||||
# Before the input is applied, so this run reads the message ids
|
||||
# that are saved with it.
|
||||
ensure_message_ids(self.specs, input_writes)
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
discard_tasks = prepare_next_tasks(
|
||||
self.checkpoint,
|
||||
@@ -1132,26 +1114,16 @@ class PregelLoop:
|
||||
self._exit_delta_writes.append(
|
||||
(self.step, NULL_TASK_ID, "", c, v)
|
||||
)
|
||||
# A DeltaChannel reads its input from the writes stored on the
|
||||
# checkpoint this run starts from, under a task id of their own:
|
||||
# readers apply a checkpoint's NULL_TASK_ID writes as its own state.
|
||||
# A new thread has no checkpoint to store them on, and one a
|
||||
# `checkpoint_id` addressed may have children that would read them,
|
||||
# so then the input checkpoint snapshots the channel instead.
|
||||
# Persist delta-channel input writes so sub-freq inputs are
|
||||
# recoverable via ancestor walk (mirrors the Command input path).
|
||||
if self.durability != "exit":
|
||||
delta_input = [
|
||||
(c, v)
|
||||
for c, v in input_writes
|
||||
if isinstance(self.specs.get(c), DeltaChannel)
|
||||
]
|
||||
if delta_input and self._has_persisted_parent and self._loaded_latest:
|
||||
self.put_writes(
|
||||
str(uuid5(UUID(self.checkpoint["id"]), INPUT)), delta_input
|
||||
)
|
||||
else:
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
c for c, _ in delta_input
|
||||
)
|
||||
if delta_input:
|
||||
self.put_writes(NULL_TASK_ID, delta_input)
|
||||
# save input checkpoint
|
||||
self.updated_channels = updated_channels
|
||||
self._put_checkpoint({"source": "input"})
|
||||
@@ -1281,7 +1253,6 @@ class PregelLoop:
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=self.checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
for k in channels_to_snapshot:
|
||||
new_counters[k] = (0, 0)
|
||||
@@ -1328,17 +1299,12 @@ class PregelLoop:
|
||||
)
|
||||
self.checkpoint_previous_versions = channel_versions
|
||||
|
||||
# Take this checkpoint's writes now: saves run in the background
|
||||
# and can start out of order, so a save that took them itself
|
||||
# could get another checkpoint's writes.
|
||||
delta_write_futs, self._delta_write_futs = self._delta_write_futs, []
|
||||
# save it, without blocking
|
||||
# if there's a previous checkpoint save in progress, wait for it
|
||||
# ensuring checkpointers receive checkpoints in order
|
||||
self._put_checkpoint_fut = self.submit(
|
||||
self._checkpointer_put_after_previous,
|
||||
getattr(self, "_put_checkpoint_fut", None),
|
||||
delta_write_futs,
|
||||
self.checkpoint_config,
|
||||
copy_checkpoint(self.checkpoint),
|
||||
self.checkpoint_metadata,
|
||||
@@ -1371,14 +1337,6 @@ class PregelLoop:
|
||||
):
|
||||
return
|
||||
|
||||
# A checkpoint a `checkpoint_id` addressed may already have children,
|
||||
# which read the writes stored on it, so then the channels this run
|
||||
# wrote snapshot instead.
|
||||
if self._has_persisted_parent and not self._loaded_latest:
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
w[3] for w in self._exit_delta_writes
|
||||
)
|
||||
|
||||
counters = dict(
|
||||
self.checkpoint_metadata.get("counters_since_delta_snapshot") or {}
|
||||
)
|
||||
@@ -1416,7 +1374,6 @@ class PregelLoop:
|
||||
self._put_checkpoint_fut = self.submit(
|
||||
self._checkpointer_put_after_previous,
|
||||
getattr(self, "_put_checkpoint_fut", None),
|
||||
(),
|
||||
stub_put_config,
|
||||
stub_cp,
|
||||
{"step": -2},
|
||||
@@ -1684,19 +1641,21 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
def _checkpointer_put_after_previous(
|
||||
self,
|
||||
prev: concurrent.futures.Future | None,
|
||||
delta_write_futs: Sequence[concurrent.futures.Future],
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
for fut in delta_write_futs:
|
||||
fut.result()
|
||||
if prev is not None:
|
||||
prev.result()
|
||||
cast(BaseCheckpointSaver, self.checkpointer).put(
|
||||
config, checkpoint, metadata, new_versions
|
||||
)
|
||||
if self._delta_write_futs:
|
||||
futs, self._delta_write_futs = self._delta_write_futs, []
|
||||
concurrent.futures.wait(futs)
|
||||
try:
|
||||
if prev is not None:
|
||||
prev.result()
|
||||
finally:
|
||||
cast(BaseCheckpointSaver, self.checkpointer).put(
|
||||
config, checkpoint, metadata, new_versions
|
||||
)
|
||||
|
||||
def match_cached_writes(self) -> Sequence[PregelExecutableTask]:
|
||||
if self.cache is None:
|
||||
@@ -1718,7 +1677,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
) -> PregelExecutableTask | None:
|
||||
if pushed := super().accept_push(task, write_idx, call):
|
||||
for task in self.match_cached_writes():
|
||||
self.put_writes(task.id, task.writes, cached=True)
|
||||
self.output_writes(task.id, task.writes, cached=True)
|
||||
return pushed
|
||||
|
||||
def schedule_error_handler(
|
||||
@@ -1755,18 +1714,16 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
if self._reapplies_pending_writes:
|
||||
self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
|
||||
for task in self.match_cached_writes():
|
||||
self.put_writes(task.id, task.writes, cached=True)
|
||||
self.output_writes(task.id, task.writes, cached=True)
|
||||
return handler_task
|
||||
|
||||
def put_writes(
|
||||
self, task_id: str, writes: WritesT, *, cached: bool = False
|
||||
) -> None:
|
||||
def put_writes(self, task_id: str, writes: WritesT) -> None:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
super().put_writes(task_id, writes, cached=cached)
|
||||
if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
|
||||
super().put_writes(task_id, writes)
|
||||
if not writes or self.cache is None or not hasattr(self, "tasks"):
|
||||
return
|
||||
task = self.tasks.get(task_id)
|
||||
if task is None or task.cache_key is None or not _cacheable(writes):
|
||||
if task is None or task.cache_key is None:
|
||||
return
|
||||
self.submit(
|
||||
self.cache.set,
|
||||
@@ -1810,7 +1767,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
self._loaded_latest = True
|
||||
|
||||
# Capture before the synthetic-empty fallback below overwrites `saved`.
|
||||
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
|
||||
@@ -1940,19 +1896,23 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
async def _checkpointer_put_after_previous(
|
||||
self,
|
||||
prev: asyncio.Task | None,
|
||||
delta_write_futs: Sequence[asyncio.Future],
|
||||
config: RunnableConfig,
|
||||
checkpoint: Checkpoint,
|
||||
metadata: CheckpointMetadata,
|
||||
new_versions: ChannelVersions,
|
||||
) -> RunnableConfig:
|
||||
if delta_write_futs:
|
||||
await asyncio.gather(*delta_write_futs)
|
||||
if prev is not None:
|
||||
await prev
|
||||
await cast(BaseCheckpointSaver, self.checkpointer).aput(
|
||||
config, checkpoint, metadata, new_versions
|
||||
)
|
||||
# Drain DeltaChannel write futures before committing the checkpoint so
|
||||
# ancestor walks never see a checkpoint without its backing writes.
|
||||
if self._delta_write_futs:
|
||||
futs, self._delta_write_futs = self._delta_write_futs, []
|
||||
await asyncio.gather(*futs)
|
||||
try:
|
||||
if prev is not None:
|
||||
await prev
|
||||
finally:
|
||||
await cast(BaseCheckpointSaver, self.checkpointer).aput(
|
||||
config, checkpoint, metadata, new_versions
|
||||
)
|
||||
|
||||
async def amatch_cached_writes(self) -> Sequence[PregelExecutableTask]:
|
||||
if self.cache is None:
|
||||
@@ -1974,7 +1934,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
) -> PregelExecutableTask | None:
|
||||
if pushed := super().accept_push(task, write_idx, call):
|
||||
for task in await self.amatch_cached_writes():
|
||||
self.put_writes(task.id, task.writes, cached=True)
|
||||
self.output_writes(task.id, task.writes, cached=True)
|
||||
return pushed
|
||||
|
||||
async def aschedule_error_handler(
|
||||
@@ -2011,18 +1971,19 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
if self._reapplies_pending_writes:
|
||||
self._reapply_writes_to_succeeded_nodes({handler_task.id: handler_task})
|
||||
for task in await self.amatch_cached_writes():
|
||||
self.put_writes(task.id, task.writes, cached=True)
|
||||
self.output_writes(task.id, task.writes, cached=True)
|
||||
return handler_task
|
||||
|
||||
def put_writes(
|
||||
self, task_id: str, writes: WritesT, *, cached: bool = False
|
||||
) -> None:
|
||||
def put_writes(self, task_id: str, writes: WritesT) -> None:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
super().put_writes(task_id, writes, cached=cached)
|
||||
if cached or not writes or self.cache is None or not hasattr(self, "tasks"):
|
||||
super().put_writes(task_id, writes)
|
||||
if not writes or self.cache is None or not hasattr(self, "tasks"):
|
||||
return
|
||||
task = self.tasks.get(task_id)
|
||||
if task is None or task.cache_key is None or not _cacheable(writes):
|
||||
if task is None or task.cache_key is None:
|
||||
return
|
||||
if writes[0][0] in (INTERRUPT, ERROR):
|
||||
# only cache successful tasks
|
||||
return
|
||||
self.submit(
|
||||
self.cache.aset,
|
||||
@@ -2066,7 +2027,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
self._loaded_latest = True
|
||||
|
||||
# Capture before the synthetic-empty fallback below overwrites `saved`.
|
||||
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
|
||||
|
||||
@@ -1,13 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import (
|
||||
AsyncIterator,
|
||||
Callable,
|
||||
Iterator,
|
||||
Mapping,
|
||||
MutableSequence,
|
||||
Sequence,
|
||||
)
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from dataclasses import fields, is_dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -17,18 +10,16 @@ from typing import (
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.messages import AnyMessage, BaseMessage, RemoveMessage, ToolMessage
|
||||
from langchain_core.messages import BaseMessage, ToolMessage
|
||||
from langchain_core.messages.utils import convert_to_messages
|
||||
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph._internal._config import filter_to_user_tags
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
from langgraph.types import Command, Overwrite
|
||||
from langgraph.types import Command
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
@@ -417,74 +408,6 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
|
||||
self.stream((meta[0], "messages", (event, v2_meta)))
|
||||
|
||||
|
||||
# Defined here rather than in `langgraph.graph.message`, which re-exports it:
|
||||
# `ensure_message_ids` has to recognize it, and importing that module from the
|
||||
# loop would be an import cycle.
|
||||
def _messages_delta_reducer(
|
||||
state: list[AnyMessage], writes: list[list[AnyMessage]]
|
||||
) -> list[AnyMessage]:
|
||||
"""**Experimental.** Batch reducer for use with `DeltaChannel`.
|
||||
|
||||
Processes all writes in one pass — dedup by ID, `RemoveMessage`
|
||||
tombstoning — without calling `add_messages`.
|
||||
|
||||
This reducer is batching-invariant, as required by `DeltaChannel`:
|
||||
`reducer(reducer(state, xs), ys) == reducer(state, xs + ys)`.
|
||||
|
||||
Raw dict / string / tuple inputs are coerced to typed `BaseMessage`
|
||||
objects so that HTTP-driven graphs work without a separate coercion
|
||||
step. This is not full `add_messages` parity — `REMOVE_ALL_MESSAGES`,
|
||||
unknown-id `RemoveMessage` errors, missing-id UUID assignment, and
|
||||
`BaseMessageChunk` conversion are not handled here.
|
||||
|
||||
Example::
|
||||
|
||||
from typing import Annotated
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
"""
|
||||
|
||||
# Each write is either a list of message-likes or a single message-like
|
||||
# (BaseMessage / dict / str / tuple). Only lists flatten; everything
|
||||
# else is one message.
|
||||
flat: list[Any] = []
|
||||
for w in writes:
|
||||
if isinstance(w, list):
|
||||
flat.extend(w)
|
||||
else:
|
||||
flat.append(w)
|
||||
# Steady state: the reducer's own output is already typed, so skip
|
||||
# `convert_to_messages` on state when the first element is a BaseMessage.
|
||||
# Only raw input (initial dicts, deserialized blobs) hits the slow path.
|
||||
if state and isinstance(state[0], BaseMessage):
|
||||
state_msgs = state
|
||||
else:
|
||||
state_msgs = cast("list[AnyMessage]", convert_to_messages(state))
|
||||
msgs = cast("list[AnyMessage]", convert_to_messages(flat))
|
||||
|
||||
index: dict[str, int] = {
|
||||
m.id: i for i, m in enumerate(state_msgs) if m.id is not None
|
||||
}
|
||||
result: list[AnyMessage | None] = list(state_msgs)
|
||||
for msg in msgs:
|
||||
mid = msg.id
|
||||
if mid is None:
|
||||
result.append(msg)
|
||||
elif isinstance(msg, RemoveMessage):
|
||||
if mid in index:
|
||||
result[index[mid]] = None
|
||||
del index[mid]
|
||||
elif mid in index:
|
||||
result[index[mid]] = msg
|
||||
else:
|
||||
index[mid] = len(result)
|
||||
result.append(msg)
|
||||
return [m for m in result if m is not None]
|
||||
|
||||
|
||||
# Known role values (OpenAI-style) and type values (LangChain serialisation)
|
||||
# that identify a dict as a message. Checked before coercing to BaseMessage so
|
||||
# we don't accidentally touch unrelated dicts that happen to have a "role" key.
|
||||
@@ -500,103 +423,39 @@ def _is_message_dict(item: dict) -> bool:
|
||||
return item.get("role") in _MESSAGE_ROLES or item.get("type") in _MESSAGE_TYPES
|
||||
|
||||
|
||||
def ensure_message_ids(
|
||||
channels: Mapping[str, Any], writes: Sequence[tuple[str, Any]]
|
||||
) -> None:
|
||||
"""Give the messages written to DeltaChannels stable ids.
|
||||
def ensure_message_ids(value: Any) -> None:
|
||||
"""Coerce message-like write values to typed BaseMessages with stable IDs.
|
||||
|
||||
A `DeltaChannel` replays its saved writes on every read, and its reducer
|
||||
may not assign missing ids (`_messages_delta_reducer` doesn't), so a
|
||||
message saved without an id comes back without one, or with a new one on
|
||||
every read, and nothing can target it by id, for example `RemoveMessage`.
|
||||
Called in put_writes() before DeltaChannel writes are submitted to the
|
||||
checkpointer. Without this the checkpoint may store raw dicts or id=None
|
||||
BaseMessages; every get_state() replay then produces a different UUID and
|
||||
the same message appears with a different ID in each LangSmith trace.
|
||||
|
||||
In a channel whose reducer is `_messages_delta_reducer`, every message-like
|
||||
it accepts (a message, a dict, a string, a `(role, content)` tuple; alone,
|
||||
in a list or in an `Overwrite`) becomes a message with an id. Other
|
||||
DeltaChannels may hold other data, so there only message objects get ids,
|
||||
plus message dicts outside an `Overwrite`: one in a list becomes a message.
|
||||
Anything that can't become a message is left as it is, so it fails where
|
||||
it always has: a message dict in a list raises here, anything else when
|
||||
the reducer reads it. An empty id counts as missing.
|
||||
Handles three input shapes:
|
||||
- BaseMessage objects: assign a UUID if id is None.
|
||||
- Dicts with a known "role" (OpenAI-style) or "type" (LangChain format) at
|
||||
the root level: stamp "id" into the dict in-place. The reducer's
|
||||
convert_to_messages call will forward the id to the resulting BaseMessage.
|
||||
- Lists of the above: apply the same logic to each element, replacing dict
|
||||
items with coerced BaseMessages so the shared list reference seen by
|
||||
checkpoint_pending_writes and the background thread both get typed messages.
|
||||
|
||||
A write that changes is replaced in `writes`: lists and dicts the caller
|
||||
passed in are never changed, only the ids of message objects.
|
||||
Mutating synchronously here (before the background thread is submitted) is
|
||||
safe: the serialised bytes always reflect the post-coercion state.
|
||||
"""
|
||||
for i, (ch, value) in enumerate(writes):
|
||||
spec = channels.get(ch)
|
||||
if not isinstance(spec, DeltaChannel):
|
||||
continue
|
||||
if spec.reducer is _messages_delta_reducer:
|
||||
new_value = _with_message_ids(value)
|
||||
else:
|
||||
new_value = _with_base_message_ids(value)
|
||||
if new_value is not value:
|
||||
cast(MutableSequence[tuple[str, Any]], writes)[i] = (ch, new_value)
|
||||
|
||||
|
||||
def _with_message_ids(value: Any) -> Any:
|
||||
is_overwrite, overwrite_value = _get_overwrite(value)
|
||||
if not is_overwrite:
|
||||
if isinstance(value, list):
|
||||
return _map_messages(value, _list_item_with_id)
|
||||
return _message_with_id(value)
|
||||
if isinstance(overwrite_value, (list, tuple)):
|
||||
items = _map_messages(overwrite_value, _message_with_id)
|
||||
if items is not overwrite_value:
|
||||
return Overwrite(items)
|
||||
return value
|
||||
|
||||
|
||||
def _with_base_message_ids(value: Any) -> Any:
|
||||
is_overwrite, overwrite_value = _get_overwrite(value)
|
||||
if is_overwrite:
|
||||
if isinstance(overwrite_value, (list, tuple)):
|
||||
for item in overwrite_value:
|
||||
_stamp_message_id(item)
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return _map_messages(value, _base_list_item_with_id)
|
||||
if isinstance(value, dict) and _is_message_dict(value) and not value.get("id"):
|
||||
return {**value, "id": str(uuid4())}
|
||||
_stamp_message_id(value)
|
||||
return value
|
||||
|
||||
|
||||
def _map_messages(items: list | tuple, fn: Callable[[Any], Any]) -> list | tuple:
|
||||
mapped = [fn(item) for item in items]
|
||||
if all(new is old for new, old in zip(mapped, items)):
|
||||
return items
|
||||
return mapped if isinstance(items, list) else tuple(mapped)
|
||||
|
||||
|
||||
def _list_item_with_id(item: Any) -> Any:
|
||||
if isinstance(item, dict) and _is_message_dict(item):
|
||||
return _message_with_id(item, strict=True)
|
||||
return _message_with_id(item)
|
||||
|
||||
|
||||
def _base_list_item_with_id(item: Any) -> Any:
|
||||
if isinstance(item, dict) and _is_message_dict(item):
|
||||
return _message_with_id(item, strict=True)
|
||||
_stamp_message_id(item)
|
||||
return item
|
||||
|
||||
|
||||
def _message_with_id(item: Any, *, strict: bool = False) -> Any:
|
||||
if isinstance(item, BaseMessage):
|
||||
_stamp_message_id(item)
|
||||
return item
|
||||
try:
|
||||
message = convert_to_messages([item])[0]
|
||||
except Exception:
|
||||
if strict:
|
||||
raise
|
||||
return item
|
||||
if not message.id:
|
||||
message.id = str(uuid4())
|
||||
return message
|
||||
|
||||
|
||||
def _stamp_message_id(item: Any) -> None:
|
||||
if isinstance(item, BaseMessage) and item.id is None:
|
||||
item.id = str(uuid4())
|
||||
if isinstance(value, BaseMessage):
|
||||
if value.id is None:
|
||||
value.id = str(uuid4())
|
||||
elif isinstance(value, dict) and _is_message_dict(value):
|
||||
if not value.get("id"):
|
||||
value["id"] = str(uuid4())
|
||||
elif isinstance(value, list):
|
||||
for i, item in enumerate(value):
|
||||
if isinstance(item, BaseMessage):
|
||||
if item.id is None:
|
||||
item.id = str(uuid4())
|
||||
elif isinstance(item, dict) and _is_message_dict(item):
|
||||
msg = convert_to_messages([item])[0]
|
||||
if msg.id is None:
|
||||
msg.id = str(uuid4())
|
||||
value[i] = msg
|
||||
|
||||
@@ -7,7 +7,7 @@ import sys
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
@@ -686,6 +686,8 @@ async def arun_with_retry(
|
||||
task: PregelExecutableTask,
|
||||
retry_policy: Sequence[RetryPolicy] | None,
|
||||
stream: bool = False,
|
||||
match_cached_writes: Callable[[], Awaitable[Sequence[PregelExecutableTask]]]
|
||||
| None = None,
|
||||
configurable: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Run a task asynchronously with retries."""
|
||||
@@ -709,6 +711,11 @@ async def arun_with_retry(
|
||||
)
|
||||
},
|
||||
)
|
||||
if match_cached_writes is not None and task.cache_key is not None:
|
||||
for t in await match_cached_writes():
|
||||
if t is task:
|
||||
# if the task is already cached, return
|
||||
return
|
||||
while True:
|
||||
runtime = config.get(CONF, {}).get(CONFIG_KEY_RUNTIME)
|
||||
if isinstance(runtime, Runtime):
|
||||
|
||||
@@ -137,7 +137,6 @@ from langgraph.pregel._checkpoint import (
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
create_checkpoint_plan_for_update_state_api,
|
||||
delta_channels_overwritten,
|
||||
delta_channels_with_pending_writes,
|
||||
empty_checkpoint,
|
||||
get_updated_channels_from_tasks,
|
||||
@@ -153,7 +152,6 @@ from langgraph.pregel._loop import (
|
||||
from langgraph.pregel._messages import (
|
||||
StreamMessagesHandler,
|
||||
StreamMessagesHandlerV2,
|
||||
ensure_message_ids,
|
||||
)
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._retry import RetryPolicy
|
||||
@@ -1768,7 +1766,6 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -1791,22 +1788,6 @@ class Pregel(
|
||||
)
|
||||
|
||||
if input_writes := deque(map_input(self.input_channels, values)):
|
||||
_store_or_fork_delta_writes(
|
||||
checkpointer,
|
||||
config,
|
||||
saved,
|
||||
checkpoint_config,
|
||||
self.channels,
|
||||
[
|
||||
(
|
||||
str(uuid5(UUID(checkpoint["id"]), INPUT)),
|
||||
input_writes,
|
||||
None,
|
||||
)
|
||||
],
|
||||
fork_pending,
|
||||
is_first=is_first,
|
||||
)
|
||||
updated_channels = apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
@@ -1814,9 +1795,6 @@ class Pregel(
|
||||
checkpointer.get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, input_writes
|
||||
)
|
||||
|
||||
# apply input write to channels
|
||||
next_step = (
|
||||
@@ -1844,7 +1822,6 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -1856,6 +1833,13 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
|
||||
# store the writes
|
||||
checkpointer.put_writes(
|
||||
next_config,
|
||||
input_writes,
|
||||
str(uuid5(UUID(checkpoint["id"]), INPUT)),
|
||||
)
|
||||
|
||||
return patch_checkpoint_map(
|
||||
next_config, saved.metadata if saved else None
|
||||
)
|
||||
@@ -2049,22 +2033,30 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
updated_channels = get_updated_channels_from_tasks(run_tasks)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, (w for t in run_tasks for w in t.writes)
|
||||
)
|
||||
_store_or_fork_delta_writes(
|
||||
checkpointer,
|
||||
config,
|
||||
saved,
|
||||
checkpoint_config,
|
||||
self.channels,
|
||||
[
|
||||
(task_id, task.writes, task)
|
||||
for task_id, task in zip(run_task_ids, run_tasks)
|
||||
],
|
||||
fork_pending,
|
||||
is_first=is_first,
|
||||
)
|
||||
# The base's other children replay whatever is stored on it, so an
|
||||
# edit of an older checkpoint stores none of its writes there: the
|
||||
# checkpoint written here carries them, its delta channels
|
||||
# snapshotted. Later supersteps address the checkpoint just written.
|
||||
if (
|
||||
is_first
|
||||
and saved is not None
|
||||
and checkpoint_superseded(checkpointer, config, saved)
|
||||
):
|
||||
fork_pending.update(
|
||||
ch
|
||||
for ch in updated_channels
|
||||
if isinstance(self.channels.get(ch), DeltaChannel)
|
||||
)
|
||||
elif saved is not None:
|
||||
for task_id, task in zip(run_task_ids, run_tasks):
|
||||
channel_writes = [w for w in task.writes if w[0] != PUSH]
|
||||
if channel_writes:
|
||||
checkpointer.put_writes(
|
||||
checkpoint_config,
|
||||
channel_writes,
|
||||
task_id,
|
||||
**_task_path_kwarg(checkpointer.put_writes, task),
|
||||
)
|
||||
apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
@@ -2094,7 +2086,6 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -2267,7 +2258,6 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
@@ -2290,22 +2280,6 @@ class Pregel(
|
||||
)
|
||||
|
||||
if input_writes := deque(map_input(self.input_channels, values)):
|
||||
await _astore_or_fork_delta_writes(
|
||||
checkpointer,
|
||||
config,
|
||||
saved,
|
||||
checkpoint_config,
|
||||
self.channels,
|
||||
[
|
||||
(
|
||||
str(uuid5(UUID(checkpoint["id"]), INPUT)),
|
||||
input_writes,
|
||||
None,
|
||||
)
|
||||
],
|
||||
fork_pending,
|
||||
is_first=is_first,
|
||||
)
|
||||
updated_channels = apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
@@ -2313,9 +2287,6 @@ class Pregel(
|
||||
checkpointer.get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, input_writes
|
||||
)
|
||||
|
||||
# apply input write to channels
|
||||
next_step = (
|
||||
@@ -2343,7 +2314,6 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
@@ -2355,6 +2325,13 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
|
||||
# store the writes
|
||||
await checkpointer.aput_writes(
|
||||
next_config,
|
||||
input_writes,
|
||||
str(uuid5(UUID(checkpoint["id"]), INPUT)),
|
||||
)
|
||||
|
||||
return patch_checkpoint_map(
|
||||
next_config, saved.metadata if saved else None
|
||||
)
|
||||
@@ -2547,22 +2524,30 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
updated_channels = get_updated_channels_from_tasks(run_tasks)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, (w for t in run_tasks for w in t.writes)
|
||||
)
|
||||
await _astore_or_fork_delta_writes(
|
||||
checkpointer,
|
||||
config,
|
||||
saved,
|
||||
checkpoint_config,
|
||||
self.channels,
|
||||
[
|
||||
(task_id, task.writes, task)
|
||||
for task_id, task in zip(run_task_ids, run_tasks)
|
||||
],
|
||||
fork_pending,
|
||||
is_first=is_first,
|
||||
)
|
||||
# The base's other children replay whatever is stored on it, so an
|
||||
# edit of an older checkpoint stores none of its writes there: the
|
||||
# checkpoint written here carries them, its delta channels
|
||||
# snapshotted. Later supersteps address the checkpoint just written.
|
||||
if (
|
||||
is_first
|
||||
and saved is not None
|
||||
and await acheckpoint_superseded(checkpointer, config, saved)
|
||||
):
|
||||
fork_pending.update(
|
||||
ch
|
||||
for ch in updated_channels
|
||||
if isinstance(self.channels.get(ch), DeltaChannel)
|
||||
)
|
||||
elif saved is not None:
|
||||
for task_id, task in zip(run_task_ids, run_tasks):
|
||||
channel_writes = [w for w in task.writes if w[0] != PUSH]
|
||||
if channel_writes:
|
||||
await checkpointer.aput_writes(
|
||||
checkpoint_config,
|
||||
channel_writes,
|
||||
task_id,
|
||||
**_task_path_kwarg(checkpointer.aput_writes, task),
|
||||
)
|
||||
apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
@@ -2592,7 +2577,6 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
@@ -3054,7 +3038,7 @@ class Pregel(
|
||||
# with channel updates applied only at the transition between steps.
|
||||
while loop.tick():
|
||||
for task in loop.match_cached_writes():
|
||||
loop.put_writes(task.id, task.writes, cached=True)
|
||||
loop.output_writes(task.id, task.writes, cached=True)
|
||||
for _ in runner.tick(
|
||||
[t for t in loop.tasks.values() if not t.writes],
|
||||
timeout=self.step_timeout,
|
||||
@@ -3525,7 +3509,7 @@ class Pregel(
|
||||
try:
|
||||
while loop.tick():
|
||||
for task in await loop.amatch_cached_writes():
|
||||
loop.put_writes(task.id, task.writes, cached=True)
|
||||
loop.output_writes(task.id, task.writes, cached=True)
|
||||
async for _ in runner.atick(
|
||||
[t for t in loop.tasks.values() if not t.writes],
|
||||
timeout=self.step_timeout,
|
||||
@@ -4271,87 +4255,6 @@ def _task_path_kwarg(put_writes: Callable[..., Any], task: PregelTaskWrites) ->
|
||||
return {"task_path": task_path_str(task.path)}
|
||||
|
||||
|
||||
_UpdateWrites = Sequence[tuple[str, Sequence[tuple[str, Any]], PregelTaskWrites | None]]
|
||||
|
||||
|
||||
def _store_or_fork_delta_writes(
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
checkpoint_config: RunnableConfig,
|
||||
channels: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
writes: _UpdateWrites,
|
||||
fork_pending: set[str],
|
||||
*,
|
||||
is_first: bool,
|
||||
) -> None:
|
||||
"""Save an update's writes where its DeltaChannels will read them.
|
||||
|
||||
A DeltaChannel rebuilds its value from the writes saved on a checkpoint's
|
||||
ancestors, so the writes go on the checkpoint the update builds on. If the
|
||||
thread already moved past that checkpoint, its other children would read
|
||||
them too, so the new checkpoint snapshots those channels instead. `writes`
|
||||
holds `(task_id, writes, task)` per task; `task` is `None` for input. A
|
||||
task's `PUSH` writes are left out: the update saves them on the checkpoint
|
||||
it creates. The messages in `writes` get their ids first, in place, as the
|
||||
loop's `put_writes` does, so pass the writes the update then applies.
|
||||
"""
|
||||
for _, task_writes, _ in writes:
|
||||
ensure_message_ids(channels, task_writes)
|
||||
if saved is None:
|
||||
return
|
||||
if is_first and checkpoint_superseded(checkpointer, config, saved):
|
||||
fork_pending.update(
|
||||
ch
|
||||
for _, task_writes, _ in writes
|
||||
for ch, _ in task_writes
|
||||
if isinstance(channels.get(ch), DeltaChannel)
|
||||
)
|
||||
return
|
||||
for task_id, task_writes, task in writes:
|
||||
if channel_writes := [w for w in task_writes if w[0] != PUSH]:
|
||||
checkpointer.put_writes(
|
||||
checkpoint_config,
|
||||
channel_writes,
|
||||
task_id,
|
||||
**(_task_path_kwarg(checkpointer.put_writes, task) if task else {}),
|
||||
)
|
||||
|
||||
|
||||
async def _astore_or_fork_delta_writes(
|
||||
checkpointer: BaseCheckpointSaver,
|
||||
config: RunnableConfig,
|
||||
saved: CheckpointTuple | None,
|
||||
checkpoint_config: RunnableConfig,
|
||||
channels: Mapping[str, BaseChannel | ManagedValueSpec],
|
||||
writes: _UpdateWrites,
|
||||
fork_pending: set[str],
|
||||
*,
|
||||
is_first: bool,
|
||||
) -> None:
|
||||
"""Async `_store_or_fork_delta_writes`."""
|
||||
for _, task_writes, _ in writes:
|
||||
ensure_message_ids(channels, task_writes)
|
||||
if saved is None:
|
||||
return
|
||||
if is_first and await acheckpoint_superseded(checkpointer, config, saved):
|
||||
fork_pending.update(
|
||||
ch
|
||||
for _, task_writes, _ in writes
|
||||
for ch, _ in task_writes
|
||||
if isinstance(channels.get(ch), DeltaChannel)
|
||||
)
|
||||
return
|
||||
for task_id, task_writes, task in writes:
|
||||
if channel_writes := [w for w in task_writes if w[0] != PUSH]:
|
||||
await checkpointer.aput_writes(
|
||||
checkpoint_config,
|
||||
channel_writes,
|
||||
task_id,
|
||||
**(_task_path_kwarg(checkpointer.aput_writes, task) if task else {}),
|
||||
)
|
||||
|
||||
|
||||
def _update_task_id(checkpoint_id: str, i: int) -> str:
|
||||
"""Task id for the `i`th update of a superstep that has no task to reuse.
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ from warnings import warn
|
||||
from langchain_core.messages import AnyMessage
|
||||
from langchain_core.runnables import Runnable, RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import (
|
||||
NotRequired,
|
||||
TypeAliasType,
|
||||
@@ -882,16 +882,6 @@ class Command(Generic[N], ToolOutputMixin):
|
||||
PARENT: ClassVar[Literal["__parent__"]] = "__parent__"
|
||||
|
||||
|
||||
def _validate_resume(adapter: TypeAdapter[Any], value: Any) -> Any:
|
||||
from langgraph.errors import _mark_invalid_resume
|
||||
|
||||
try:
|
||||
return adapter.validate_python(value)
|
||||
except ValidationError as exc:
|
||||
_mark_invalid_resume(exc)
|
||||
raise
|
||||
|
||||
|
||||
@overload
|
||||
def interrupt(value: Any, *, response_schema: type[ResponseT]) -> ResponseT: ...
|
||||
|
||||
@@ -999,7 +989,6 @@ def interrupt(
|
||||
Raises:
|
||||
GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client.
|
||||
pydantic.ValidationError: When a resume value does not match a Pydantic model, `TypedDict`, or dataclass `response_schema`.
|
||||
Nothing is saved, so the interrupt can be answered again. `is_invalid_resume` identifies it.
|
||||
"""
|
||||
from langgraph._internal._constants import (
|
||||
CONFIG_KEY_CHECKPOINT_NS,
|
||||
@@ -1023,14 +1012,14 @@ def interrupt(
|
||||
if scratchpad.resume:
|
||||
if idx < len(scratchpad.resume):
|
||||
v = scratchpad.resume[idx]
|
||||
validated = _validate_resume(adapter, v) if adapter else v
|
||||
validated = adapter.validate_python(v) if adapter else v
|
||||
conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume[: idx + 1])])
|
||||
return validated
|
||||
# find current resume value
|
||||
v = scratchpad.get_null_resume(True)
|
||||
if v is not None:
|
||||
assert len(scratchpad.resume) == idx, (scratchpad.resume, idx)
|
||||
validated = _validate_resume(adapter, v) if adapter else v
|
||||
validated = adapter.validate_python(v) if adapter else v
|
||||
scratchpad.resume.append(v)
|
||||
conf[CONFIG_KEY_SEND]([(RESUME, scratchpad.resume)])
|
||||
return validated
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.2.15"
|
||||
version = "1.2.14"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -625,33 +625,3 @@ def test_exit_resume_retried_after_its_final_checkpoint_fails_reruns_the_resumed
|
||||
state = graph.get_state(config)
|
||||
assert state.values["log"] == state.values["plain"]
|
||||
assert sorted(state.values["plain"]) == ["ask", "done", "in"]
|
||||
|
||||
|
||||
class _TurnState(_ResumeState):
|
||||
turn: int
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exit_runs", [1, 2])
|
||||
def test_exit_runs_from_a_checkpoint_the_thread_moved_past_leave_its_newer_branch_alone(
|
||||
sync_checkpointer: BaseCheckpointSaver, exit_runs: int
|
||||
) -> None:
|
||||
builder = StateGraph(_TurnState)
|
||||
builder.add_node("a", lambda state: _both(f"a{state['turn']}"))
|
||||
builder.add_node("b", lambda state: _both(f"b{state['turn']}"))
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge("a", "b")
|
||||
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"turn": 1, **_both("in1")}, config)
|
||||
end_of_turn_1 = graph.get_state(config).config
|
||||
graph.invoke({"turn": 2, **_both("in2")}, config)
|
||||
|
||||
for turn in range(3, 3 + exit_runs):
|
||||
graph.invoke(
|
||||
{"turn": turn, **_both(f"in{turn}")}, end_of_turn_1, durability="exit"
|
||||
)
|
||||
|
||||
history = list(graph.get_state_history(config))
|
||||
assert [s.values.get("log", []) for s in history] == [
|
||||
s.values.get("plain", []) for s in history
|
||||
]
|
||||
|
||||
@@ -16,11 +16,7 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Command, Durability, Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -38,19 +34,18 @@ def _append_reducer(
|
||||
return result
|
||||
|
||||
|
||||
def _build_graph(
|
||||
checkpointer: Any,
|
||||
*,
|
||||
reducer: Any = _append_reducer,
|
||||
write: Any = lambda: [AIMessage(content="reply", id="ai-1")],
|
||||
) -> Any:
|
||||
def _build_graph(checkpointer: Any) -> Any:
|
||||
State = TypedDict( # noqa: UP013
|
||||
"State",
|
||||
{"messages": Annotated[list, DeltaChannel(reducer, snapshot_frequency=50)]},
|
||||
{
|
||||
"messages": Annotated[
|
||||
list, DeltaChannel(_append_reducer, snapshot_frequency=50)
|
||||
]
|
||||
},
|
||||
) # type: ignore[call-overload]
|
||||
|
||||
def agent(state: dict) -> dict:
|
||||
return {"messages": write()}
|
||||
return {"messages": [AIMessage(content="reply", id="ai-1")]}
|
||||
|
||||
return (
|
||||
StateGraph(State)
|
||||
@@ -143,193 +138,3 @@ def test_delta_channel_dict_style_message_gets_stable_id() -> None:
|
||||
assert len(set(ids)) == 1, (
|
||||
f"dict-style HumanMessage id must be stable across get_state() calls; got {ids}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("container", [list, tuple])
|
||||
def test_delta_channel_message_in_an_overwrite_gets_stable_id(container: type) -> None:
|
||||
graph = _build_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "overwrite-id-stability"}}
|
||||
|
||||
graph.invoke(
|
||||
{"messages": Overwrite(container([HumanMessage(content="hello")]))}, config
|
||||
)
|
||||
|
||||
ids = [
|
||||
next(
|
||||
m.id
|
||||
for m in graph.get_state(config).values["messages"]
|
||||
if isinstance(m, HumanMessage)
|
||||
)
|
||||
for _ in range(3)
|
||||
]
|
||||
|
||||
assert ids[0] is not None, "a message inside an Overwrite should get a UUID"
|
||||
assert len(set(ids)) == 1, (
|
||||
f"its id must be stable across get_state() calls; got {ids}"
|
||||
)
|
||||
|
||||
|
||||
# Writes of "hello" in every shape `_messages_delta_reducer` accepts.
|
||||
MESSAGE_SHAPES = [
|
||||
pytest.param(lambda: [HumanMessage("hello")], id="message"),
|
||||
pytest.param(lambda: HumanMessage("hello"), id="bare message"),
|
||||
pytest.param(lambda: [{"role": "user", "content": "hello"}], id="dict"),
|
||||
pytest.param(lambda: {"type": "human", "content": "hello"}, id="bare dict"),
|
||||
pytest.param(
|
||||
lambda: [{"role": "user", "content": "hello", "id": ""}], id="dict, empty id"
|
||||
),
|
||||
pytest.param(lambda: [("user", "hello")], id="tuple"),
|
||||
pytest.param(lambda: ("user", "hello"), id="bare tuple"),
|
||||
pytest.param(lambda: ["hello"], id="string"),
|
||||
pytest.param(lambda: "hello", id="bare string"),
|
||||
pytest.param(lambda: Overwrite(["hello"]), id="Overwrite"),
|
||||
pytest.param(lambda: Overwrite(("hello",)), id="Overwrite of a tuple"),
|
||||
pytest.param(lambda: {"__overwrite__": [("user", "hello")]}, id="Overwrite dict"),
|
||||
]
|
||||
|
||||
|
||||
def _hello_ids(graph: Any, config: dict) -> list[str | None]:
|
||||
return [
|
||||
next(
|
||||
m.id
|
||||
for m in graph.get_state(config).values["messages"]
|
||||
if m.content == "hello"
|
||||
)
|
||||
for _ in range(2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
|
||||
def test_input_message_of_any_shape_gets_stable_id(write: Any) -> None:
|
||||
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"messages": write()}, config)
|
||||
|
||||
first, second = _hello_ids(graph, config)
|
||||
assert first, "the message should get a UUID"
|
||||
assert first == second, "its id must be the same on every read"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
|
||||
def test_node_message_of_any_shape_gets_stable_id(write: Any) -> None:
|
||||
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer, write=write)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"messages": []}, config)
|
||||
|
||||
first, second = _hello_ids(graph, config)
|
||||
assert first, "the message should get a UUID"
|
||||
assert first == second, "its id must be the same on every read"
|
||||
|
||||
|
||||
def test_a_list_written_twice_adds_its_messages_twice() -> None:
|
||||
turn = [("user", "hello")]
|
||||
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"messages": turn}, config)
|
||||
graph.invoke({"messages": turn}, config)
|
||||
|
||||
hellos = [
|
||||
m.id for m in graph.get_state(config).values["messages"] if m.content == "hello"
|
||||
]
|
||||
assert len(set(hellos)) == 2, f"each write should add a message; got {hellos}"
|
||||
assert turn == [("user", "hello")], "the caller's list should be left as it was"
|
||||
|
||||
|
||||
def test_a_forked_command_update_gets_a_stable_id() -> None:
|
||||
graph = _build_graph(InMemorySaver(), reducer=_messages_delta_reducer)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
|
||||
before_agent = next(
|
||||
s.config for s in graph.get_state_history(config) if s.next == ("agent",)
|
||||
)
|
||||
|
||||
graph.invoke(Command(update={"messages": "hello"}), before_agent)
|
||||
|
||||
first, second = _hello_ids(graph, config)
|
||||
assert first, "the message should get a UUID"
|
||||
assert first == second, "its id must be the same on every read"
|
||||
|
||||
|
||||
def test_a_message_dict_the_reducer_rejects_fails_the_run_before_it_is_saved() -> None:
|
||||
graph = _build_graph(
|
||||
InMemorySaver(),
|
||||
reducer=_messages_delta_reducer,
|
||||
write=lambda: [{"role": "tool", "content": "x"}],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
|
||||
|
||||
assert [m.id for m in graph.get_state(config).values["messages"]] == ["a"]
|
||||
|
||||
|
||||
def _pregel_graph(seen: list[list[str | None]]) -> Any:
|
||||
def look(state: dict) -> list:
|
||||
seen.append([m.id for m in state["messages"]])
|
||||
return []
|
||||
|
||||
node = (
|
||||
NodeBuilder()
|
||||
.subscribe_to("go", read=False)
|
||||
.read_from("messages")
|
||||
.do(look)
|
||||
.write_to("messages")
|
||||
)
|
||||
return Pregel(
|
||||
nodes={"look": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
"go": LastValue(int),
|
||||
},
|
||||
input_channels=["messages", "go"],
|
||||
output_channels=["messages"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
def test_pregel_input_message_gets_stable_id_on_a_new_thread(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _pregel_graph([])
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke(
|
||||
{"messages": [HumanMessage("hello")], "go": 1}, config, durability=durability
|
||||
)
|
||||
|
||||
first, second = _hello_ids(graph, config)
|
||||
assert first, "the message should get a UUID"
|
||||
assert first == second, "its id must be the same on every read"
|
||||
|
||||
|
||||
def test_a_node_reads_the_pregel_input_message_ids_that_are_saved() -> None:
|
||||
seen: list[list[str | None]] = []
|
||||
graph = _pregel_graph(seen)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage("first", id="first")], "go": 1}, config)
|
||||
|
||||
graph.invoke({"messages": [{"role": "user", "content": "hello"}], "go": 1}, config)
|
||||
|
||||
saved = [m.id for m in graph.get_state(config).values["messages"]]
|
||||
assert None not in saved
|
||||
assert seen[-1] == saved, "the node must see the ids that are saved"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[
|
||||
pytest.param({"type": "tool", "name": "search"}, id="not a message"),
|
||||
pytest.param({"role": "user", "content": "x"}, id="message-shaped"),
|
||||
],
|
||||
)
|
||||
def test_other_delta_channels_keep_the_dicts_in_an_overwrite(value: dict) -> None:
|
||||
graph = _build_graph(InMemorySaver(), write=lambda: Overwrite([dict(value)]))
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"messages": []}, config)
|
||||
|
||||
assert graph.get_state(config).values["messages"] == [value]
|
||||
|
||||
@@ -1,122 +0,0 @@
|
||||
"""A node served from the node cache saves its writes like a node that ran,
|
||||
without writing them back to the cache, and a node that fails isn't cached."""
|
||||
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langgraph.cache.memory import InMemoryCache
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.types import CachePolicy, Durability
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
INPUT = {"log": [], "plain": []}
|
||||
|
||||
|
||||
def _append(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
log: Annotated[list, DeltaChannel(_append)]
|
||||
plain: Annotated[list, operator.add]
|
||||
|
||||
|
||||
class _CountsSets(InMemoryCache):
|
||||
sets = 0
|
||||
|
||||
def set(self, keys: Any) -> None:
|
||||
self.sets += 1
|
||||
super().set(keys)
|
||||
|
||||
|
||||
def _a_then_cached_b_then_c(runs: list[str], cache: InMemoryCache) -> Any:
|
||||
def node(name: str) -> Any:
|
||||
def run(state: _State) -> dict:
|
||||
runs.append(name)
|
||||
return {"log": [name], "plain": [name]}
|
||||
|
||||
return run
|
||||
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("a", node("a"))
|
||||
builder.add_node("b", node("b"), cache_policy=CachePolicy())
|
||||
builder.add_node("c", node("c"))
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge("a", "b")
|
||||
builder.add_edge("b", "c")
|
||||
return builder.compile(checkpointer=InMemorySaver(), cache=cache)
|
||||
|
||||
|
||||
def test_a_cache_hit_saves_its_writes_without_caching_them_again(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
runs: list[str] = []
|
||||
cache = _CountsSets()
|
||||
graph = _a_then_cached_b_then_c(runs, cache)
|
||||
graph.invoke(INPUT, {"configurable": {"thread_id": "1"}}, durability=durability)
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
|
||||
graph.invoke(INPUT, config, durability=durability)
|
||||
|
||||
assert runs == ["a", "b", "c", "a", "c"]
|
||||
assert cache.sets == 1, "the cache hit was written back to the cache"
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
async def test_a_cache_hit_saves_its_writes_without_caching_them_again_async(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
runs: list[str] = []
|
||||
cache = _CountsSets()
|
||||
graph = _a_then_cached_b_then_c(runs, cache)
|
||||
await graph.ainvoke(
|
||||
INPUT, {"configurable": {"thread_id": "1"}}, durability=durability
|
||||
)
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
|
||||
await graph.ainvoke(INPUT, config, durability=durability)
|
||||
|
||||
assert runs == ["a", "b", "c", "a", "c"]
|
||||
assert cache.sets == 1, "the cache hit was written back to the cache"
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
def _cached_node_that_fails(runs: list[str]) -> Any:
|
||||
def fail(state: _State) -> dict:
|
||||
runs.append("b")
|
||||
raise ValueError("b failed")
|
||||
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("b", fail, cache_policy=CachePolicy())
|
||||
builder.add_edge(START, "b")
|
||||
return builder.compile(checkpointer=InMemorySaver(), cache=InMemoryCache())
|
||||
|
||||
|
||||
def test_a_node_that_fails_is_not_cached() -> None:
|
||||
runs: list[str] = []
|
||||
graph = _cached_node_that_fails(runs)
|
||||
|
||||
for thread in ("1", "2"):
|
||||
with pytest.raises(ValueError, match="b failed"):
|
||||
graph.invoke(INPUT, {"configurable": {"thread_id": thread}})
|
||||
|
||||
assert runs == ["b", "b"]
|
||||
|
||||
|
||||
async def test_a_node_that_fails_is_not_cached_async() -> None:
|
||||
runs: list[str] = []
|
||||
graph = _cached_node_that_fails(runs)
|
||||
|
||||
for thread in ("1", "2"):
|
||||
with pytest.raises(ValueError, match="b failed"):
|
||||
await graph.ainvoke(INPUT, {"configurable": {"thread_id": thread}})
|
||||
|
||||
assert runs == ["b", "b"]
|
||||
@@ -1,103 +0,0 @@
|
||||
"""A run's input to a DeltaChannel input channel reads back on the checkpoints
|
||||
built from it, and on no others."""
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Durability
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _sorted_extend(current: list, writes: list) -> list:
|
||||
return sorted([*current, *(item for write in writes for item in write)])
|
||||
|
||||
|
||||
def _delta_input_graph() -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend),
|
||||
"plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
|
||||
"go": LastValue(int),
|
||||
},
|
||||
input_channels=["log", "plain", "go"],
|
||||
output_channels=["log", "plain"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
def test_each_run_input_reads_back_on_its_own_checkpoints(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"log": [0], "plain": [0], "go": 1}, config, durability=durability)
|
||||
graph.invoke({"log": [5], "plain": [5], "go": 1}, config, durability=durability)
|
||||
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
async def test_each_run_input_reads_back_on_its_own_checkpoints_async(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
await graph.ainvoke(
|
||||
{"log": [0], "plain": [0], "go": 1}, config, durability=durability
|
||||
)
|
||||
await graph.ainvoke(
|
||||
{"log": [5], "plain": [5], "go": 1}, config, durability=durability
|
||||
)
|
||||
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
OTHER_BRANCH_INPUTS = pytest.mark.parametrize(
|
||||
"other_branch_input",
|
||||
[{"go": 1}, {"log": [5], "plain": [5], "go": 1}],
|
||||
ids=["other-branch-without-delta-input", "other-branch-with-delta-input"],
|
||||
)
|
||||
|
||||
|
||||
@OTHER_BRANCH_INPUTS
|
||||
def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch(
|
||||
durability: Durability, other_branch_input: dict
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"go": 1}, config, durability=durability)
|
||||
older = graph.get_state(config).config
|
||||
graph.invoke(other_branch_input, config, durability=durability)
|
||||
|
||||
graph.invoke({"log": [7], "plain": [7], "go": 1}, older, durability=durability)
|
||||
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
@OTHER_BRANCH_INPUTS
|
||||
async def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch_async(
|
||||
durability: Durability, other_branch_input: dict
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"go": 1}, config, durability=durability)
|
||||
older = (await graph.aget_state(config)).config
|
||||
await graph.ainvoke(other_branch_input, config, durability=durability)
|
||||
|
||||
await graph.ainvoke(
|
||||
{"log": [7], "plain": [7], "go": 1}, older, durability=durability
|
||||
)
|
||||
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
@@ -1,78 +0,0 @@
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
CONFIG = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
|
||||
def _extend(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
def _graph(reads: list) -> Pregel:
|
||||
writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
|
||||
reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
|
||||
return Pregel(
|
||||
nodes={"writer": writer, "reader": reader},
|
||||
channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
|
||||
input_channels="a",
|
||||
output_channels=["d"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
def test_an_input_update_from_before_the_first_write_starts_only_the_writer() -> None:
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
graph.invoke("go", CONFIG)
|
||||
first = next(
|
||||
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
|
||||
)
|
||||
|
||||
fork = graph.update_state(first, {"a": "go"}, as_node="__input__")
|
||||
|
||||
assert graph.get_state(fork).next == ("writer",)
|
||||
graph.invoke(None, fork)
|
||||
assert reads == [[1], [1]]
|
||||
|
||||
|
||||
def test_a_replay_from_before_the_first_write_forks_with_only_the_writer_next() -> None:
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
graph.invoke("go", CONFIG)
|
||||
first = next(
|
||||
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
|
||||
)
|
||||
|
||||
graph.invoke(None, first, durability="sync")
|
||||
|
||||
fork = next(
|
||||
s for s in graph.get_state_history(CONFIG) if s.metadata["source"] == "fork"
|
||||
)
|
||||
assert fork.next == ("writer",)
|
||||
graph.invoke(None, fork.config)
|
||||
assert reads == [[1], [1], [1]]
|
||||
|
||||
|
||||
async def test_an_ainput_update_from_before_the_first_write_starts_only_the_writer() -> (
|
||||
None
|
||||
):
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
await graph.ainvoke("go", CONFIG)
|
||||
first = [
|
||||
s.config
|
||||
async for s in graph.aget_state_history(CONFIG)
|
||||
if s.metadata["step"] == -1
|
||||
][0]
|
||||
|
||||
fork = await graph.aupdate_state(first, {"a": "go"}, as_node="__input__")
|
||||
|
||||
assert (await graph.aget_state(fork)).next == ("writer",)
|
||||
await graph.ainvoke(None, fork)
|
||||
assert reads == [[1], [1]]
|
||||
@@ -1,110 +0,0 @@
|
||||
"""An `Overwrite` through `update_state` snapshots its DeltaChannel on the
|
||||
checkpoint the update saves, as a node's `Overwrite` does on the loop's."""
|
||||
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
|
||||
def _messages_graph(saver: InMemorySaver) -> Any:
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("model", lambda state: {})
|
||||
builder.add_edge(START, "model")
|
||||
return builder.compile(checkpointer=saver)
|
||||
|
||||
|
||||
def _extend(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
def _delta_input_graph(saver: InMemorySaver) -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log")
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={"log": DeltaChannel(_extend), "go": LastValue(int)},
|
||||
input_channels=["log", "go"],
|
||||
output_channels=["log"],
|
||||
checkpointer=saver,
|
||||
)
|
||||
|
||||
|
||||
def test_update_state_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _messages_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="a", id="1")]}, config)
|
||||
|
||||
graph.update_state(
|
||||
config,
|
||||
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
|
||||
as_node="model",
|
||||
)
|
||||
|
||||
head = saver.get_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
|
||||
assert [m.content for m in graph.get_state(config).values["messages"]] == ["b"]
|
||||
|
||||
|
||||
async def test_aupdate_state_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _messages_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"messages": [HumanMessage(content="a", id="1")]}, config)
|
||||
|
||||
await graph.aupdate_state(
|
||||
config,
|
||||
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
|
||||
as_node="model",
|
||||
)
|
||||
|
||||
head = await saver.aget_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
|
||||
values = (await graph.aget_state(config)).values
|
||||
assert [m.content for m in values["messages"]] == ["b"]
|
||||
|
||||
|
||||
def test_update_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _delta_input_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"log": [0], "go": 1}, config)
|
||||
|
||||
graph.update_state(config, {"log": Overwrite([1])}, as_node="__input__")
|
||||
|
||||
head = saver.get_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
|
||||
assert graph.get_state(config).values["log"] == [1]
|
||||
|
||||
|
||||
async def test_aupdate_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _delta_input_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"log": [0], "go": 1}, config)
|
||||
|
||||
await graph.aupdate_state(config, {"log": Overwrite([1])}, as_node="__input__")
|
||||
|
||||
head = await saver.aget_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
|
||||
assert (await graph.aget_state(config)).values["log"] == [1]
|
||||
@@ -25,13 +25,10 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite, StateSnapshot, StateUpdate
|
||||
from langgraph.types import StateSnapshot, StateUpdate
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -541,207 +538,3 @@ def test_update_state_that_snapshots_keeps_a_deferred_node_pending() -> None:
|
||||
|
||||
assert [m.content for m in final["messages"]] == ["s", "a", "u", "b"]
|
||||
assert graph.get_state(config).next == ()
|
||||
|
||||
|
||||
def _sorted_extend(current: list, writes: list) -> list:
|
||||
return sorted([*current, *(item for write in writes for item in write)])
|
||||
|
||||
|
||||
def _delta_input_graph(snapshot_frequency: int = 1000) -> Any:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency),
|
||||
"plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
|
||||
"go": LastValue(int),
|
||||
},
|
||||
input_channels=["log", "plain", "go"],
|
||||
output_channels=["log", "plain"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snapshot_frequency", [1, 2])
|
||||
def test_update_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
|
||||
snapshot_frequency: int,
|
||||
) -> None:
|
||||
graph = _delta_input_graph(snapshot_frequency)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"log": [0], "plain": [0], "go": 1}, config)
|
||||
|
||||
graph.update_state(config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__")
|
||||
after_update = graph.get_state(config).values
|
||||
graph.invoke(None, config)
|
||||
after_run = graph.get_state(config).values
|
||||
|
||||
assert after_update["log"] == after_update["plain"]
|
||||
assert after_run["log"] == after_run["plain"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snapshot_frequency", [1, 2])
|
||||
async def test_aupdate_as_input_reads_back_on_its_checkpoint_and_after_the_next_run(
|
||||
snapshot_frequency: int,
|
||||
) -> None:
|
||||
graph = _delta_input_graph(snapshot_frequency)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"log": [0], "plain": [0], "go": 1}, config)
|
||||
|
||||
await graph.aupdate_state(
|
||||
config, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
|
||||
)
|
||||
after_update = (await graph.aget_state(config)).values
|
||||
await graph.ainvoke(None, config)
|
||||
after_run = (await graph.aget_state(config)).values
|
||||
|
||||
assert after_update["log"] == after_update["plain"]
|
||||
assert after_run["log"] == after_run["plain"]
|
||||
|
||||
|
||||
def test_update_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"go": 1}, config)
|
||||
older = graph.get_state(config).config
|
||||
graph.invoke({"go": 1}, config)
|
||||
other_branch = graph.get_state(config)
|
||||
|
||||
edited = graph.update_state(
|
||||
older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
|
||||
)
|
||||
|
||||
values = graph.get_state(edited).values
|
||||
assert values["log"] == values["plain"]
|
||||
assert graph.get_state(other_branch.config).values == other_branch.values
|
||||
|
||||
|
||||
async def test_aupdate_as_input_to_an_older_checkpoint_stays_out_of_its_other_branch() -> (
|
||||
None
|
||||
):
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"go": 1}, config)
|
||||
older = (await graph.aget_state(config)).config
|
||||
await graph.ainvoke({"go": 1}, config)
|
||||
other_branch = await graph.aget_state(config)
|
||||
|
||||
edited = await graph.aupdate_state(
|
||||
older, {"log": [1], "plain": [1], "go": 1}, as_node="__input__"
|
||||
)
|
||||
|
||||
values = (await graph.aget_state(edited)).values
|
||||
assert values["log"] == values["plain"]
|
||||
assert (await graph.aget_state(other_branch.config)).values == other_branch.values
|
||||
|
||||
|
||||
def _message_ids(graph: Any, config: dict) -> list[str | None]:
|
||||
return [m.id for m in graph.get_state(config).values["messages"]]
|
||||
|
||||
|
||||
async def _amessage_ids(graph: Any, config: dict) -> list[str | None]:
|
||||
return [m.id for m in (await graph.aget_state(config)).values["messages"]]
|
||||
|
||||
|
||||
def _messages_input_graph() -> Any:
|
||||
node = (
|
||||
NodeBuilder()
|
||||
.subscribe_only("go")
|
||||
.do(lambda _: [HumanMessage("n", id="n")])
|
||||
.write_to("messages")
|
||||
)
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
"go": LastValue(int),
|
||||
},
|
||||
input_channels=["messages", "go"],
|
||||
output_channels=["messages"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
# Writes of "b" in shapes `_messages_delta_reducer` accepts.
|
||||
MESSAGE_SHAPES = [
|
||||
pytest.param(lambda: [HumanMessage("b")], id="message"),
|
||||
pytest.param(lambda: ["b"], id="string"),
|
||||
pytest.param(lambda: "b", id="bare string"),
|
||||
pytest.param(lambda: [("user", "b")], id="tuple"),
|
||||
pytest.param(lambda: {"role": "user", "content": "b"}, id="bare dict"),
|
||||
pytest.param(lambda: Overwrite([HumanMessage("b")]), id="Overwrite"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
|
||||
def test_update_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
|
||||
write: Any,
|
||||
) -> None:
|
||||
graph = _build_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
|
||||
|
||||
graph.update_state(config, {"messages": write()})
|
||||
|
||||
first, second = _message_ids(graph, config), _message_ids(graph, config)
|
||||
assert first[-1] is not None
|
||||
assert first == second
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
|
||||
async def test_aupdate_state_gives_a_message_of_any_shape_an_id_that_every_read_keeps(
|
||||
write: Any,
|
||||
) -> None:
|
||||
graph = _build_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"messages": [HumanMessage("a", id="a")]}, config)
|
||||
|
||||
await graph.aupdate_state(config, {"messages": write()})
|
||||
|
||||
first, second = (
|
||||
await _amessage_ids(graph, config),
|
||||
await _amessage_ids(graph, config),
|
||||
)
|
||||
assert first[-1] is not None
|
||||
assert first == second
|
||||
|
||||
|
||||
@pytest.mark.parametrize("write", MESSAGE_SHAPES)
|
||||
def test_update_as_input_gives_a_message_of_any_shape_an_id(write: Any) -> None:
|
||||
graph = _messages_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
|
||||
|
||||
graph.update_state(config, {"messages": write()}, as_node="__input__")
|
||||
|
||||
first, second = _message_ids(graph, config), _message_ids(graph, config)
|
||||
assert first[-1] is not None
|
||||
assert first == second
|
||||
|
||||
|
||||
def test_update_state_on_an_older_checkpoint_gives_a_message_an_id() -> None:
|
||||
graph = _build_graph(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage("a", id="a")]}, config)
|
||||
older = graph.get_state(config).config
|
||||
graph.invoke({"messages": [HumanMessage("c", id="c")]}, config)
|
||||
|
||||
branch = graph.update_state(older, {"messages": [HumanMessage("b")]})
|
||||
|
||||
assert _message_ids(graph, branch)[-1] is not None
|
||||
|
||||
|
||||
async def test_aupdate_as_input_gives_a_message_an_id_that_every_read_keeps() -> None:
|
||||
graph = _messages_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"messages": [HumanMessage("a", id="a")], "go": 1}, config)
|
||||
|
||||
await graph.aupdate_state(
|
||||
config, {"messages": [HumanMessage("b")]}, as_node="__input__"
|
||||
)
|
||||
|
||||
first, second = (
|
||||
await _amessage_ids(graph, config),
|
||||
await _amessage_ids(graph, config),
|
||||
)
|
||||
assert first[-1] is not None
|
||||
assert first == second
|
||||
|
||||
@@ -1,168 +0,0 @@
|
||||
"""A checkpoint must never be saved without the `DeltaChannel` writes it reads."""
|
||||
|
||||
import operator
|
||||
import threading
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.types import Durability
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
INPUT = {"log": [], "plain": []}
|
||||
FINAL = {"log": ["a", "b", "c"], "plain": ["a", "b", "c"]}
|
||||
|
||||
# Exit mode saves nothing before the failed write, so its retry starts over.
|
||||
RETRIES = [
|
||||
pytest.param("sync", None, id="sync"),
|
||||
pytest.param("async", None, id="async"),
|
||||
pytest.param("exit", INPUT, id="exit"),
|
||||
]
|
||||
|
||||
|
||||
def _append(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
log: Annotated[list, DeltaChannel(_append)]
|
||||
plain: Annotated[list, operator.add]
|
||||
|
||||
|
||||
class _FailsTheWriteOfBOnce(InMemorySaver):
|
||||
failed = False
|
||||
|
||||
def _fail_once(self, writes: Any) -> None:
|
||||
if not self.failed and ("log", ["b"]) in writes:
|
||||
self.failed = True
|
||||
raise ConnectionError("b's write was not saved")
|
||||
|
||||
def put_writes(
|
||||
self, config: Any, writes: Any, task_id: str, task_path: str = ""
|
||||
) -> None:
|
||||
self._fail_once(writes)
|
||||
super().put_writes(config, writes, task_id, task_path)
|
||||
|
||||
async def aput_writes(
|
||||
self, config: Any, writes: Any, task_id: str, task_path: str = ""
|
||||
) -> None:
|
||||
self._fail_once(writes)
|
||||
await super().aput_writes(config, writes, task_id, task_path)
|
||||
|
||||
|
||||
class _FailsTheSaveAfterAOnce(InMemorySaver):
|
||||
failed = False
|
||||
|
||||
def _fail_once(self, checkpoint: Any) -> None:
|
||||
if not self.failed and checkpoint["channel_values"].get("plain") == ["a"]:
|
||||
self.failed = True
|
||||
raise ConnectionError("the checkpoint after a was not saved")
|
||||
|
||||
def put(
|
||||
self, config: Any, checkpoint: Any, metadata: Any, new_versions: Any
|
||||
) -> Any:
|
||||
self._fail_once(checkpoint)
|
||||
return super().put(config, checkpoint, metadata, new_versions)
|
||||
|
||||
async def aput(
|
||||
self, config: Any, checkpoint: Any, metadata: Any, new_versions: Any
|
||||
) -> Any:
|
||||
self._fail_once(checkpoint)
|
||||
return await super().aput(config, checkpoint, metadata, new_versions)
|
||||
|
||||
|
||||
def _a_then_b_then_c(saver: InMemorySaver) -> Any:
|
||||
builder = StateGraph(_State)
|
||||
for name in "abc":
|
||||
builder.add_node(
|
||||
name, lambda state, name=name: {"log": [name], "plain": [name]}
|
||||
)
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge("a", "b")
|
||||
builder.add_edge("b", "c")
|
||||
return builder.compile(checkpointer=saver)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("durability", "retry_input"), RETRIES)
|
||||
def test_a_failed_delta_write_is_rerun_not_lost(
|
||||
durability: Durability, retry_input: dict | None
|
||||
) -> None:
|
||||
graph = _a_then_b_then_c(_FailsTheWriteOfBOnce())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
graph.invoke(INPUT, config, durability=durability)
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
graph.invoke(retry_input, config, durability=durability)
|
||||
assert graph.get_state(config).values == FINAL
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("durability", "retry_input"), RETRIES)
|
||||
async def test_a_failed_delta_write_is_rerun_not_lost_async(
|
||||
durability: Durability, retry_input: dict | None
|
||||
) -> None:
|
||||
graph = _a_then_b_then_c(_FailsTheWriteOfBOnce())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
await graph.ainvoke(INPUT, config, durability=durability)
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
await graph.ainvoke(retry_input, config, durability=durability)
|
||||
assert (await graph.aget_state(config)).values == FINAL
|
||||
|
||||
|
||||
@pytest.mark.parametrize("durability", ["sync", "async"])
|
||||
def test_a_failed_checkpoint_save_is_rerun_not_built_on(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _a_then_b_then_c(_FailsTheSaveAfterAOnce())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
graph.invoke(INPUT, config, durability=durability)
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
graph.invoke(None, config, durability=durability)
|
||||
assert graph.get_state(config).values == FINAL
|
||||
|
||||
|
||||
@pytest.mark.parametrize("durability", ["sync", "async"])
|
||||
async def test_a_failed_checkpoint_save_is_rerun_not_built_on_async(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _a_then_b_then_c(_FailsTheSaveAfterAOnce())
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
with pytest.raises(ConnectionError):
|
||||
await graph.ainvoke(INPUT, config, durability=durability)
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
await graph.ainvoke(None, config, durability=durability)
|
||||
assert (await graph.aget_state(config)).values == FINAL
|
||||
|
||||
|
||||
def test_a_delta_graph_finishes_on_a_single_background_thread() -> None:
|
||||
graph = _a_then_b_then_c(InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "t"}, "max_concurrency": 1}
|
||||
result: dict = {}
|
||||
run = threading.Thread(
|
||||
target=lambda: result.update(graph.invoke(INPUT, config, durability="async")),
|
||||
daemon=True,
|
||||
)
|
||||
|
||||
run.start()
|
||||
run.join(timeout=10)
|
||||
|
||||
assert not run.is_alive(), "invoke hung"
|
||||
assert result == FINAL
|
||||
@@ -6,7 +6,6 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.errors import is_invalid_resume
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.types import Command, Durability, Interrupt, interrupt
|
||||
from tests.any_str import AnyStr
|
||||
@@ -203,22 +202,14 @@ def test_interrupt_response_schema_rejects_invalid_resume(
|
||||
def resume(value: dict[str, Any]) -> Command:
|
||||
return Command(resume=value if resume_style == "null" else {pending.id: value})
|
||||
|
||||
with pytest.raises(ValidationError, match="approved") as exc_info:
|
||||
with pytest.raises(ValidationError, match="approved"):
|
||||
graph.invoke(resume({"approved": "nope"}), config)
|
||||
assert is_invalid_resume(exc_info.value)
|
||||
|
||||
assert graph.invoke(resume({"approved": False}), config) == {
|
||||
"answer": Decision(approved=False)
|
||||
}
|
||||
|
||||
|
||||
def test_is_invalid_resume_ignores_other_errors() -> None:
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
Decision.model_validate({"approved": "nope"})
|
||||
assert not is_invalid_resume(exc_info.value)
|
||||
assert not is_invalid_resume(ValueError("nope"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resume_style", ["null", "id_map"])
|
||||
def test_interrupt_response_schema_invalid_resume_after_earlier_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver, resume_style: str
|
||||
|
||||
@@ -5932,9 +5932,9 @@ def test_no_redundant_put_writes_for_cached_task(
|
||||
put_writes_task_ids: list[str] = []
|
||||
orig = PregelLoop.put_writes
|
||||
|
||||
def spy(self, task_id, writes, **kwargs):
|
||||
def spy(self, task_id, writes):
|
||||
put_writes_task_ids.append(task_id)
|
||||
return orig(self, task_id, writes, **kwargs)
|
||||
return orig(self, task_id, writes)
|
||||
|
||||
with patch.object(PregelLoop, "put_writes", spy):
|
||||
result = workflow.invoke(Command(resume="ans"), config=config)
|
||||
|
||||
@@ -8157,9 +8157,9 @@ async def test_no_redundant_put_writes_for_cached_task(
|
||||
put_writes_task_ids: list[str] = []
|
||||
orig = PregelLoop.put_writes
|
||||
|
||||
def spy(self, task_id, writes, **kwargs):
|
||||
def spy(self, task_id, writes):
|
||||
put_writes_task_ids.append(task_id)
|
||||
return orig(self, task_id, writes, **kwargs)
|
||||
return orig(self, task_id, writes)
|
||||
|
||||
with patch.object(PregelLoop, "put_writes", spy):
|
||||
result = await workflow.ainvoke(Command(resume="ans"), config=config)
|
||||
|
||||
@@ -248,95 +248,6 @@ def test_drain_from_subgraph_can_resume_parent() -> None:
|
||||
}
|
||||
|
||||
|
||||
class _BreakpointState(TypedDict, total=False):
|
||||
first: str
|
||||
second: str
|
||||
|
||||
|
||||
def _drain_before_breakpoint_builder(
|
||||
second_runs: list[str],
|
||||
) -> StateGraph[_BreakpointState]:
|
||||
def first(state: _BreakpointState, runtime: Runtime) -> _BreakpointState:
|
||||
runtime.control.request_drain("rollout")
|
||||
return {"first": "done"}
|
||||
|
||||
def second(state: _BreakpointState) -> _BreakpointState:
|
||||
second_runs.append("second")
|
||||
return {"second": "done"}
|
||||
|
||||
builder = StateGraph(_BreakpointState)
|
||||
builder.add_node("first", first)
|
||||
builder.add_node("second", second)
|
||||
builder.add_edge(START, "first")
|
||||
builder.add_edge("first", "second")
|
||||
return builder
|
||||
|
||||
|
||||
@pytest.mark.parametrize("at_compile", [True, False])
|
||||
def test_drain_before_interrupt_before_stops_at_breakpoint(at_compile: bool) -> None:
|
||||
# A resume passes the next step's breakpoints, so a drain right before a
|
||||
# breakpoint has to stop there instead, or the resumed run skips it.
|
||||
second_runs: list[str] = []
|
||||
breakpoints = {"interrupt_before": ["second"]}
|
||||
compiled = _drain_before_breakpoint_builder(second_runs).compile(
|
||||
checkpointer=MemorySaver(), **(breakpoints if at_compile else {})
|
||||
)
|
||||
run_kwargs = {} if at_compile else breakpoints
|
||||
config = {"configurable": {"thread_id": "drain-breakpoint"}}
|
||||
|
||||
assert compiled.invoke({}, config, control=RunControl(), **run_kwargs) == {
|
||||
"first": "done"
|
||||
}
|
||||
assert compiled.get_state(config).next == ("second",)
|
||||
assert second_runs == []
|
||||
|
||||
assert compiled.invoke(None, config, **run_kwargs) == {
|
||||
"first": "done",
|
||||
"second": "done",
|
||||
}
|
||||
assert second_runs == ["second"]
|
||||
|
||||
|
||||
def test_drain_before_subgraph_interrupt_before_stops_at_breakpoint() -> None:
|
||||
second_runs: list[str] = []
|
||||
child = _drain_before_breakpoint_builder(second_runs).compile(
|
||||
interrupt_before=["second"]
|
||||
)
|
||||
parent = StateGraph(_BreakpointState)
|
||||
parent.add_node("child", child)
|
||||
parent.add_edge(START, "child")
|
||||
compiled = parent.compile(checkpointer=MemorySaver())
|
||||
config = {"configurable": {"thread_id": "drain-subgraph-breakpoint"}}
|
||||
|
||||
compiled.invoke({}, config, control=RunControl())
|
||||
state = compiled.get_state(config, subgraphs=True)
|
||||
assert state.next == ("child",)
|
||||
assert state.tasks[0].state.next == ("second",)
|
||||
assert second_runs == []
|
||||
|
||||
assert compiled.invoke(None, config) == {"first": "done", "second": "done"}
|
||||
assert second_runs == ["second"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_drain_before_interrupt_before_stops_at_breakpoint_async() -> None:
|
||||
second_runs: list[str] = []
|
||||
compiled = _drain_before_breakpoint_builder(second_runs).compile(
|
||||
checkpointer=MemorySaver(), interrupt_before=["second"]
|
||||
)
|
||||
config = {"configurable": {"thread_id": "drain-breakpoint-async"}}
|
||||
|
||||
assert await compiled.ainvoke({}, config, control=RunControl()) == {"first": "done"}
|
||||
assert (await compiled.aget_state(config)).next == ("second",)
|
||||
assert second_runs == []
|
||||
|
||||
assert await compiled.ainvoke(None, config) == {
|
||||
"first": "done",
|
||||
"second": "done",
|
||||
}
|
||||
assert second_runs == ["second"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_drain_requested_in_terminal_step_finishes_normally_async() -> None:
|
||||
class State(TypedDict, total=False):
|
||||
|
||||
Generated
+3
-3
@@ -1451,7 +1451,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.15"
|
||||
version = "1.2.14"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1685,7 +1685,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.3"
|
||||
version = "3.1.2"
|
||||
source = { editable = "../checkpoint-postgres" }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
@@ -1734,7 +1734,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.1"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -99,14 +99,6 @@ if TYPE_CHECKING:
|
||||
from langgraph.runtime import Runtime
|
||||
from pydantic_core import ErrorDetails
|
||||
|
||||
try:
|
||||
from langgraph.errors import is_invalid_resume
|
||||
except ImportError: # `langgraph` before `is_invalid_resume` never marks resume errors
|
||||
|
||||
def is_invalid_resume(error: BaseException) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
# right now we use a dict as the default, can change this to AgentState, but depends
|
||||
# on if this lives in LangChain or LangGraph... ideally would have some typed
|
||||
# messages key
|
||||
@@ -965,11 +957,6 @@ class ToolNode(RunnableCallable):
|
||||
try:
|
||||
response = tool.invoke(call_args, config)
|
||||
except ValidationError as exc:
|
||||
if is_invalid_resume(exc):
|
||||
# An `interrupt()` in this tool, or in a graph it ran, got a resume
|
||||
# value that doesn't match its `response_schema`. That's not a bad
|
||||
# tool argument: fail the run so the interrupt can be answered again.
|
||||
raise
|
||||
# Filter out errors for injected arguments
|
||||
injected = self._injected_args.get(call["name"])
|
||||
filtered_errors = _filter_validation_errors(exc, injected)
|
||||
@@ -995,10 +982,6 @@ class ToolNode(RunnableCallable):
|
||||
except GraphBubbleUp:
|
||||
raise
|
||||
except Exception as e:
|
||||
# The model can't fix a resume value that doesn't match an interrupt's
|
||||
# `response_schema`, so no `handle_tool_errors` setting handles it.
|
||||
if is_invalid_resume(e):
|
||||
raise
|
||||
# Determine which exception types are handled
|
||||
handled_types: tuple[type[Exception], ...]
|
||||
if isinstance(self._handle_tool_errors, type) and issubclass(
|
||||
@@ -1070,13 +1053,9 @@ class ToolNode(RunnableCallable):
|
||||
# Call wrapper with request and execute callable
|
||||
try:
|
||||
return self._wrap_tool_call(tool_request, execute)
|
||||
except GraphBubbleUp:
|
||||
# Interrupts always propagate, as they do without a wrapper.
|
||||
raise
|
||||
except Exception as e:
|
||||
# Wrapper threw an exception. The model can't fix a resume value that
|
||||
# doesn't match an interrupt's `response_schema`, so it's never handled.
|
||||
if not self._handle_tool_errors or is_invalid_resume(e):
|
||||
# Wrapper threw an exception
|
||||
if not self._handle_tool_errors:
|
||||
raise
|
||||
# Convert to error message
|
||||
content = _handle_tool_error(e, flag=self._handle_tool_errors)
|
||||
@@ -1125,11 +1104,6 @@ class ToolNode(RunnableCallable):
|
||||
try:
|
||||
response = await tool.ainvoke(call_args, config)
|
||||
except ValidationError as exc:
|
||||
if is_invalid_resume(exc):
|
||||
# An `interrupt()` in this tool, or in a graph it ran, got a resume
|
||||
# value that doesn't match its `response_schema`. That's not a bad
|
||||
# tool argument: fail the run so the interrupt can be answered again.
|
||||
raise
|
||||
# Filter out errors for injected arguments
|
||||
injected = self._injected_args.get(call["name"])
|
||||
filtered_errors = _filter_validation_errors(exc, injected)
|
||||
@@ -1155,10 +1129,6 @@ class ToolNode(RunnableCallable):
|
||||
except GraphBubbleUp:
|
||||
raise
|
||||
except Exception as e:
|
||||
# The model can't fix a resume value that doesn't match an interrupt's
|
||||
# `response_schema`, so no `handle_tool_errors` setting handles it.
|
||||
if is_invalid_resume(e):
|
||||
raise
|
||||
# Determine which exception types are handled
|
||||
handled_types: tuple[type[Exception], ...]
|
||||
if isinstance(self._handle_tool_errors, type) and issubclass(
|
||||
@@ -1238,13 +1208,9 @@ class ToolNode(RunnableCallable):
|
||||
# None check was performed above already
|
||||
self._wrap_tool_call = cast("ToolCallWrapper", self._wrap_tool_call)
|
||||
return self._wrap_tool_call(tool_request, _sync_execute)
|
||||
except GraphBubbleUp:
|
||||
# Interrupts always propagate, as they do without a wrapper.
|
||||
raise
|
||||
except Exception as e:
|
||||
# Wrapper threw an exception. The model can't fix a resume value that
|
||||
# doesn't match an interrupt's `response_schema`, so it's never handled.
|
||||
if not self._handle_tool_errors or is_invalid_resume(e):
|
||||
# Wrapper threw an exception
|
||||
if not self._handle_tool_errors:
|
||||
raise
|
||||
# Convert to error message
|
||||
content = _handle_tool_error(e, flag=self._handle_tool_errors)
|
||||
|
||||
@@ -25,7 +25,6 @@ from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.tools import BaseTool, InjectedToolArg, ToolException
|
||||
from langchain_core.tools import tool as dec_tool
|
||||
from langchain_core.tools.base import InjectedToolCallId
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.errors import GraphBubbleUp, GraphInterrupt
|
||||
from langgraph.graph import START, MessagesState, StateGraph
|
||||
@@ -33,8 +32,8 @@ from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
|
||||
from langgraph.runtime import ExecutionInfo, ServerInfo
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
from langgraph.types import Command, Send, interrupt
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from langgraph.types import Command, Send
|
||||
from pydantic import BaseModel
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
@@ -627,149 +626,6 @@ def test_tool_node_node_interrupt() -> None:
|
||||
assert exc_info.value == "foo"
|
||||
|
||||
|
||||
class _Approval(BaseModel):
|
||||
approved: bool
|
||||
|
||||
|
||||
class _AskState(TypedDict, total=False):
|
||||
answer: str
|
||||
|
||||
|
||||
def _approval_graph():
|
||||
"""A graph that asks a human for approval with a typed interrupt."""
|
||||
|
||||
def ask(state: _AskState) -> _AskState:
|
||||
approval = interrupt("Approve?", response_schema=_Approval)
|
||||
return {"answer": f"approved={approval.approved}"}
|
||||
|
||||
return StateGraph(_AskState).add_node("ask", ask).add_edge(START, "ask").compile()
|
||||
|
||||
|
||||
def _ask_human_call() -> dict[str, list[AnyMessage]]:
|
||||
call = ToolCall(name="ask_human", args={}, id="call_1")
|
||||
return {"messages": [AIMessage("", tool_calls=[call])]}
|
||||
|
||||
|
||||
def _handle_any(e): # no annotation: handles every error
|
||||
return "handled"
|
||||
|
||||
|
||||
# A bad answer to an interrupt must fail the run whatever `handle_tool_errors` is,
|
||||
# including settings that cover `ValidationError` (a `ValueError`), with or without
|
||||
# a wrapper. `create_agent` always runs tools through a wrapper (its middleware).
|
||||
_TOOL_NODES = pytest.mark.parametrize(
|
||||
("wrapped", "handle_tool_errors"),
|
||||
[
|
||||
(wrapped, handler)
|
||||
for wrapped in (False, True)
|
||||
for handler in (None, True, (ValueError,), _handle_any)
|
||||
],
|
||||
ids=[
|
||||
f"{wrapped}-{handler}"
|
||||
for wrapped in ("plain", "wrapped")
|
||||
for handler in ("default", "handle_true", "handle_value_error", "untyped")
|
||||
],
|
||||
)
|
||||
# The interrupt either runs in a graph the tool starts (a subagent) or in the tool.
|
||||
_SHAPE = pytest.mark.parametrize("nested", [True, False], ids=["nested", "direct"])
|
||||
|
||||
|
||||
@_TOOL_NODES
|
||||
@_SHAPE
|
||||
def test_tool_node_reraises_invalid_resume(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
wrapped: bool,
|
||||
nested: bool,
|
||||
handle_tool_errors: Any,
|
||||
) -> None:
|
||||
asker = _approval_graph()
|
||||
|
||||
@dec_tool
|
||||
def ask_human() -> str:
|
||||
"""Ask a human for approval."""
|
||||
if nested:
|
||||
return asker.invoke({})["answer"]
|
||||
approval = interrupt("Approve?", response_schema=_Approval)
|
||||
return f"approved={approval.approved}"
|
||||
|
||||
def pass_through(request, handler):
|
||||
return handler(request)
|
||||
|
||||
errors = (
|
||||
{} if handle_tool_errors is None else {"handle_tool_errors": handle_tool_errors}
|
||||
)
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node(
|
||||
"tools",
|
||||
ToolNode(
|
||||
[ask_human], wrap_tool_call=pass_through if wrapped else None, **errors
|
||||
),
|
||||
)
|
||||
.add_edge(START, "tools")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
[pending] = graph.invoke(_ask_human_call(), config)["__interrupt__"]
|
||||
|
||||
# A bad answer isn't a bad tool argument: the run fails without saving, so
|
||||
# the same interrupt can be answered again.
|
||||
with pytest.raises(ValidationError, match="approved"):
|
||||
graph.invoke(Command(resume={pending.id: {"approved": "maybe"}}), config)
|
||||
assert [i.id for i in graph.get_state(config).interrupts] == [pending.id]
|
||||
|
||||
result = graph.invoke(Command(resume={pending.id: {"approved": True}}), config)
|
||||
assert result["messages"][-1].content == "approved=True"
|
||||
|
||||
|
||||
@_TOOL_NODES
|
||||
@_SHAPE
|
||||
async def test_tool_node_reraises_invalid_resume_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
wrapped: bool,
|
||||
nested: bool,
|
||||
handle_tool_errors: Any,
|
||||
) -> None:
|
||||
asker = _approval_graph()
|
||||
|
||||
@dec_tool
|
||||
async def ask_human() -> str:
|
||||
"""Ask a human for approval."""
|
||||
if nested:
|
||||
return (await asker.ainvoke({}))["answer"]
|
||||
approval = interrupt("Approve?", response_schema=_Approval)
|
||||
return f"approved={approval.approved}"
|
||||
|
||||
async def pass_through(request, handler):
|
||||
return await handler(request)
|
||||
|
||||
errors = (
|
||||
{} if handle_tool_errors is None else {"handle_tool_errors": handle_tool_errors}
|
||||
)
|
||||
graph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node(
|
||||
"tools",
|
||||
ToolNode(
|
||||
[ask_human], awrap_tool_call=pass_through if wrapped else None, **errors
|
||||
),
|
||||
)
|
||||
.add_edge(START, "tools")
|
||||
.compile(checkpointer=async_checkpointer)
|
||||
)
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
[pending] = (await graph.ainvoke(_ask_human_call(), config))["__interrupt__"]
|
||||
|
||||
with pytest.raises(ValidationError, match="approved"):
|
||||
await graph.ainvoke(Command(resume={pending.id: {"approved": "maybe"}}), config)
|
||||
state = await graph.aget_state(config)
|
||||
assert [i.id for i in state.interrupts] == [pending.id]
|
||||
|
||||
resume = Command(resume={pending.id: {"approved": True}})
|
||||
result = await graph.ainvoke(resume, config)
|
||||
assert result["messages"][-1].content == "approved=True"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("input_type", ["dict", "tool_calls"])
|
||||
async def test_tool_node_command(input_type: str) -> None:
|
||||
|
||||
|
||||
Generated
+3
-3
@@ -286,7 +286,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.15"
|
||||
version = "1.2.14"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -418,7 +418,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.3"
|
||||
version = "3.1.2"
|
||||
source = { editable = "../checkpoint-postgres" }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
@@ -467,7 +467,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.1"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -3,7 +3,7 @@ from langgraph_sdk.client import get_client, get_sync_client
|
||||
from langgraph_sdk.encryption import Encryption
|
||||
from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
|
||||
|
||||
__version__ = "0.4.7"
|
||||
__version__ = "0.4.6"
|
||||
|
||||
__all__ = [
|
||||
"Auth",
|
||||
|
||||
Generated
+1
-1
@@ -301,7 +301,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.15"
|
||||
version = "1.2.14"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Reference in New Issue
Block a user