mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-21 23:41:53 +02:00
Compare commits
20
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4ac46ba186 | ||
|
|
e14f4277dc | ||
|
|
4e2dfe4a24 | ||
|
|
e77a5cb08c | ||
|
|
b0482d0627 | ||
|
|
889a5ff16f | ||
|
|
5eb60aef6f | ||
|
|
d4149cc1f5 | ||
|
|
36bf3887f0 | ||
|
|
e5b422455d | ||
|
|
6cd26d0e21 | ||
|
|
53cf85bece | ||
|
|
69d98cc0d9 | ||
|
|
ccb851effb | ||
|
|
a6731f4367 | ||
|
|
8a77332d0d | ||
|
|
a1ce197234 | ||
|
|
afbe7e48b1 | ||
|
|
1074714c26 | ||
|
|
740461ed50 |
@@ -27,10 +27,15 @@ from psycopg_pool import ConnectionPool
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
_DELTA_PAGE_SIZE,
|
||||
PRUNE_DELETE_THREADS_SQL,
|
||||
PRUNE_STRATEGIES,
|
||||
BasePostgresSaver,
|
||||
_build_delta_stage1_sql,
|
||||
_build_delta_stage2_sql,
|
||||
_DeltaStage2Row,
|
||||
_execute,
|
||||
_keep_latest,
|
||||
_refuse_delta_threads,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
@@ -401,6 +406,73 @@ class PostgresSaver(BasePostgresSaver):
|
||||
(str(thread_id),),
|
||||
)
|
||||
|
||||
def prune(
|
||||
self,
|
||||
thread_ids: Sequence[str],
|
||||
*,
|
||||
strategy: str = "keep_latest",
|
||||
) -> None:
|
||||
"""Prune checkpoints for the given threads.
|
||||
|
||||
`"delete"` removes every checkpoint, write and blob of the threads, in
|
||||
one transaction.
|
||||
|
||||
`"keep_latest"` keeps the latest checkpoint of each namespace (the one
|
||||
`get_tuple` loads, plus its parent when it predates v4 and has pending
|
||||
sends), with its pending writes and the blobs it references, and
|
||||
deletes the rest. A kept checkpoint whose parent it deletes is left with
|
||||
no parent. A subgraph call's namespace is deleted with the root
|
||||
checkpoint the call started from, except a `checkpointer=True`
|
||||
subgraph's, which keeps its latest checkpoint. Each thread is pruned in
|
||||
its own transaction, on one snapshot of it, so a run appending to the
|
||||
thread meanwhile is safe; forking or replaying from a checkpoint being
|
||||
deleted is not. A thread that fails because another transaction changed
|
||||
it leaves the threads before it pruned, and retrying is safe.
|
||||
|
||||
`"keep_latest"` refuses threads that use a `DeltaChannel`, since it
|
||||
rebuilds the latest value from older checkpoints: if any requested
|
||||
thread does, nothing is pruned. A thread that starts using one before
|
||||
its turn in the batch is refused then, after the threads before it.
|
||||
|
||||
Args:
|
||||
thread_ids: The thread IDs to prune.
|
||||
strategy: `"keep_latest"` or `"delete"`.
|
||||
|
||||
Raises:
|
||||
TypeError: If `thread_ids` is a single string.
|
||||
ValueError: If `strategy` is unknown.
|
||||
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
|
||||
uses a `DeltaChannel`.
|
||||
psycopg.errors.SerializationFailure: If another transaction, such
|
||||
as a second prune, changes a thread's rows while `"keep_latest"`
|
||||
prunes it.
|
||||
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
|
||||
at the same time.
|
||||
"""
|
||||
if isinstance(thread_ids, str):
|
||||
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
|
||||
if strategy not in PRUNE_STRATEGIES:
|
||||
raise ValueError(f"Unknown prune strategy: {strategy!r}")
|
||||
thread_ids = [str(t) for t in thread_ids]
|
||||
if not thread_ids:
|
||||
return
|
||||
if strategy == "delete":
|
||||
with self._cursor() as cur, cur.connection.transaction():
|
||||
for sql in PRUNE_DELETE_THREADS_SQL:
|
||||
cur.execute(sql, (thread_ids,))
|
||||
return
|
||||
with self._cursor() as cur:
|
||||
_execute(cur, _refuse_delta_threads(thread_ids))
|
||||
for thread_id in thread_ids:
|
||||
with self._cursor() as cur:
|
||||
if self.pipe:
|
||||
# The transaction has to start with its isolation level, and
|
||||
# in pipeline mode it would take in what ran since the last
|
||||
# sync.
|
||||
self.pipe.sync()
|
||||
with cur.connection.transaction():
|
||||
_execute(cur, _keep_latest(thread_id))
|
||||
|
||||
@contextmanager
|
||||
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
|
||||
"""Create a database cursor as a context manager.
|
||||
|
||||
@@ -27,10 +27,15 @@ from psycopg_pool import AsyncConnectionPool
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
_DELTA_PAGE_SIZE,
|
||||
PRUNE_DELETE_THREADS_SQL,
|
||||
PRUNE_STRATEGIES,
|
||||
BasePostgresSaver,
|
||||
_aexecute,
|
||||
_build_delta_stage1_sql,
|
||||
_build_delta_stage2_sql,
|
||||
_DeltaStage2Row,
|
||||
_keep_latest,
|
||||
_refuse_delta_threads,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
@@ -360,6 +365,55 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
(str(thread_id),),
|
||||
)
|
||||
|
||||
async def aprune(
|
||||
self,
|
||||
thread_ids: Sequence[str],
|
||||
*,
|
||||
strategy: str = "keep_latest",
|
||||
) -> None:
|
||||
"""Prune checkpoints for the given threads.
|
||||
|
||||
See `PostgresSaver.prune` for what each strategy keeps.
|
||||
|
||||
Args:
|
||||
thread_ids: The thread IDs to prune.
|
||||
strategy: `"keep_latest"` or `"delete"`.
|
||||
|
||||
Raises:
|
||||
TypeError: If `thread_ids` is a single string.
|
||||
ValueError: If `strategy` is unknown.
|
||||
NotImplementedError: If `strategy` is `"keep_latest"` and a thread
|
||||
uses a `DeltaChannel`.
|
||||
psycopg.errors.SerializationFailure: If another transaction, such
|
||||
as a second prune, changes a thread's rows while `"keep_latest"`
|
||||
prunes it.
|
||||
psycopg.errors.DeadlockDetected: If a second prune of a thread runs
|
||||
at the same time.
|
||||
"""
|
||||
if isinstance(thread_ids, str):
|
||||
raise TypeError("thread_ids must be a sequence of thread IDs, not a str")
|
||||
if strategy not in PRUNE_STRATEGIES:
|
||||
raise ValueError(f"Unknown prune strategy: {strategy!r}")
|
||||
thread_ids = [str(t) for t in thread_ids]
|
||||
if not thread_ids:
|
||||
return
|
||||
if strategy == "delete":
|
||||
async with self._cursor() as cur, cur.connection.transaction():
|
||||
for sql in PRUNE_DELETE_THREADS_SQL:
|
||||
await cur.execute(sql, (thread_ids,))
|
||||
return
|
||||
async with self._cursor() as cur:
|
||||
await _aexecute(cur, _refuse_delta_threads(thread_ids))
|
||||
for thread_id in thread_ids:
|
||||
async with self._cursor() as cur:
|
||||
if self.pipe:
|
||||
# The transaction has to start with its isolation level, and
|
||||
# in pipeline mode it would take in what ran since the last
|
||||
# sync.
|
||||
await self.pipe.sync()
|
||||
async with cur.connection.transaction():
|
||||
await _aexecute(cur, _keep_latest(thread_id))
|
||||
|
||||
@asynccontextmanager
|
||||
async def _cursor(
|
||||
self, *, pipeline: bool = False
|
||||
@@ -688,6 +742,33 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
self.adelete_thread(thread_id), self.loop
|
||||
).result()
|
||||
|
||||
def prune(
|
||||
self,
|
||||
thread_ids: Sequence[str],
|
||||
*,
|
||||
strategy: str = "keep_latest",
|
||||
) -> None:
|
||||
"""Prune checkpoints for the given threads.
|
||||
|
||||
See `PostgresSaver.prune` for what each strategy keeps.
|
||||
|
||||
Args:
|
||||
thread_ids: The thread IDs to prune.
|
||||
strategy: `"keep_latest"` or `"delete"`.
|
||||
"""
|
||||
try:
|
||||
if asyncio.get_running_loop() is self.loop:
|
||||
raise asyncio.InvalidStateError(
|
||||
"Synchronous calls to AsyncPostgresSaver are only allowed from a "
|
||||
"different thread. From the main thread, use the async interface. "
|
||||
"For example, use `await checkpointer.aprune(...)`."
|
||||
)
|
||||
except RuntimeError:
|
||||
pass
|
||||
return asyncio.run_coroutine_threadsafe(
|
||||
self.aprune(thread_ids, strategy=strategy), self.loop
|
||||
).result()
|
||||
|
||||
def get_delta_channel_history(
|
||||
self, *, config: RunnableConfig, channels: Sequence[str]
|
||||
) -> Mapping[str, DeltaChannelHistory]:
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import random
|
||||
import warnings
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Generator, Mapping, Sequence
|
||||
from importlib.metadata import version as get_version
|
||||
from typing import Any, TypedDict, cast
|
||||
|
||||
@@ -17,6 +17,8 @@ from langgraph.checkpoint.base import (
|
||||
writes_sort_key,
|
||||
)
|
||||
from langgraph.checkpoint.serde.types import TASKS
|
||||
from psycopg import AsyncCursor, Cursor
|
||||
from psycopg.rows import DictRow
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
# Page size for stage-1 paged scan in `get_delta_channel_history`. Internal
|
||||
@@ -160,6 +162,218 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
"""
|
||||
|
||||
|
||||
PRUNE_STRATEGIES = ("keep_latest", "delete")
|
||||
|
||||
PRUNE_DELETE_THREADS_SQL = (
|
||||
"DELETE FROM checkpoints WHERE thread_id = ANY(%s)",
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = ANY(%s)",
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = ANY(%s)",
|
||||
)
|
||||
|
||||
# The prune statements below never join two sets of a thread's rows: a join is
|
||||
# planned from per-thread row estimates, which are worst for the few huge
|
||||
# threads worth pruning, and can go quadratic. They look rows up by primary key
|
||||
# instead (`LATERAL ... LIMIT 1`, `= ANY`).
|
||||
|
||||
# A `DeltaChannel` leaves update counters in checkpoint metadata, or a
|
||||
# snapshot: an inline `true` marker with a blob behind it, which a plain `True`
|
||||
# value never has. LangGraph.js inlines no values, but marks the threads its
|
||||
# DeltaChannel graphs start with `delta_writes_versioned`.
|
||||
PRUNE_DELTA_THREADS_SQL = """
|
||||
SELECT thread_id FROM checkpoints
|
||||
WHERE thread_id = ANY(%s) AND (
|
||||
metadata ? 'counters_since_delta_snapshot'
|
||||
OR metadata ? 'delta_writes_versioned'
|
||||
)
|
||||
UNION
|
||||
SELECT c.thread_id
|
||||
FROM checkpoints c
|
||||
CROSS JOIN LATERAL jsonb_each(c.checkpoint -> 'channel_values') AS v
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT 1 FROM checkpoint_blobs b
|
||||
WHERE b.thread_id = c.thread_id
|
||||
AND b.checkpoint_ns = c.checkpoint_ns
|
||||
AND b.channel = v.key
|
||||
AND b.version = c.checkpoint -> 'channel_versions' ->> v.key
|
||||
LIMIT 1
|
||||
) AS snapshot
|
||||
WHERE c.thread_id = ANY(%s) AND v.value = 'true'::jsonb
|
||||
"""
|
||||
|
||||
# Per namespace: the latest checkpoint, the parent it reads its pending sends
|
||||
# from when it predates v4, and the root graph's checkpoint its subgraph call
|
||||
# started from, with whether that one is saved.
|
||||
PRUNE_HEADS_SQL = f"""
|
||||
SELECT
|
||||
h.checkpoint_ns,
|
||||
h.checkpoint_id AS head,
|
||||
CASE WHEN (c.checkpoint ->> 'v')::int < 4 AND EXISTS (
|
||||
SELECT 1 FROM checkpoint_writes w
|
||||
WHERE w.thread_id = %s
|
||||
AND w.checkpoint_ns = h.checkpoint_ns
|
||||
AND w.checkpoint_id = c.parent_checkpoint_id
|
||||
AND w.channel = '{TASKS}'
|
||||
) THEN c.parent_checkpoint_id END AS sends_parent,
|
||||
c.metadata -> 'parents' ->> '' AS root,
|
||||
EXISTS (
|
||||
SELECT 1 FROM checkpoints r
|
||||
WHERE r.thread_id = %s AND r.checkpoint_ns = ''
|
||||
AND r.checkpoint_id = c.metadata -> 'parents' ->> ''
|
||||
) AS root_saved
|
||||
FROM (
|
||||
SELECT DISTINCT ON (checkpoint_ns) checkpoint_ns, checkpoint_id
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s
|
||||
ORDER BY checkpoint_ns DESC, checkpoint_id DESC
|
||||
) AS h
|
||||
CROSS JOIN LATERAL (
|
||||
SELECT checkpoint, parent_checkpoint_id, metadata FROM checkpoints
|
||||
WHERE thread_id = %s
|
||||
AND checkpoint_ns = h.checkpoint_ns
|
||||
AND checkpoint_id = h.checkpoint_id
|
||||
) AS c
|
||||
"""
|
||||
|
||||
PRUNE_SNAPSHOT_SQL = "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"
|
||||
|
||||
PRUNE_UNKEPT_SQL = """
|
||||
SELECT checkpoint_id FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
|
||||
"""
|
||||
|
||||
# A run saves a superstep's writes in the background, where they can commit
|
||||
# before their checkpoint does, so they go by the checkpoints deleted.
|
||||
PRUNE_WRITES_SQL = """
|
||||
DELETE FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = ANY(%s)
|
||||
"""
|
||||
|
||||
PRUNE_CHECKPOINTS_SQL = """
|
||||
WITH deleted AS (
|
||||
DELETE FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id <> ALL(%s)
|
||||
RETURNING checkpoint -> 'channel_versions' AS versions
|
||||
)
|
||||
SELECT DISTINCT v.key AS channel, v.value AS version
|
||||
FROM deleted
|
||||
CROSS JOIN LATERAL jsonb_each_text(deleted.versions) AS v
|
||||
"""
|
||||
|
||||
PRUNE_PARENTS_SQL = """
|
||||
UPDATE checkpoints SET parent_checkpoint_id = NULL
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
AND checkpoint_id = ANY(%s) AND parent_checkpoint_id <> ALL(%s)
|
||||
"""
|
||||
|
||||
PRUNE_REFERENCED_BLOBS_SQL = """
|
||||
SELECT DISTINCT v.key AS channel, v.value AS version
|
||||
FROM checkpoints c
|
||||
CROSS JOIN LATERAL jsonb_each_text(c.checkpoint -> 'channel_versions') AS v
|
||||
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
|
||||
"""
|
||||
|
||||
PRUNE_BLOBS_SQL = """
|
||||
DELETE FROM checkpoint_blobs
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s AND version = ANY(%s)
|
||||
"""
|
||||
|
||||
PRUNE_SUBGRAPH_CALLS_SQL = (
|
||||
"DELETE FROM checkpoints WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
|
||||
"DELETE FROM checkpoint_writes WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
|
||||
"DELETE FROM checkpoint_blobs WHERE thread_id = %s AND checkpoint_ns = ANY(%s)",
|
||||
)
|
||||
|
||||
|
||||
def _garbage_blobs(
|
||||
freed: Sequence[Mapping[str, Any]], referenced: Sequence[Mapping[str, Any]]
|
||||
) -> dict[str, list[str]]:
|
||||
"""Blobs a deleted checkpoint used and no kept one does, by channel."""
|
||||
kept = {(r["channel"], r["version"]) for r in referenced}
|
||||
garbage: dict[str, list[str]] = {}
|
||||
for r in freed:
|
||||
if (r["channel"], r["version"]) not in kept:
|
||||
garbage.setdefault(r["channel"], []).append(r["version"])
|
||||
return garbage
|
||||
|
||||
|
||||
_Statements = Generator[tuple[str, Sequence[Any], bool], list[DictRow], None]
|
||||
|
||||
|
||||
def _refuse_delta_threads(thread_ids: Sequence[str]) -> _Statements:
|
||||
rows = yield PRUNE_DELTA_THREADS_SQL, (thread_ids, thread_ids), True
|
||||
if delta := sorted(r["thread_id"] for r in rows):
|
||||
raise NotImplementedError(
|
||||
f"keep_latest can't prune threads that use a DeltaChannel: {', '.join(delta)}"
|
||||
)
|
||||
|
||||
|
||||
def _keep_latest(thread_id: str) -> _Statements:
|
||||
"""The statements of one thread's keep_latest transaction, for `_execute`
|
||||
and `_aexecute` to run. Each comes with whether its rows are sent back.
|
||||
|
||||
They all read one snapshot, so what a run saves meanwhile is neither seen
|
||||
nor deleted, whatever its checkpoint ID.
|
||||
"""
|
||||
yield PRUNE_SNAPSHOT_SQL, (), False
|
||||
heads = yield PRUNE_HEADS_SQL, (thread_id,) * 4, True
|
||||
# Catches a thread that started using a DeltaChannel after the batch check.
|
||||
yield from _refuse_delta_threads([thread_id])
|
||||
kept = {
|
||||
h["checkpoint_ns"]: [i for i in (h["head"], h["sends_parent"]) if i]
|
||||
for h in heads
|
||||
}
|
||||
# A subgraph call's namespace has its task ID after a ":", unlike a
|
||||
# `checkpointer=True` subgraph's, whose state outlives each call. A call
|
||||
# goes with the root checkpoint it started from, but only once that one is
|
||||
# saved: runs save in the background, so a call can commit first.
|
||||
calls = {
|
||||
h["checkpoint_ns"]
|
||||
for h in heads
|
||||
if ":" in h["checkpoint_ns"]
|
||||
and h["root_saved"]
|
||||
and h["root"] not in kept.get("", ())
|
||||
}
|
||||
if calls:
|
||||
for sql in PRUNE_SUBGRAPH_CALLS_SQL:
|
||||
yield sql, (thread_id, list(calls)), False
|
||||
for ns, ids in kept.items():
|
||||
if ns in calls:
|
||||
continue
|
||||
unkept = yield PRUNE_UNKEPT_SQL, (thread_id, ns, ids), True
|
||||
if unkept:
|
||||
deleted = [r["checkpoint_id"] for r in unkept]
|
||||
yield PRUNE_WRITES_SQL, (thread_id, ns, deleted), False
|
||||
freed = yield PRUNE_CHECKPOINTS_SQL, (thread_id, ns, ids), True
|
||||
yield PRUNE_PARENTS_SQL, (thread_id, ns, ids, ids), False
|
||||
if freed:
|
||||
referenced = yield PRUNE_REFERENCED_BLOBS_SQL, (thread_id, ns), True
|
||||
for channel, versions in _garbage_blobs(freed, referenced).items():
|
||||
yield PRUNE_BLOBS_SQL, (thread_id, ns, channel, versions), False
|
||||
|
||||
|
||||
# Unprepared, so Postgres plans each statement knowing its arrays: a cached
|
||||
# generic plan checks a row against an `= ANY(...)` array one element at a time.
|
||||
def _execute(cur: Cursor[DictRow], statements: _Statements) -> None:
|
||||
try:
|
||||
sql, params, fetch = next(statements)
|
||||
while True:
|
||||
cur.execute(sql, params, prepare=False)
|
||||
sql, params, fetch = statements.send(cur.fetchall() if fetch else [])
|
||||
except StopIteration:
|
||||
pass
|
||||
|
||||
|
||||
async def _aexecute(cur: AsyncCursor[DictRow], statements: _Statements) -> None:
|
||||
try:
|
||||
sql, params, fetch = next(statements)
|
||||
while True:
|
||||
await cur.execute(sql, params, prepare=False)
|
||||
rows = await cur.fetchall() if fetch else []
|
||||
sql, params, fetch = statements.send(rows)
|
||||
except StopIteration:
|
||||
pass
|
||||
|
||||
|
||||
class _DeltaStage2Row(TypedDict, total=False):
|
||||
"""One row from `_build_delta_stage2_sql` (a UNION ALL of writes and blobs)."""
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.2"
|
||||
version = "3.2.0"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1,540 @@
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import empty_checkpoint
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
from langgraph.checkpoint.conformance import validate
|
||||
from langgraph.checkpoint.conformance.initializer import checkpointer_test
|
||||
from langgraph.checkpoint.serde.types import TASKS, _DeltaSnapshot
|
||||
from psycopg import AsyncConnection
|
||||
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
PRUNE_DELTA_THREADS_SQL,
|
||||
PRUNE_HEADS_SQL,
|
||||
PRUNE_REFERENCED_BLOBS_SQL,
|
||||
)
|
||||
from tests.conftest import DEFAULT_POSTGRES_URI, DEFAULT_URI
|
||||
from tests.test_async import _saver as _async_saver
|
||||
from tests.test_sync import _saver as _sync_saver
|
||||
|
||||
SAVERS = ["base", "pool", "pipe", "aio_base", "aio_pool", "aio_pipe"]
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _saver(name: str):
|
||||
if name.startswith("aio_"):
|
||||
async with _async_saver(name.removeprefix("aio_")) as saver:
|
||||
yield saver
|
||||
else:
|
||||
with _sync_saver(name) as saver:
|
||||
yield saver
|
||||
|
||||
|
||||
async def _call(saver: Any, method: str, *args: Any, **kwargs: Any) -> Any:
|
||||
if isinstance(saver, AsyncPostgresSaver):
|
||||
return await getattr(saver, f"a{method}")(*args, **kwargs)
|
||||
return getattr(saver, method)(*args, **kwargs)
|
||||
|
||||
|
||||
async def _rows(saver: Any, thread_id: str) -> dict[str, list[tuple]]:
|
||||
queries = {
|
||||
"checkpoints": "SELECT checkpoint_ns, checkpoint_id FROM checkpoints"
|
||||
" WHERE thread_id = %s ORDER BY 1, 2",
|
||||
"writes": "SELECT checkpoint_ns, checkpoint_id, task_id, idx"
|
||||
" FROM checkpoint_writes WHERE thread_id = %s ORDER BY 1, 2, 3, 4",
|
||||
"blobs": "SELECT checkpoint_ns, channel, version FROM checkpoint_blobs"
|
||||
" WHERE thread_id = %s ORDER BY 1, 2, 3",
|
||||
}
|
||||
rows = {}
|
||||
for table, sql in queries.items():
|
||||
if isinstance(saver, AsyncPostgresSaver):
|
||||
async with saver._cursor() as cur:
|
||||
await cur.execute(sql, (thread_id,))
|
||||
rows[table] = [tuple(r.values()) for r in await cur.fetchall()]
|
||||
else:
|
||||
with saver._cursor() as cur:
|
||||
cur.execute(sql, (thread_id,))
|
||||
rows[table] = [tuple(r.values()) for r in cur.fetchall()]
|
||||
return rows
|
||||
|
||||
|
||||
async def _put(
|
||||
saver: Any,
|
||||
thread_id: str,
|
||||
parent: dict | None,
|
||||
step: int,
|
||||
values: dict[str, Any],
|
||||
versions: dict[str, str],
|
||||
*,
|
||||
ns: str = "",
|
||||
metadata: dict[str, Any] | None = None,
|
||||
v: int = 4,
|
||||
checkpoint_id: str | None = None,
|
||||
) -> dict:
|
||||
config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ns}}
|
||||
if parent is not None:
|
||||
config["configurable"]["checkpoint_id"] = parent["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint.update(
|
||||
v=v,
|
||||
id=checkpoint_id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=versions,
|
||||
)
|
||||
new_versions = {k: versions[k] for k in values}
|
||||
return await _call(
|
||||
saver,
|
||||
"put",
|
||||
config,
|
||||
checkpoint,
|
||||
{"source": "loop", "step": step, "parents": {}, **(metadata or {})},
|
||||
new_versions,
|
||||
)
|
||||
|
||||
|
||||
async def _three_checkpoints(
|
||||
saver: Any,
|
||||
thread_id: str,
|
||||
ns: str = "",
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> dict:
|
||||
parent = None
|
||||
steps = [
|
||||
({"a": [1], "s": ["shared"]}, {"a": "1", "s": "1"}),
|
||||
({"a": [1, 2]}, {"a": "2", "s": "1"}),
|
||||
({"a": [1, 2, 3]}, {"a": "3", "s": "1"}),
|
||||
]
|
||||
for step, (values, versions) in enumerate(steps):
|
||||
parent = await _put(
|
||||
saver, thread_id, parent, step, values, versions, ns=ns, metadata=metadata
|
||||
)
|
||||
await _call(saver, "put_writes", parent, [("a", step)], str(uuid4()))
|
||||
return parent
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _second_saver(saver: AsyncPostgresSaver):
|
||||
assert isinstance(saver.conn, AsyncConnection)
|
||||
uri = DEFAULT_POSTGRES_URI + saver.conn.info.dbname
|
||||
async with AsyncPostgresSaver.from_conn_string(uri) as other:
|
||||
yield other
|
||||
|
||||
|
||||
def _run_once_after(
|
||||
saver: AsyncPostgresSaver,
|
||||
sql: str,
|
||||
action: Callable[[], Awaitable[Any]],
|
||||
*,
|
||||
nth: int = 1,
|
||||
) -> None:
|
||||
seen = 0
|
||||
cursor = saver._cursor
|
||||
|
||||
class Hooked:
|
||||
def __init__(self, cur: Any) -> None:
|
||||
self.cur = cur
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self.cur, name)
|
||||
|
||||
async def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
|
||||
nonlocal seen
|
||||
await self.cur.execute(query, params, **kwargs)
|
||||
if query is sql:
|
||||
seen += 1
|
||||
if seen == nth:
|
||||
await action()
|
||||
|
||||
@asynccontextmanager
|
||||
async def hooked(**kwargs: Any):
|
||||
async with cursor(**kwargs) as cur:
|
||||
yield Hooked(cur)
|
||||
|
||||
saver._cursor = hooked
|
||||
|
||||
|
||||
async def _delta_thread(saver: Any, thread_id: str, evidence: str) -> None:
|
||||
if evidence == "counters":
|
||||
first = await _put(saver, thread_id, None, 0, {"m": [1]}, {"m": "1"})
|
||||
await _put(
|
||||
saver,
|
||||
thread_id,
|
||||
first,
|
||||
1,
|
||||
{},
|
||||
{"m": "2"},
|
||||
metadata={"counters_since_delta_snapshot": {"m": [1, 1]}},
|
||||
)
|
||||
elif evidence == "snapshot":
|
||||
first = await _put(
|
||||
saver, thread_id, None, 0, {"m": _DeltaSnapshot([1])}, {"m": "1"}
|
||||
)
|
||||
await _put(saver, thread_id, first, 1, {}, {"m": "2"})
|
||||
else:
|
||||
marker = {"delta_writes_versioned": True}
|
||||
first = await _put(
|
||||
saver, thread_id, None, 0, {"m": [1]}, {"m": "1"}, metadata=marker
|
||||
)
|
||||
await _put(
|
||||
saver, thread_id, first, 1, {"m": [1, 2]}, {"m": "2"}, metadata=marker
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_keeps_each_namespace_head_with_its_writes_and_blobs(
|
||||
name: str,
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id, other_id = str(uuid4()), str(uuid4())
|
||||
heads = {
|
||||
ns: await _three_checkpoints(saver, thread_id, ns) for ns in ("", "sub:1")
|
||||
}
|
||||
await _three_checkpoints(saver, other_id)
|
||||
before = {ns: await _call(saver, "get_tuple", h) for ns, h in heads.items()}
|
||||
other_before = await _rows(saver, other_id)
|
||||
|
||||
await _call(saver, "prune", [thread_id])
|
||||
|
||||
rows = await _rows(saver, thread_id)
|
||||
assert rows["checkpoints"] == [
|
||||
(ns, heads[ns]["configurable"]["checkpoint_id"]) for ns in ("", "sub:1")
|
||||
]
|
||||
assert [(ns, cid) for ns, cid, *_ in rows["writes"]] == rows["checkpoints"]
|
||||
assert rows["blobs"] == [
|
||||
(ns, ch, ver)
|
||||
for ns in ("", "sub:1")
|
||||
for ch, ver in (("a", "3"), ("s", "1"))
|
||||
]
|
||||
for ns, head in heads.items():
|
||||
assert await _call(saver, "get_tuple", head) == before[ns]._replace(
|
||||
parent_config=None
|
||||
)
|
||||
assert await _rows(saver, other_id) == other_before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("evidence", ["counters", "snapshot", "versioned"])
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
|
||||
name: str, evidence: str
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
ordinary_id, delta_id = str(uuid4()), str(uuid4())
|
||||
await _three_checkpoints(saver, ordinary_id)
|
||||
await _delta_thread(saver, delta_id, evidence)
|
||||
before = [await _rows(saver, t) for t in (ordinary_id, delta_id)]
|
||||
|
||||
with pytest.raises(NotImplementedError, match=delta_id):
|
||||
await _call(saver, "prune", [ordinary_id, delta_id])
|
||||
|
||||
assert [await _rows(saver, t) for t in (ordinary_id, delta_id)] == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_prunes_a_thread_storing_a_plain_true(name: str) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id = str(uuid4())
|
||||
first = await _put(
|
||||
saver, thread_id, None, 0, {"flag": True, "a": [1]}, {"flag": "1", "a": "1"}
|
||||
)
|
||||
head = await _put(
|
||||
saver, thread_id, first, 1, {"flag": True}, {"flag": "2", "a": "1"}
|
||||
)
|
||||
|
||||
await _call(saver, "prune", [thread_id])
|
||||
|
||||
assert (await _rows(saver, thread_id))["checkpoints"] == [
|
||||
("", head["configurable"]["checkpoint_id"])
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_keeps_the_parent_holding_a_pre_v4_heads_sends(
|
||||
name: str,
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id = str(uuid4())
|
||||
first = await _put(saver, thread_id, None, 0, {}, {}, v=1)
|
||||
await _call(saver, "put_writes", first, [(TASKS, "send-1")], str(uuid4()))
|
||||
second = await _put(saver, thread_id, first, 1, {}, {}, v=1)
|
||||
head = await _put(saver, thread_id, second, 2, {}, {}, v=1)
|
||||
await _call(saver, "put_writes", second, [(TASKS, "send-2")], str(uuid4()))
|
||||
|
||||
await _call(saver, "prune", [thread_id])
|
||||
|
||||
assert (await _rows(saver, thread_id))["checkpoints"] == [
|
||||
("", second["configurable"]["checkpoint_id"]),
|
||||
("", head["configurable"]["checkpoint_id"]),
|
||||
]
|
||||
loaded = await _call(saver, "get_tuple", head)
|
||||
assert loaded.checkpoint["channel_values"][TASKS] == ["send-2"]
|
||||
assert loaded.parent_config == second
|
||||
assert (await _call(saver, "get_tuple", second)).parent_config is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_delete_removes_every_row_of_the_threads(name: str) -> None:
|
||||
async with _saver(name) as saver:
|
||||
ordinary_id, delta_id, other_id = str(uuid4()), str(uuid4()), str(uuid4())
|
||||
await _three_checkpoints(saver, ordinary_id)
|
||||
await _delta_thread(saver, delta_id, "counters")
|
||||
await _three_checkpoints(saver, other_id)
|
||||
other_before = await _rows(saver, other_id)
|
||||
|
||||
await _call(saver, "prune", [ordinary_id, delta_id], strategy="delete")
|
||||
|
||||
empty = {"checkpoints": [], "writes": [], "blobs": []}
|
||||
assert await _rows(saver, ordinary_id) == empty
|
||||
assert await _rows(saver, delta_id) == empty
|
||||
assert await _rows(saver, other_id) == other_before
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_unknown_strategy_changes_nothing(name: str) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id = str(uuid4())
|
||||
await _three_checkpoints(saver, thread_id)
|
||||
before = await _rows(saver, thread_id)
|
||||
|
||||
with pytest.raises(ValueError, match="keep_all"):
|
||||
await _call(saver, "prune", [thread_id], strategy="keep_all")
|
||||
|
||||
assert await _rows(saver, thread_id) == before
|
||||
|
||||
|
||||
async def test_async_saver_prune_runs_from_another_thread_only() -> None:
|
||||
async with _saver("aio_base") as saver:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
|
||||
with pytest.raises(asyncio.InvalidStateError):
|
||||
saver.prune([thread_id])
|
||||
await asyncio.to_thread(saver.prune, [thread_id])
|
||||
|
||||
assert (await _rows(saver, thread_id))["checkpoints"] == [
|
||||
("", head["configurable"]["checkpoint_id"])
|
||||
]
|
||||
|
||||
|
||||
async def test_prune_conformance() -> None:
|
||||
@checkpointer_test(name="AsyncPostgresSaver")
|
||||
async def postgres_saver():
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
yield saver
|
||||
|
||||
report = await validate(postgres_saver, capabilities={"prune"})
|
||||
|
||||
prune = report.results["prune"]
|
||||
if not prune.passed:
|
||||
details = "\n".join(prune.failures)
|
||||
pytest.fail(f"Capability prune failed (passed={prune.passed}):\n{details}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("strategy", ["keep_latest", "delete"])
|
||||
@pytest.mark.parametrize("name", ["base", "aio_base"])
|
||||
async def test_prune_rejects_a_single_thread_id_string(
|
||||
name: str, strategy: str
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
for thread_id in ("ab", "a", "b"):
|
||||
await _three_checkpoints(saver, thread_id)
|
||||
before = [await _rows(saver, t) for t in ("ab", "a", "b")]
|
||||
|
||||
with pytest.raises(TypeError, match="not a str"):
|
||||
await _call(saver, "prune", "ab", strategy=strategy)
|
||||
|
||||
assert [await _rows(saver, t) for t in ("ab", "a", "b")] == before
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"after", [PRUNE_HEADS_SQL, PRUNE_REFERENCED_BLOBS_SQL], ids=["heads", "blobs"]
|
||||
)
|
||||
async def test_keep_latest_keeps_a_checkpoint_saved_while_it_runs(after: str) -> None:
|
||||
async with _saver("aio_base") as saver, _second_saver(saver) as other:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
saved: list[dict] = []
|
||||
|
||||
async def save_a_child() -> None:
|
||||
saved.append(
|
||||
await _put(
|
||||
other, thread_id, head, 3, {"a": [1, 2, 3, 4]}, {"a": "4", "s": "1"}
|
||||
)
|
||||
)
|
||||
|
||||
_run_once_after(saver, after, save_a_child)
|
||||
await saver.aprune([thread_id])
|
||||
|
||||
loaded = await saver.aget_tuple(saved[0])
|
||||
assert loaded is not None, "the checkpoint saved during the prune was deleted"
|
||||
assert loaded.checkpoint["channel_values"] == {
|
||||
"a": [1, 2, 3, 4],
|
||||
"s": ["shared"],
|
||||
}
|
||||
assert loaded.parent_config == head
|
||||
|
||||
|
||||
async def test_keep_latest_refuses_a_thread_that_turns_delta_after_the_batch_check() -> (
|
||||
None
|
||||
):
|
||||
async with _saver("aio_base") as saver, _second_saver(saver) as other:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
|
||||
async def save_a_delta_checkpoint() -> None:
|
||||
await _put(
|
||||
other,
|
||||
thread_id,
|
||||
head,
|
||||
3,
|
||||
{},
|
||||
{"a": "4", "s": "1"},
|
||||
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
|
||||
)
|
||||
|
||||
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint)
|
||||
with pytest.raises(NotImplementedError, match=thread_id):
|
||||
await saver.aprune([thread_id])
|
||||
|
||||
assert len((await _rows(saver, thread_id))["checkpoints"]) == 4
|
||||
|
||||
|
||||
async def test_keep_latest_keeps_the_history_of_a_delta_checkpoint_saved_after_the_thread_check() -> (
|
||||
None
|
||||
):
|
||||
async with _saver("aio_base") as saver, _second_saver(saver) as other:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
saved: list[dict] = []
|
||||
|
||||
async def save_a_delta_checkpoint() -> None:
|
||||
saved.append(
|
||||
await _put(
|
||||
other,
|
||||
thread_id,
|
||||
head,
|
||||
3,
|
||||
{},
|
||||
{"a": "4", "s": "1"},
|
||||
metadata={"counters_since_delta_snapshot": {"a": [1, 1]}},
|
||||
)
|
||||
)
|
||||
|
||||
_run_once_after(saver, PRUNE_DELTA_THREADS_SQL, save_a_delta_checkpoint, nth=2)
|
||||
await saver.aprune([thread_id])
|
||||
|
||||
assert (await _rows(saver, thread_id))["checkpoints"] == [
|
||||
("", head["configurable"]["checkpoint_id"]),
|
||||
("", saved[0]["configurable"]["checkpoint_id"]),
|
||||
]
|
||||
history = await saver.aget_delta_channel_history(
|
||||
config=saved[0], channels=["a"]
|
||||
)
|
||||
assert history["a"]["seed"] == [1, 2, 3]
|
||||
assert [value for _, _, value in history["a"]["writes"]] == [2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_deletes_a_subgraph_call_with_the_checkpoint_it_started_from(
|
||||
name: str,
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id = str(uuid4())
|
||||
older = await _put(saver, thread_id, None, 0, {}, {})
|
||||
head = await _put(saver, thread_id, older, 1, {}, {})
|
||||
unsaved = str(uuid6(clock_seq=2))
|
||||
started_from = {
|
||||
"finished:1": older["configurable"]["checkpoint_id"],
|
||||
"interrupted:2": head["configurable"]["checkpoint_id"],
|
||||
"running:3": unsaved,
|
||||
"checkpointer_true": older["configurable"]["checkpoint_id"],
|
||||
}
|
||||
for ns, root in started_from.items():
|
||||
await _three_checkpoints(
|
||||
saver, thread_id, ns, metadata={"parents": {"": root}}
|
||||
)
|
||||
|
||||
await _call(saver, "prune", [thread_id])
|
||||
|
||||
rows = await _rows(saver, thread_id)
|
||||
assert {ns for ns, *_ in rows["checkpoints"]} == {
|
||||
"",
|
||||
"interrupted:2",
|
||||
"running:3",
|
||||
"checkpointer_true",
|
||||
}
|
||||
assert {ns for ns, *_ in rows["writes"]} == {
|
||||
"interrupted:2",
|
||||
"running:3",
|
||||
"checkpointer_true",
|
||||
}
|
||||
assert {ns for ns, *_ in rows["blobs"]} == {
|
||||
"interrupted:2",
|
||||
"running:3",
|
||||
"checkpointer_true",
|
||||
}
|
||||
|
||||
|
||||
async def test_keep_latest_keeps_a_checkpoint_saved_meanwhile_whatever_its_id() -> None:
|
||||
async with _saver("aio_base") as saver, _second_saver(saver) as other:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
lower_id = "0" + head["configurable"]["checkpoint_id"][1:]
|
||||
saved: list[dict] = []
|
||||
|
||||
async def save_a_child_with_a_lower_id() -> None:
|
||||
child = await _put(
|
||||
other,
|
||||
thread_id,
|
||||
head,
|
||||
3,
|
||||
{"a": [1, 2, 3, 4]},
|
||||
{"a": "4", "s": "1"},
|
||||
checkpoint_id=lower_id,
|
||||
)
|
||||
await _call(other, "put_writes", child, [("a", 4)], str(uuid4()))
|
||||
saved.append(child)
|
||||
|
||||
_run_once_after(saver, PRUNE_HEADS_SQL, save_a_child_with_a_lower_id)
|
||||
await saver.aprune([thread_id])
|
||||
|
||||
loaded = await saver.aget_tuple(saved[0])
|
||||
assert loaded is not None, "the checkpoint saved during the prune was deleted"
|
||||
assert loaded.checkpoint["channel_values"] == {
|
||||
"a": [1, 2, 3, 4],
|
||||
"s": ["shared"],
|
||||
}
|
||||
assert [value for _, _, value in loaded.pending_writes or []] == [4]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", SAVERS)
|
||||
async def test_keep_latest_keeps_writes_saved_before_their_checkpoint(
|
||||
name: str,
|
||||
) -> None:
|
||||
async with _saver(name) as saver:
|
||||
thread_id = str(uuid4())
|
||||
head = await _three_checkpoints(saver, thread_id)
|
||||
next_id = str(uuid6(clock_seq=3))
|
||||
next_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": next_id,
|
||||
}
|
||||
}
|
||||
await _call(saver, "put_writes", next_config, [("a", 4)], str(uuid4()))
|
||||
|
||||
await _call(saver, "prune", [thread_id])
|
||||
await _put(
|
||||
saver, thread_id, head, 3, {}, {"a": "3", "s": "1"}, checkpoint_id=next_id
|
||||
)
|
||||
|
||||
loaded = await _call(saver, "get_tuple", next_config)
|
||||
assert [value for _, _, value in loaded.pending_writes] == [4]
|
||||
Generated
+1
-1
@@ -353,7 +353,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.2"
|
||||
version = "3.2.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.1"
|
||||
version = "3.1.2"
|
||||
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
Generated
+1
-1
@@ -360,7 +360,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.1"
|
||||
version = "3.1.2"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -2829,11 +2829,11 @@ graceful-fs@^4.2.11:
|
||||
integrity sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==
|
||||
|
||||
handlebars@^4.7.9:
|
||||
version "4.7.9"
|
||||
resolved "https://registry.yarnpkg.com/handlebars/-/handlebars-4.7.9.tgz#6f139082ab58dc4e5a0e51efe7db5ae890d56a0f"
|
||||
integrity sha512-4E71E0rpOaQuJR2A3xDZ+GM1HyWYv1clR58tC8emQNeQe3RH7MAzSbat+V0wG78LQBo6m6bzSG/L4pBuCsgnUQ==
|
||||
version "4.7.10"
|
||||
resolved "https://registry.yarnpkg.com/handlebars/-/handlebars-4.7.10.tgz#12419ea1b41391b3ffcf003f2164dcbce090c417"
|
||||
integrity sha512-P5VJMVM7qgBn6vjXMw8WG9uVI+ncf2pi72j4de4yz5ZULLj2RGqLYaKOYGsgyrViQ0tePOVlN1tDCCXXtFqXKg==
|
||||
dependencies:
|
||||
minimist "^1.2.5"
|
||||
minimist "^1.2.8"
|
||||
neo-async "^2.6.2"
|
||||
source-map "^0.6.1"
|
||||
wordwrap "^1.0.0"
|
||||
@@ -3920,7 +3920,7 @@ minimatch@^9.0.4:
|
||||
dependencies:
|
||||
brace-expansion "^2.0.2"
|
||||
|
||||
minimist@^1.2.0, minimist@^1.2.5, minimist@^1.2.6:
|
||||
minimist@^1.2.0, minimist@^1.2.6, minimist@^1.2.8:
|
||||
version "1.2.8"
|
||||
resolved "https://registry.yarnpkg.com/minimist/-/minimist-1.2.8.tgz#c1a464e7693302e082a075cee0c057741ac4772c"
|
||||
integrity sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==
|
||||
|
||||
@@ -1 +1 @@
|
||||
__version__ = "0.4.33"
|
||||
__version__ = "0.4.36"
|
||||
|
||||
@@ -231,15 +231,6 @@ class RequestedPlacement:
|
||||
def requested(self) -> bool:
|
||||
return self.listener_id is not None or self.k8s_namespace is not None
|
||||
|
||||
def ensure_not_requested(self, deployment_id: str) -> None:
|
||||
if self.requested:
|
||||
raise click.UsageError(
|
||||
"Listener and namespace are fixed when a deployment is created. "
|
||||
f"Deployment {deployment_id} already exists, so drop --listener-id "
|
||||
"and --k8s-namespace, or create a new deployment with a different "
|
||||
"--name."
|
||||
)
|
||||
|
||||
def on(self, listener: Listener) -> Placement:
|
||||
return OnListener(listener.id, self._namespace(listener))
|
||||
|
||||
@@ -1055,15 +1046,28 @@ def _print_deployment_result(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _strip_url_scheme(url: str) -> str:
|
||||
"""Return a registry reference without a trailing slash or URL scheme."""
|
||||
reference = url.rstrip("/")
|
||||
if "://" in reference:
|
||||
reference = reference.split("//", 1)[1]
|
||||
return reference
|
||||
|
||||
|
||||
# Artifact Registry's username for access tokens. Control planes whose registry
|
||||
# needs another one return it as `username`.
|
||||
_DEFAULT_REGISTRY_USERNAME = "oauth2accesstoken"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _docker_config_for_token(registry_host: str, token: str):
|
||||
def _docker_config_for_token(registry_host: str, token: str, username: str):
|
||||
"""Create a temporary Docker config with only the push token.
|
||||
|
||||
Yields the path to a temporary config directory that can be passed
|
||||
to `docker --config <path>` so that system credential helpers
|
||||
(e.g. gcloud) don't interfere with the push token.
|
||||
"""
|
||||
auth_b64 = base64.b64encode(f"oauth2accesstoken:{token}".encode()).decode()
|
||||
auth_b64 = base64.b64encode(f"{username}:{token}".encode()).decode()
|
||||
config_data = {"auths": {registry_host: {"auth": auth_b64}}}
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
with open(os.path.join(tmpdir, "config.json"), "w") as f:
|
||||
@@ -1201,10 +1205,37 @@ class DockerBuildCommand:
|
||||
return cls(("docker", "build"), ())
|
||||
flags: tuple[str, ...] = ("--platform", _DEPLOYMENT_PLATFORM, "--load")
|
||||
if not verbose:
|
||||
flags += ("--progress=quiet",)
|
||||
flags += ("--progress=plain",)
|
||||
return cls(("docker", "buildx", "build"), flags)
|
||||
|
||||
|
||||
def _active_docker_host(runner: CommandRunner) -> str | None:
|
||||
"""Return the daemon endpoint of the user's active Docker context.
|
||||
|
||||
None when `DOCKER_HOST` is set, which Docker honors with any config
|
||||
directory. Also None when the lookup fails (for example, `DOCKER_CONTEXT`
|
||||
names a context that doesn't exist) or the context has no Docker endpoint
|
||||
(Docker prints `<no value>`); that keeps Docker's own default in place.
|
||||
"""
|
||||
if os.environ.get("DOCKER_HOST"):
|
||||
return None
|
||||
try:
|
||||
stdout, _ = runner.run(
|
||||
subp_exec(
|
||||
"docker",
|
||||
"context",
|
||||
"inspect",
|
||||
"--format",
|
||||
"{{.Endpoints.docker.Host}}",
|
||||
collect=True,
|
||||
)
|
||||
)
|
||||
except (click.exceptions.Exit, OSError):
|
||||
return None
|
||||
host = (stdout or "").strip()
|
||||
return host if "://" in host else None
|
||||
|
||||
|
||||
def _docker_argv(docker_config_dir: str | None) -> tuple[str, ...]:
|
||||
if docker_config_dir is None:
|
||||
return ("docker",)
|
||||
@@ -1238,17 +1269,13 @@ def _push_image(
|
||||
runner: CommandRunner,
|
||||
image: str,
|
||||
*,
|
||||
docker_config_dir: str | None,
|
||||
docker: tuple[str, ...],
|
||||
verbose: bool,
|
||||
) -> None:
|
||||
for attempt in range(1, _PUSH_ATTEMPTS + 1):
|
||||
try:
|
||||
with Progress(message="Pushing...", elapsed=not verbose):
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*_docker_argv(docker_config_dir), "push", image, verbose=verbose
|
||||
)
|
||||
)
|
||||
runner.run(subp_exec(*docker, "push", image, verbose=verbose))
|
||||
return
|
||||
except click.exceptions.Exit:
|
||||
if attempt == _PUSH_ATTEMPTS:
|
||||
@@ -1317,21 +1344,27 @@ def _run_local_build(
|
||||
raise click.ClickException(
|
||||
"Push token response missing token or registry_url"
|
||||
)
|
||||
username = push_data.get("username") or _DEFAULT_REGISTRY_USERNAME
|
||||
step += 1
|
||||
|
||||
normalized_registry = registry_url.rstrip("/")
|
||||
if "://" in normalized_registry:
|
||||
normalized_registry = normalized_registry.split("//", 1)[1]
|
||||
repo_seed = image_name or name or spec.config.parent.name
|
||||
remote_image = str(
|
||||
ImageReference(
|
||||
f"{normalized_registry}/{normalize_name(repo_seed)}",
|
||||
tag,
|
||||
if push_data.get("repository_url"):
|
||||
# The control plane chose the exact repository (an ECR repository
|
||||
# is a single image name, so nothing may be appended to it).
|
||||
repository = _strip_url_scheme(push_data["repository_url"])
|
||||
else:
|
||||
repo_seed = image_name or name or spec.config.parent.name
|
||||
repository = (
|
||||
f"{_strip_url_scheme(registry_url)}/{normalize_name(repo_seed)}"
|
||||
)
|
||||
)
|
||||
registry_host = normalized_registry.split("/")[0]
|
||||
remote_image = str(ImageReference(repository, tag))
|
||||
registry_host = repository.split("/")[0]
|
||||
|
||||
with _docker_config_for_token(registry_host, deployment_token) as cfg:
|
||||
with _docker_config_for_token(registry_host, deployment_token, username) as cfg:
|
||||
# A config directory has no current context, so Docker would fall
|
||||
# back to /var/run/docker.sock; keep using the daemon the build used.
|
||||
docker_host = _active_docker_host(runner)
|
||||
host_args = ("--host", docker_host) if docker_host else ()
|
||||
docker = (*_docker_argv(cfg), *host_args)
|
||||
_log_deploy_step(step, f"Logging into {registry_host}")
|
||||
token_input = (
|
||||
deployment_token
|
||||
@@ -1340,10 +1373,10 @@ def _run_local_build(
|
||||
)
|
||||
runner.run(
|
||||
subp_exec(
|
||||
*_docker_argv(cfg),
|
||||
*docker,
|
||||
"login",
|
||||
"-u",
|
||||
"oauth2accesstoken",
|
||||
username,
|
||||
"--password-stdin",
|
||||
registry_host,
|
||||
input=token_input,
|
||||
@@ -1356,7 +1389,7 @@ def _run_local_build(
|
||||
runner.run(
|
||||
subp_exec("docker", "tag", image_to_push, remote_image, verbose=verbose)
|
||||
)
|
||||
_push_image(runner, remote_image, docker_config_dir=cfg, verbose=verbose)
|
||||
_push_image(runner, remote_image, docker=docker, verbose=verbose)
|
||||
step += 1
|
||||
|
||||
resolved_image = _resolve_pushed_image_digest(
|
||||
@@ -1647,10 +1680,72 @@ class PublishedImage:
|
||||
ImageSource = BuildAndPush | PublishedImage
|
||||
|
||||
|
||||
def _resource_spec(image_pull_secrets: tuple[str, ...]) -> Mapping[str, object]:
|
||||
if not image_pull_secrets:
|
||||
return _OPERATOR_DEFAULT_RESOURCE_SPEC
|
||||
return {
|
||||
"image_pull_secrets": [{"name": name} for name in image_pull_secrets],
|
||||
}
|
||||
|
||||
|
||||
# Mirrors the control plane's own LocalObjectReference.name validation (a
|
||||
# Kubernetes DNS subdomain), so a bad secret name is rejected here instead of
|
||||
# after a build and push.
|
||||
_K8S_SECRET_NAME_RE = re.compile(
|
||||
r"^[a-z0-9]([-a-z0-9]*[a-z0-9])?(\.[a-z0-9]([-a-z0-9]*[a-z0-9])?)*$"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CreateOnlySettings:
|
||||
"""Options for an external-image deployment that only apply when
|
||||
creating one, never when updating an existing deployment."""
|
||||
|
||||
placement: RequestedPlacement
|
||||
image_pull_secrets: tuple[str, ...] = ()
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
for name in self.image_pull_secrets:
|
||||
if len(name) > 253 or not _K8S_SECRET_NAME_RE.match(name):
|
||||
raise click.UsageError(
|
||||
f"--image-pull-secret {name!r} is not a valid Kubernetes "
|
||||
"secret name (lowercase alphanumeric characters, '-', or "
|
||||
"'.' only, max 253 characters)."
|
||||
)
|
||||
|
||||
@property
|
||||
def requested(self) -> bool:
|
||||
return self.placement.requested or bool(self.image_pull_secrets)
|
||||
|
||||
def ensure_creatable(self, push_to: str | None, image_uri: str | None) -> None:
|
||||
if push_to is None and image_uri is None and self.requested:
|
||||
raise click.UsageError(
|
||||
"--listener-id, --k8s-namespace, and --image-pull-secret only "
|
||||
"apply when creating a deployment with --push-to or --image-uri."
|
||||
)
|
||||
|
||||
def ensure_not_existing(self, deployment_id: str) -> None:
|
||||
if self.placement.requested:
|
||||
raise click.UsageError(
|
||||
"Listener and namespace are fixed when a deployment is "
|
||||
f"created. Deployment {deployment_id} already exists, so "
|
||||
"drop --listener-id and --k8s-namespace, or create a new "
|
||||
"deployment with a different --name."
|
||||
)
|
||||
if self.image_pull_secrets:
|
||||
raise click.UsageError(
|
||||
"langgraph deploy does not update image pull secrets on an "
|
||||
f"existing deployment. Deployment {deployment_id} already "
|
||||
"exists, so drop --image-pull-secret and update resource_spec "
|
||||
"through the API or the LangSmith UI instead, or create a "
|
||||
"new deployment with a different --name."
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CustomerRegistrySource:
|
||||
image: ImageSource
|
||||
requested_placement: RequestedPlacement
|
||||
create_only: CreateOnlySettings
|
||||
|
||||
def run(self, ctx: DeployContext) -> DeployOutcome:
|
||||
if isinstance(ctx.selector, ById):
|
||||
@@ -1672,7 +1767,7 @@ class CustomerRegistrySource:
|
||||
self, ctx: DeployContext, existing: ExistingDeployment, step: int
|
||||
) -> DeployOutcome:
|
||||
_ensure_customer_registry_source(existing)
|
||||
self.requested_placement.ensure_not_requested(existing.id)
|
||||
self.create_only.ensure_not_existing(existing.id)
|
||||
image_uri, step = self._publish(ctx, step)
|
||||
_log_deploy_step(step, f"Updating deployment {existing.id}")
|
||||
updated = ctx.client.update_deployment(
|
||||
@@ -1687,7 +1782,7 @@ class CustomerRegistrySource:
|
||||
)
|
||||
|
||||
def _resolve_placement(self, ctx: DeployContext) -> Placement:
|
||||
requested = self.requested_placement
|
||||
requested = self.create_only.placement
|
||||
if requested.listener_id is not None:
|
||||
return requested.on(_requested_listener(ctx.client, requested.listener_id))
|
||||
if not (ctx.endpoints.is_cloud or requested.requested):
|
||||
@@ -1716,7 +1811,9 @@ class CustomerRegistrySource:
|
||||
else None,
|
||||
source=_CUSTOMER_REGISTRY_SOURCE,
|
||||
source_config={
|
||||
"resource_spec": _OPERATOR_DEFAULT_RESOURCE_SPEC,
|
||||
"resource_spec": _resource_spec(
|
||||
self.create_only.image_pull_secrets
|
||||
),
|
||||
**placement.source_config(),
|
||||
},
|
||||
source_revision_config={"image_uri": image_uri},
|
||||
@@ -1757,7 +1854,7 @@ class CustomerRegistrySource:
|
||||
_build_image(runner, ctx.spec, image, verbose=ctx.verbose)
|
||||
step += 1
|
||||
_log_deploy_step(step, f"Pushing image {image}")
|
||||
_push_image(runner, image, docker_config_dir=None, verbose=ctx.verbose)
|
||||
_push_image(runner, image, docker=_docker_argv(None), verbose=ctx.verbose)
|
||||
step += 1
|
||||
digest = _resolve_pushed_image_digest(
|
||||
runner, remote_image=image, docker_config_dir=None, verbose=ctx.verbose
|
||||
@@ -1810,19 +1907,13 @@ def _select_source(
|
||||
tag: str | None,
|
||||
remote_build_flag: bool | None,
|
||||
placement: RequestedPlacement,
|
||||
image_pull_secrets: tuple[str, ...],
|
||||
selector: DeploymentSelector,
|
||||
) -> DeploymentSource:
|
||||
if push_to is None and image_uri is None and placement.requested:
|
||||
raise click.UsageError(
|
||||
"--listener-id and --k8s-namespace only apply when creating a "
|
||||
"deployment with --push-to or --image-uri."
|
||||
)
|
||||
if placement.requested and isinstance(selector, ById):
|
||||
raise click.UsageError(
|
||||
"Listener and namespace are fixed when a deployment is created, so "
|
||||
"they cannot be set for an existing --deployment-id. Drop them, or "
|
||||
"create a new deployment with --name."
|
||||
)
|
||||
create_only = CreateOnlySettings(placement, image_pull_secrets)
|
||||
create_only.ensure_creatable(push_to, image_uri)
|
||||
if create_only.requested and isinstance(selector, ById):
|
||||
create_only.ensure_not_existing(selector.deployment_id)
|
||||
if image_uri is not None:
|
||||
if push_to is not None:
|
||||
raise click.UsageError("--image-uri cannot be combined with --push-to.")
|
||||
@@ -1834,7 +1925,7 @@ def _select_source(
|
||||
raise click.UsageError("--image-uri cannot be combined with --remote.")
|
||||
return CustomerRegistrySource(
|
||||
image=PublishedImage(_validate_image_uri(image_uri)),
|
||||
requested_placement=placement,
|
||||
create_only=create_only,
|
||||
)
|
||||
if push_to is not None:
|
||||
if remote_build_flag is True:
|
||||
@@ -1844,7 +1935,7 @@ def _select_source(
|
||||
_require_local_docker()
|
||||
return CustomerRegistrySource(
|
||||
image=BuildAndPush(reference=reference, prebuilt_image=image),
|
||||
requested_placement=placement,
|
||||
create_only=create_only,
|
||||
)
|
||||
if image and remote_build_flag is True:
|
||||
raise click.UsageError("--image cannot be combined with --remote builds.")
|
||||
@@ -1874,7 +1965,10 @@ def _create_host_backend_client(
|
||||
host_url: str | None,
|
||||
api_key: str | None,
|
||||
env_vars: dict[str, str] | None = None,
|
||||
*,
|
||||
no_input: bool | None = None,
|
||||
) -> HostBackendClient:
|
||||
effective_no_input = _no_input if no_input is None else no_input
|
||||
if env_vars is None:
|
||||
env_vars = _parse_env_from_config({}, pathlib.Path.cwd() / DEFAULT_CONFIG)
|
||||
resolved_api_key = api_key
|
||||
@@ -1889,7 +1983,7 @@ def _create_host_backend_client(
|
||||
resolved_api_key = val
|
||||
break
|
||||
if not resolved_api_key:
|
||||
if _no_input:
|
||||
if effective_no_input:
|
||||
raise click.ClickException(
|
||||
"No LangSmith API key found. Set LANGSMITH_API_KEY in the "
|
||||
"environment or .env file."
|
||||
@@ -1914,6 +2008,8 @@ def _create_host_backend_client(
|
||||
def _call_host_backend_with_optional_tenant(
|
||||
client: HostBackendClient,
|
||||
operation: Callable[[HostBackendClient], _T],
|
||||
*,
|
||||
no_input: bool | None = None,
|
||||
) -> _T:
|
||||
"""Run *operation*, prompting for a workspace ID on org-scoped 403s.
|
||||
|
||||
@@ -1922,6 +2018,7 @@ def _call_host_backend_with_optional_tenant(
|
||||
in-place so all subsequent calls through the same instance are
|
||||
tenant-aware.
|
||||
"""
|
||||
effective_no_input = _no_input if no_input is None else no_input
|
||||
prompted_for_tenant = False
|
||||
|
||||
while True:
|
||||
@@ -1933,7 +2030,7 @@ def _call_host_backend_with_optional_tenant(
|
||||
and err.status_code == 403
|
||||
and "requires workspace specification" in err.message
|
||||
):
|
||||
if _no_input:
|
||||
if effective_no_input:
|
||||
raise click.ClickException(
|
||||
"API key is org-scoped and requires a workspace ID. "
|
||||
"Set LANGSMITH_TENANT_ID in your .env file or "
|
||||
@@ -2116,11 +2213,13 @@ def _deploy_base_options(
|
||||
),
|
||||
click.option(
|
||||
"--deployment-type",
|
||||
type=click.Choice(["dev", "prod"]),
|
||||
type=click.Choice(["serverless", "dedicated", "dev", "prod"]),
|
||||
default="dev",
|
||||
show_default=True,
|
||||
help=(
|
||||
"Deployment type (used when creating a new deployment). "
|
||||
"Deployment type (used when creating a new deployment): "
|
||||
"serverless (scale to zero) or dedicated (prod); "
|
||||
"legacy dev and prod are also accepted. "
|
||||
"Ignored with --push-to."
|
||||
),
|
||||
),
|
||||
@@ -2182,6 +2281,18 @@ def _deploy_base_options(
|
||||
"creating a deployment with --push-to or --image-uri."
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--image-pull-secret",
|
||||
"image_pull_secrets",
|
||||
multiple=True,
|
||||
help=(
|
||||
"Name of a Kubernetes secret (type kubernetes.io/"
|
||||
"dockerconfigjson) in the deployment's namespace to use for "
|
||||
"pulling the image from a private registry. Repeat to give "
|
||||
"more than one. Only used when creating a deployment with "
|
||||
"--push-to or --image-uri."
|
||||
),
|
||||
),
|
||||
click.option(
|
||||
"--config",
|
||||
"-c",
|
||||
@@ -2295,6 +2406,7 @@ def _deploy_cmd(
|
||||
image_uri: str | None,
|
||||
listener_id: str | None,
|
||||
k8s_namespace: str | None,
|
||||
image_pull_secrets: tuple[str, ...],
|
||||
tag: str | None,
|
||||
base_image: str | None,
|
||||
install_command: str | None,
|
||||
@@ -2371,6 +2483,7 @@ def _deploy_cmd(
|
||||
tag=tag,
|
||||
remote_build_flag=remote_build_flag,
|
||||
placement=RequestedPlacement(listener_id, k8s_namespace),
|
||||
image_pull_secrets=image_pull_secrets,
|
||||
selector=selector,
|
||||
)
|
||||
|
||||
@@ -2397,7 +2510,9 @@ def _deploy_cmd(
|
||||
),
|
||||
verbose=verbose,
|
||||
selector=selector,
|
||||
deployment_type=deployment_type,
|
||||
deployment_type={"serverless": "dev_zero", "dedicated": "prod"}.get(
|
||||
deployment_type, deployment_type
|
||||
),
|
||||
secrets=secrets,
|
||||
tracked_packages=tracked_packages,
|
||||
)
|
||||
@@ -2518,6 +2633,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 +2648,16 @@ def deploy_listeners() -> None:
|
||||
"--listener-id <id>` to deploy through it."
|
||||
),
|
||||
)
|
||||
def deploy_listeners_list(api_key: str | None, host_url: str | None) -> None:
|
||||
client = _create_host_backend_client(host_url, api_key)
|
||||
def deploy_listeners_list(
|
||||
api_key: str | None, host_url: str | None, json_output: bool
|
||||
) -> None:
|
||||
client = _create_host_backend_client(host_url, api_key, no_input=json_output)
|
||||
listeners = _call_host_backend_with_optional_tenant(
|
||||
client, lambda c: c.list_listeners()
|
||||
client, lambda c: c.list_listeners(), no_input=json_output
|
||||
)
|
||||
if json_output:
|
||||
click.echo(json_mod.dumps(listeners))
|
||||
return
|
||||
if not listeners:
|
||||
click.echo("No listeners found for this workspace.")
|
||||
return
|
||||
|
||||
@@ -249,9 +249,12 @@ class HostBackendClient:
|
||||
return self._request("DELETE", f"/v2/deployments/{deployment_id}")
|
||||
|
||||
def request_push_token(self, deployment_id: str) -> dict[str, Any]:
|
||||
# Tells the control plane this client logs in with the returned
|
||||
# `username` and pushes to `repository_url`. LangSmith AWS requires it.
|
||||
return self._request(
|
||||
"POST",
|
||||
f"/v2/deployments/{deployment_id}/push-token",
|
||||
{"supports_registry_username": True},
|
||||
)
|
||||
|
||||
def request_upload_url(self, deployment_id: str) -> dict[str, Any]:
|
||||
|
||||
@@ -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,121 @@ 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:
|
||||
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:
|
||||
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."""
|
||||
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] = {}
|
||||
|
||||
@@ -52,6 +52,8 @@ ARCHIVE = ("/tmp/src.tgz", 2048, "langgraph.json")
|
||||
OBJECT_PATH = "tarballs/src.tgz"
|
||||
PLATFORM_FORMAT = "{{.Os}}/{{.Architecture}}"
|
||||
DIGESTS_FORMAT = "{{json .RepoDigests}}"
|
||||
CONTEXT_HOST_FORMAT = "{{.Endpoints.docker.Host}}"
|
||||
CONTEXT_HOST = "unix:///Users/dev/.docker/run/docker.sock"
|
||||
NOT_A_CLI_DEPLOYMENT = (
|
||||
"push token is only available for 'internal_docker' source deployments"
|
||||
)
|
||||
@@ -93,6 +95,7 @@ class ControlPlaneDouble:
|
||||
timeline: list[str]
|
||||
existing_deployments: list[dict] = field(default_factory=list)
|
||||
push_token_status: int = 200
|
||||
push_token_extra: dict = field(default_factory=dict)
|
||||
create_error: str | None = None
|
||||
listeners: list[dict] = field(default_factory=list)
|
||||
listeners_by_id: dict[str, dict] = field(default_factory=dict)
|
||||
@@ -147,7 +150,12 @@ class ControlPlaneDouble:
|
||||
if self.push_token_status != 200:
|
||||
return httpx.Response(self.push_token_status, text=NOT_A_CLI_DEPLOYMENT)
|
||||
return httpx.Response(
|
||||
200, json={"token": PUSH_TOKEN, "registry_url": REGISTRY_URL}
|
||||
200,
|
||||
json={
|
||||
"token": PUSH_TOKEN,
|
||||
"registry_url": REGISTRY_URL,
|
||||
**self.push_token_extra,
|
||||
},
|
||||
)
|
||||
if path.endswith("/upload-url"):
|
||||
return httpx.Response(
|
||||
@@ -186,6 +194,8 @@ class DockerCommand:
|
||||
class DockerDouble:
|
||||
timeline: list[str]
|
||||
failing_pushes: int = 0
|
||||
context_lookup_fails: bool = False
|
||||
context_host: str = CONTEXT_HOST
|
||||
builds: list[dict] = field(default_factory=list)
|
||||
commands: list[DockerCommand] = field(default_factory=list)
|
||||
|
||||
@@ -229,6 +239,10 @@ class DockerDouble:
|
||||
if "push" in args and self.failing_pushes > 0:
|
||||
self.failing_pushes -= 1
|
||||
raise click.exceptions.Exit(1)
|
||||
if CONTEXT_HOST_FORMAT in args:
|
||||
if self.context_lookup_fails:
|
||||
raise click.exceptions.Exit(1)
|
||||
return f"{self.context_host}\n", None
|
||||
if PLATFORM_FORMAT in args:
|
||||
return "linux/amd64\n", None
|
||||
if DIGESTS_FORMAT in args:
|
||||
@@ -238,6 +252,8 @@ class DockerDouble:
|
||||
|
||||
@staticmethod
|
||||
def _verb(args: tuple[str, ...]) -> str:
|
||||
if CONTEXT_HOST_FORMAT in args:
|
||||
return "inspect-context"
|
||||
if PLATFORM_FORMAT in args:
|
||||
return "inspect-platform"
|
||||
if DIGESTS_FORMAT in args:
|
||||
@@ -286,7 +302,12 @@ def deploy_project(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> DeployPro
|
||||
json.dumps({"graphs": {"agent": "agent.py:graph"}, "dependencies": ["."]})
|
||||
)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
for name in ("LANGSMITH_TENANT_ID", "LANGSMITH_ENDPOINT", "LANGGRAPH_HOST_URL"):
|
||||
for name in (
|
||||
"LANGSMITH_TENANT_ID",
|
||||
"LANGSMITH_ENDPOINT",
|
||||
"LANGGRAPH_HOST_URL",
|
||||
"DOCKER_HOST",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
|
||||
timeline: list[str] = []
|
||||
@@ -334,6 +355,7 @@ def test_first_local_deploy_creates_then_builds_pushes_and_updates_in_order(
|
||||
CREATE_DEPLOYMENT,
|
||||
"docker build",
|
||||
_push_token(CREATED_ID),
|
||||
"docker inspect-context",
|
||||
"docker login",
|
||||
"docker tag",
|
||||
"docker push",
|
||||
@@ -343,6 +365,27 @@ def test_first_local_deploy_creates_then_builds_pushes_and_updates_in_order(
|
||||
assert "Deployment updated" in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_type", "wire_type"),
|
||||
[
|
||||
("serverless", "dev_zero"),
|
||||
("dedicated", "prod"),
|
||||
("dev", "dev"),
|
||||
("prod", "prod"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("build_mode", ["--no-remote", "--remote"])
|
||||
def test_deployment_type_aliases(
|
||||
deploy_project: DeployProject, deployment_type: str, wire_type: str, build_mode: str
|
||||
) -> None:
|
||||
result = deploy_project.run(build_mode, "--deployment-type", deployment_type)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert deploy_project.control_plane.bodies[CREATE_DEPLOYMENT]["source_config"] == {
|
||||
"deployment_type": wire_type
|
||||
}
|
||||
|
||||
|
||||
def test_first_local_deploy_creates_an_internal_docker_deployment(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
@@ -363,7 +406,7 @@ def test_first_local_deploy_creates_an_internal_docker_deployment(
|
||||
pytest.param(
|
||||
"arm64",
|
||||
("docker", "buildx", "build"),
|
||||
("--platform", "linux/amd64", "--load", "--progress=quiet"),
|
||||
("--platform", "linux/amd64", "--load", "--progress=plain"),
|
||||
id="apple_silicon_cross_builds_for_linux_amd64",
|
||||
),
|
||||
pytest.param(
|
||||
@@ -400,7 +443,7 @@ def test_local_deploy_logs_in_with_the_control_plane_push_token(
|
||||
|
||||
login = deploy_project.docker.command("login")
|
||||
assert login.args[:2] == ("docker", "--config")
|
||||
assert login.args[3:] == (
|
||||
assert login.args[-5:] == (
|
||||
"login",
|
||||
"-u",
|
||||
"oauth2accesstoken",
|
||||
@@ -410,6 +453,95 @@ def test_local_deploy_logs_in_with_the_control_plane_push_token(
|
||||
assert login.kwargs["input"] == f"{PUSH_TOKEN}\n"
|
||||
|
||||
|
||||
def test_local_deploy_logs_in_and_pushes_through_the_active_docker_context(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
"""The temporary config has no current context, so without the active
|
||||
context's endpoint Docker would push through /var/run/docker.sock."""
|
||||
result = deploy_project.run("--no-remote")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
for verb in ("login", "push"):
|
||||
args = deploy_project.docker.command(verb).args
|
||||
assert args[3:5] == ("--host", CONTEXT_HOST), verb
|
||||
|
||||
|
||||
def test_local_deploy_leaves_an_explicit_docker_host_to_docker(
|
||||
deploy_project: DeployProject, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("DOCKER_HOST", "tcp://127.0.0.1:2375")
|
||||
|
||||
result = deploy_project.run("--no-remote")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
for verb in ("login", "push"):
|
||||
assert "--host" not in deploy_project.docker.command(verb).args, verb
|
||||
|
||||
|
||||
def test_local_deploy_keeps_dockers_default_when_the_context_lookup_fails(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.docker.context_lookup_fails = True
|
||||
|
||||
result = deploy_project.run("--no-remote")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "--host" not in deploy_project.docker.command("push").args
|
||||
|
||||
|
||||
def test_local_deploy_ignores_a_context_without_a_docker_endpoint(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.docker.context_host = "<no value>"
|
||||
|
||||
result = deploy_project.run("--no-remote")
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "--host" not in deploy_project.docker.command("push").args
|
||||
|
||||
|
||||
def test_local_deploy_tells_the_control_plane_it_honors_registry_usernames(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.run("--no-remote")
|
||||
|
||||
assert deploy_project.control_plane.bodies[_push_token(CREATED_ID)] == {
|
||||
"supports_registry_username": True
|
||||
}
|
||||
|
||||
|
||||
ECR_REPOSITORY = "123456789012.dkr.ecr.us-east-2.amazonaws.com/lsd/agent-1"
|
||||
ECR_PUSH_TOKEN_FIELDS = {"username": "AWS", "repository_url": ECR_REPOSITORY}
|
||||
|
||||
|
||||
def test_local_deploy_logs_in_with_the_username_the_control_plane_returns(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.control_plane.push_token_extra = ECR_PUSH_TOKEN_FIELDS
|
||||
|
||||
deploy_project.run("--no-remote")
|
||||
|
||||
login = deploy_project.docker.command("login")
|
||||
assert login.args[-5:] == (
|
||||
"login",
|
||||
"-u",
|
||||
"AWS",
|
||||
"--password-stdin",
|
||||
"123456789012.dkr.ecr.us-east-2.amazonaws.com",
|
||||
)
|
||||
|
||||
|
||||
def test_local_deploy_pushes_to_exactly_the_returned_repository(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.control_plane.push_token_extra = ECR_PUSH_TOKEN_FIELDS
|
||||
|
||||
deploy_project.run("--no-remote")
|
||||
|
||||
tag = deploy_project.docker.command("tag")
|
||||
assert tag.args[-1] == f"{ECR_REPOSITORY}:latest"
|
||||
|
||||
|
||||
def test_local_deploy_tags_the_build_into_the_token_registry(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
@@ -458,6 +590,7 @@ def test_prebuilt_image_is_validated_and_pushed_without_a_build(
|
||||
assert deploy_project.docker.builds == []
|
||||
assert deploy_project.docker.verbs() == [
|
||||
"docker inspect-platform",
|
||||
"docker inspect-context",
|
||||
"docker login",
|
||||
"docker tag",
|
||||
"docker push",
|
||||
@@ -918,6 +1051,100 @@ def test_self_hosted_control_plane_places_when_asked(
|
||||
}
|
||||
|
||||
|
||||
def test_image_pull_secret_is_included_in_resource_spec_on_create(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
result = deploy_project.run(
|
||||
"--push-to",
|
||||
PUSH_REPOSITORY,
|
||||
"--image-pull-secret",
|
||||
"regcred",
|
||||
"--image-pull-secret",
|
||||
"other-secret",
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert deploy_project.control_plane.bodies[CREATE_DEPLOYMENT]["source_config"] == {
|
||||
"resource_spec": {
|
||||
"image_pull_secrets": [
|
||||
{"name": "regcred"},
|
||||
{"name": "other-secret"},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_image_pull_secret_works_with_image_uri_too(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
result = deploy_project.run(
|
||||
"--image-uri", EXTERNAL_DIGEST, "--image-pull-secret", "regcred"
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert deploy_project.control_plane.bodies[CREATE_DEPLOYMENT]["source_config"] == {
|
||||
"resource_spec": {"image_pull_secrets": [{"name": "regcred"}]},
|
||||
}
|
||||
|
||||
|
||||
def test_image_pull_secret_without_push_to_or_image_uri_is_refused(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
result = deploy_project.run("--image-pull-secret", "regcred")
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "only apply when creating a deployment with --push-to" in result.output
|
||||
assert deploy_project.timeline == []
|
||||
|
||||
|
||||
def test_an_invalid_image_pull_secret_name_is_rejected_before_any_docker_work(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
result = deploy_project.run(
|
||||
"--push-to", PUSH_REPOSITORY, "--image-pull-secret", "Not_Valid!"
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "not a valid Kubernetes" in result.output
|
||||
assert deploy_project.docker.verbs() == []
|
||||
assert deploy_project.timeline == []
|
||||
|
||||
|
||||
def test_image_pull_secret_cannot_update_an_existing_deployment(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
result = deploy_project.run(
|
||||
"--deployment-id",
|
||||
"dep-1",
|
||||
"--push-to",
|
||||
PUSH_REPOSITORY,
|
||||
"--image-pull-secret",
|
||||
"regcred",
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "--image-pull-secret" in result.output
|
||||
assert "dep-1" in result.output
|
||||
assert deploy_project.timeline == []
|
||||
|
||||
|
||||
def test_image_pull_secret_cannot_update_a_deployment_found_by_name(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
deploy_project.control_plane.existing_deployments = [
|
||||
{"id": "dep-ext", "name": "my-app", "source": "external_docker"}
|
||||
]
|
||||
|
||||
result = deploy_project.run(
|
||||
"--image-uri", EXTERNAL_DIGEST, "--image-pull-secret", "regcred"
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "--image-pull-secret" in result.output
|
||||
assert deploy_project.docker.verbs() == []
|
||||
assert _patch("dep-ext") not in deploy_project.timeline
|
||||
|
||||
|
||||
def test_updating_a_deployment_never_looks_up_listeners(
|
||||
deploy_project: DeployProject,
|
||||
) -> None:
|
||||
|
||||
@@ -16,6 +16,7 @@ from langgraph_cli.deploy import (
|
||||
BuildAndPush,
|
||||
ById,
|
||||
ByName,
|
||||
CreateOnlySettings,
|
||||
CustomerRegistrySource,
|
||||
DockerBuildCommand,
|
||||
ExistingDeployment,
|
||||
@@ -34,6 +35,7 @@ from langgraph_cli.deploy import (
|
||||
_parse_env_from_config,
|
||||
_resolve_env_path,
|
||||
_resolve_pushed_image_digest,
|
||||
_resource_spec,
|
||||
_select_source,
|
||||
_validate_prebuilt_image,
|
||||
find_deployment_by_name,
|
||||
@@ -46,7 +48,9 @@ from langgraph_cli.image_reference import ImageReference
|
||||
|
||||
class TestDockerConfigForToken:
|
||||
def test_creates_config_json(self):
|
||||
with _docker_config_for_token("us-docker.pkg.dev", "my-token") as cfg:
|
||||
with _docker_config_for_token(
|
||||
"us-docker.pkg.dev", "my-token", "oauth2accesstoken"
|
||||
) as cfg:
|
||||
config_path = os.path.join(cfg, "config.json")
|
||||
assert os.path.isfile(config_path)
|
||||
with open(config_path) as f:
|
||||
@@ -54,13 +58,26 @@ class TestDockerConfigForToken:
|
||||
expected_auth = base64.b64encode(b"oauth2accesstoken:my-token").decode()
|
||||
assert data == {"auths": {"us-docker.pkg.dev": {"auth": expected_auth}}}
|
||||
|
||||
def test_uses_the_given_username(self):
|
||||
with _docker_config_for_token(
|
||||
"123456789012.dkr.ecr.us-east-2.amazonaws.com", "ecr-password", "AWS"
|
||||
) as cfg:
|
||||
with open(os.path.join(cfg, "config.json")) as f:
|
||||
data = json.load(f)
|
||||
expected_auth = base64.b64encode(b"AWS:ecr-password").decode()
|
||||
assert data["auths"]["123456789012.dkr.ecr.us-east-2.amazonaws.com"] == {
|
||||
"auth": expected_auth
|
||||
}
|
||||
|
||||
def test_tempdir_cleaned_up(self):
|
||||
with _docker_config_for_token("registry.example.com", "tok") as cfg:
|
||||
with _docker_config_for_token(
|
||||
"registry.example.com", "tok", "oauth2accesstoken"
|
||||
) as cfg:
|
||||
assert os.path.isdir(cfg)
|
||||
assert not os.path.exists(cfg)
|
||||
|
||||
def test_different_registries(self):
|
||||
with _docker_config_for_token("gcr.io", "token123") as cfg:
|
||||
with _docker_config_for_token("gcr.io", "token123", "oauth2accesstoken") as cfg:
|
||||
with open(os.path.join(cfg, "config.json")) as f:
|
||||
data = json.load(f)
|
||||
assert "gcr.io" in data["auths"]
|
||||
@@ -593,9 +610,9 @@ class TestDockerBuildCommand:
|
||||
False,
|
||||
DockerBuildCommand(
|
||||
("docker", "buildx", "build"),
|
||||
("--platform", "linux/amd64", "--load", "--progress=quiet"),
|
||||
("--platform", "linux/amd64", "--load", "--progress=plain"),
|
||||
),
|
||||
id="other_hosts_cross_build_quietly",
|
||||
id="other_hosts_capture_plain_build_logs",
|
||||
),
|
||||
pytest.param(
|
||||
"arm64",
|
||||
@@ -621,6 +638,7 @@ class TestSelectSource:
|
||||
"tag": None,
|
||||
"remote_build_flag": None,
|
||||
"placement": RequestedPlacement(),
|
||||
"image_pull_secrets": (),
|
||||
"selector": ByName("my-app"),
|
||||
}
|
||||
REPOSITORY = "registry.example.com/app"
|
||||
@@ -636,7 +654,7 @@ class TestSelectSource:
|
||||
reference=ImageReference(REPOSITORY, "latest"),
|
||||
prebuilt_image=None,
|
||||
),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="push_to_selects_the_external_source_with_the_default_tag",
|
||||
),
|
||||
@@ -648,7 +666,7 @@ class TestSelectSource:
|
||||
reference=ImageReference(REPOSITORY, "v2"),
|
||||
prebuilt_image=None,
|
||||
),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="push_to_keeps_a_tag_given_in_the_reference",
|
||||
),
|
||||
@@ -660,7 +678,7 @@ class TestSelectSource:
|
||||
reference=ImageReference(REPOSITORY, "v3"),
|
||||
prebuilt_image=None,
|
||||
),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="tag_flag_composes_with_push_to",
|
||||
),
|
||||
@@ -672,7 +690,7 @@ class TestSelectSource:
|
||||
reference=ImageReference(REPOSITORY, "latest"),
|
||||
prebuilt_image="app:dev",
|
||||
),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="prebuilt_image_is_retagged_for_push_to_without_docker_checks",
|
||||
),
|
||||
@@ -687,16 +705,35 @@ class TestSelectSource:
|
||||
reference=ImageReference(REPOSITORY, "latest"),
|
||||
prebuilt_image=None,
|
||||
),
|
||||
requested_placement=RequestedPlacement("listener-1", "agents"),
|
||||
create_only=CreateOnlySettings(
|
||||
RequestedPlacement("listener-1", "agents")
|
||||
),
|
||||
),
|
||||
id="push_to_carries_the_requested_placement",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"push_to": REPOSITORY,
|
||||
"image_pull_secrets": ("regcred", "other-secret"),
|
||||
},
|
||||
True,
|
||||
CustomerRegistrySource(
|
||||
image=BuildAndPush(
|
||||
reference=ImageReference(REPOSITORY, "latest"),
|
||||
prebuilt_image=None,
|
||||
),
|
||||
create_only=CreateOnlySettings(
|
||||
RequestedPlacement(), ("regcred", "other-secret")
|
||||
),
|
||||
),
|
||||
id="push_to_carries_the_image_pull_secrets",
|
||||
),
|
||||
pytest.param(
|
||||
{"image_uri": f"{REPOSITORY}@sha256:abc123"},
|
||||
True,
|
||||
CustomerRegistrySource(
|
||||
image=PublishedImage(f"{REPOSITORY}@sha256:abc123"),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="image_uri_selects_the_published_image_source_by_digest",
|
||||
),
|
||||
@@ -705,7 +742,7 @@ class TestSelectSource:
|
||||
False,
|
||||
CustomerRegistrySource(
|
||||
image=PublishedImage(f"{REPOSITORY}@sha256:abc123"),
|
||||
requested_placement=RequestedPlacement(),
|
||||
create_only=CreateOnlySettings(RequestedPlacement()),
|
||||
),
|
||||
id="image_uri_needs_no_local_docker_and_is_trimmed",
|
||||
),
|
||||
@@ -717,7 +754,9 @@ class TestSelectSource:
|
||||
False,
|
||||
CustomerRegistrySource(
|
||||
image=PublishedImage(f"{REPOSITORY}@sha256:abc123"),
|
||||
requested_placement=RequestedPlacement("listener-1", "agents"),
|
||||
create_only=CreateOnlySettings(
|
||||
RequestedPlacement("listener-1", "agents")
|
||||
),
|
||||
),
|
||||
id="image_uri_carries_the_requested_placement",
|
||||
),
|
||||
@@ -806,6 +845,11 @@ class TestSelectSource:
|
||||
"only apply when creating a deployment with --push-to",
|
||||
id="namespace_without_push_to",
|
||||
),
|
||||
pytest.param(
|
||||
{"image_pull_secrets": ("regcred",)},
|
||||
"only apply when creating a deployment with --push-to",
|
||||
id="image_pull_secret_without_push_to",
|
||||
),
|
||||
pytest.param(
|
||||
{"image_uri": REPOSITORY, "push_to": REPOSITORY},
|
||||
"--image-uri cannot be combined with --push-to.",
|
||||
@@ -1243,7 +1287,7 @@ def test_a_deployment_id_with_listener_flags_is_refused_without_probing_docker(
|
||||
|
||||
monkeypatch.setattr(deploy_mod, "can_build_locally", explode)
|
||||
|
||||
with pytest.raises(click.UsageError, match="--deployment-id"):
|
||||
with pytest.raises(click.UsageError, match="dep-1"):
|
||||
_select_source(
|
||||
push_to="registry.example.com/app",
|
||||
image=None,
|
||||
@@ -1252,10 +1296,77 @@ def test_a_deployment_id_with_listener_flags_is_refused_without_probing_docker(
|
||||
tag=None,
|
||||
remote_build_flag=None,
|
||||
placement=RequestedPlacement(listener_id="listener-1"),
|
||||
image_pull_secrets=(),
|
||||
selector=ById("dep-1"),
|
||||
)
|
||||
|
||||
|
||||
def test_a_deployment_id_with_image_pull_secrets_is_refused_without_probing_docker(
|
||||
monkeypatch,
|
||||
):
|
||||
def explode() -> tuple[bool, str | None]:
|
||||
raise AssertionError("docker must not be probed for an argv-only conflict")
|
||||
|
||||
monkeypatch.setattr(deploy_mod, "can_build_locally", explode)
|
||||
|
||||
with pytest.raises(click.UsageError, match="dep-1"):
|
||||
_select_source(
|
||||
push_to="registry.example.com/app",
|
||||
image=None,
|
||||
image_uri=None,
|
||||
image_name=None,
|
||||
tag=None,
|
||||
remote_build_flag=None,
|
||||
placement=RequestedPlacement(),
|
||||
image_pull_secrets=("regcred",),
|
||||
selector=ById("dep-1"),
|
||||
)
|
||||
|
||||
|
||||
class TestResourceSpec:
|
||||
def test_no_secrets_keeps_the_operator_default(self):
|
||||
assert _resource_spec(()) == {}
|
||||
|
||||
def test_secrets_become_local_object_references_in_order(self):
|
||||
assert _resource_spec(("regcred", "other-secret")) == {
|
||||
"image_pull_secrets": [
|
||||
{"name": "regcred"},
|
||||
{"name": "other-secret"},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class TestCreateOnlySettings:
|
||||
@pytest.mark.parametrize(
|
||||
"name",
|
||||
[
|
||||
pytest.param("Bad_Name", id="uppercase_and_underscore"),
|
||||
pytest.param("", id="empty"),
|
||||
pytest.param("-leading-hyphen", id="leading_hyphen"),
|
||||
pytest.param("a" * 254, id="too_long"),
|
||||
],
|
||||
)
|
||||
def test_an_invalid_secret_name_is_rejected_up_front(self, name):
|
||||
with pytest.raises(click.UsageError, match="not a valid Kubernetes"):
|
||||
CreateOnlySettings(RequestedPlacement(), (name,))
|
||||
|
||||
def test_a_valid_secret_name_is_accepted(self):
|
||||
CreateOnlySettings(RequestedPlacement(), ("regcred-1.example",))
|
||||
|
||||
def test_ensure_not_existing_mentions_the_api_and_ui_for_pull_secrets(self):
|
||||
settings = CreateOnlySettings(RequestedPlacement(), ("regcred",))
|
||||
with pytest.raises(click.UsageError, match="API or the LangSmith UI"):
|
||||
settings.ensure_not_existing("dep-1")
|
||||
|
||||
def test_ensure_not_existing_keeps_the_placement_message_unchanged(self):
|
||||
settings = CreateOnlySettings(RequestedPlacement(listener_id="listener-1"))
|
||||
with pytest.raises(
|
||||
click.UsageError,
|
||||
match="Listener and namespace are fixed when a deployment is created",
|
||||
):
|
||||
settings.ensure_not_existing("dep-1")
|
||||
|
||||
|
||||
class TestPlacementOnAKnownListener:
|
||||
@pytest.mark.parametrize(
|
||||
("request_", "listener", "expected"),
|
||||
|
||||
@@ -2,6 +2,7 @@ import collections.abc
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any, Generic
|
||||
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import NotRequired, Required, Self
|
||||
|
||||
from langgraph._internal._constants import OVERWRITE
|
||||
@@ -116,6 +117,9 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
|
||||
def from_checkpoint(self, checkpoint: Value) -> Self:
|
||||
empty = self.__class__(self.typ, self.operator)
|
||||
empty.key = self.key
|
||||
# A thread written while this field was a DeltaChannel stores snapshots.
|
||||
if isinstance(checkpoint, _DeltaSnapshot):
|
||||
checkpoint = checkpoint.value
|
||||
if checkpoint is not MISSING:
|
||||
empty.value = checkpoint
|
||||
return empty
|
||||
|
||||
@@ -4,7 +4,7 @@ import uuid
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from inspect import signature
|
||||
from typing import Any, Literal, cast
|
||||
from typing import Any, Literal
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -12,6 +12,7 @@ from langgraph.checkpoint.base import (
|
||||
ChannelVersions,
|
||||
Checkpoint,
|
||||
CheckpointTuple,
|
||||
DeltaChannelHistory,
|
||||
PendingWrite,
|
||||
)
|
||||
from langgraph.checkpoint.base.id import uuid6
|
||||
@@ -32,7 +33,7 @@ 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.binop import BinaryOperatorAggregate, _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
@@ -409,23 +410,39 @@ def versions_seen_without_bumps(
|
||||
def _delta_channels_to_replay(
|
||||
specs: Mapping[str, BaseChannel], checkpoint: Checkpoint
|
||||
) -> list[str]:
|
||||
"""DeltaChannels whose value at `checkpoint` the ancestor walk rebuilds.
|
||||
"""Channels whose value at `checkpoint` the ancestor walk rebuilds.
|
||||
|
||||
A `_DeltaSnapshot` blob or a plain value (migration) resolves directly via
|
||||
`from_checkpoint`, so only a channel with nothing stored here needs the
|
||||
walk. A channel with no version was never written, so it is empty without
|
||||
one; a walk for it would find no snapshot to stop at and read every
|
||||
ancestor, every time the thread is loaded.
|
||||
|
||||
A `BinaryOperatorAggregate` stores its value at every version, so one with
|
||||
nothing stored was written as a `DeltaChannel` before its annotation
|
||||
changed, and is rebuilt the same way.
|
||||
"""
|
||||
return [
|
||||
k
|
||||
for k, spec in specs.items()
|
||||
if isinstance(spec, DeltaChannel)
|
||||
if isinstance(spec, (DeltaChannel, BinaryOperatorAggregate))
|
||||
and k in checkpoint["channel_versions"]
|
||||
and checkpoint["channel_values"].get(k, MISSING) is MISSING
|
||||
]
|
||||
|
||||
|
||||
def _replay(spec: BaseChannel, history: DeltaChannelHistory) -> BaseChannel:
|
||||
ch = spec.from_checkpoint(history.get("seed", MISSING))
|
||||
if isinstance(ch, DeltaChannel):
|
||||
ch.replay_writes(history["writes"])
|
||||
else:
|
||||
# One update per write: the writes span many supersteps, and a single
|
||||
# call would reject a second Overwrite and drop the writes after one.
|
||||
for _, _, value in history["writes"]:
|
||||
ch.update([value])
|
||||
return ch
|
||||
|
||||
|
||||
def _require_saver_for_history(
|
||||
delta_channels: list[str],
|
||||
saver: BaseCheckpointSaver | None,
|
||||
@@ -465,7 +482,7 @@ def channels_from_checkpoint(
|
||||
|
||||
delta_channels = _delta_channels_to_replay(channel_specs, checkpoint)
|
||||
_require_saver_for_history(delta_channels, saver, config)
|
||||
histories: Mapping[str, Any] = {}
|
||||
histories: Mapping[str, DeltaChannelHistory] = {}
|
||||
if delta_channels and saver is not None and config is not None:
|
||||
histories = saver.get_delta_channel_history(
|
||||
config=config, channels=delta_channels
|
||||
@@ -475,11 +492,7 @@ def channels_from_checkpoint(
|
||||
for k, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
if k in histories:
|
||||
delta_spec = cast(DeltaChannel, spec)
|
||||
history = histories[k]
|
||||
replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING))
|
||||
replay_ch.replay_writes(history["writes"])
|
||||
ch = replay_ch
|
||||
ch = _replay(spec, histories[k])
|
||||
else:
|
||||
ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
channels[k] = ch
|
||||
@@ -504,7 +517,7 @@ async def achannels_from_checkpoint(
|
||||
|
||||
delta_channels = _delta_channels_to_replay(channel_specs, checkpoint)
|
||||
_require_saver_for_history(delta_channels, saver, config)
|
||||
histories: Mapping[str, Any] = {}
|
||||
histories: Mapping[str, DeltaChannelHistory] = {}
|
||||
if delta_channels and saver is not None and config is not None:
|
||||
histories = await saver.aget_delta_channel_history(
|
||||
config=config, channels=delta_channels
|
||||
@@ -514,11 +527,7 @@ async def achannels_from_checkpoint(
|
||||
for k, spec in channel_specs.items():
|
||||
ch: BaseChannel
|
||||
if k in histories:
|
||||
delta_spec = cast(DeltaChannel, spec)
|
||||
history = histories[k]
|
||||
replay_ch = delta_spec.from_checkpoint(history.get("seed", MISSING))
|
||||
replay_ch.replay_writes(history["writes"])
|
||||
ch = replay_ch
|
||||
ch = _replay(spec, histories[k])
|
||||
else:
|
||||
ch = spec.from_checkpoint(checkpoint["channel_values"].get(k, MISSING))
|
||||
channels[k] = ch
|
||||
|
||||
@@ -257,8 +257,10 @@ class PregelLoop:
|
||||
_resume_loaded_writes: Sequence[PendingWrite] = ()
|
||||
# Tasks `_reapply_writes_to_succeeded_nodes` handed loaded writes back to.
|
||||
_reapplied_task_ids: set[str]
|
||||
# An update or fork checkpoint addressed by `checkpoint_id` that the thread
|
||||
# has moved past, so it may already have children.
|
||||
# Tasks that wrote during this run, the only ones exit mode saves at exit.
|
||||
_written_task_ids: set[str]
|
||||
# A checkpoint addressed by `checkpoint_id` that the thread has moved past,
|
||||
# so it may already have children.
|
||||
_addressed_checkpoint_superseded: bool = False
|
||||
|
||||
# The checkpoint_config that points at the parent loaded at `__enter__`
|
||||
@@ -349,7 +351,6 @@ class PregelLoop:
|
||||
self.interrupt_before = interrupt_before
|
||||
self.manager = manager
|
||||
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
|
||||
self.is_replaying = CONFIG_KEY_CHECKPOINT_ID in config[CONF]
|
||||
self._migrate_checkpoint = migrate_checkpoint
|
||||
self.trigger_to_nodes = trigger_to_nodes
|
||||
self.retry_policy = retry_policy
|
||||
@@ -407,6 +408,14 @@ class PregelLoop:
|
||||
else ()
|
||||
)
|
||||
self.prev_checkpoint_config = None
|
||||
# Whether the run points this loop at a checkpoint: a truthy
|
||||
# `checkpoint_id`, this graph's `checkpoint_map` entry, or a replaying
|
||||
# parent. Every task config carries `checkpoint_id: None`, so a key
|
||||
# check would count every subgraph.
|
||||
self._addressed = bool(
|
||||
self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
|
||||
or self.config[CONF].get(CONFIG_KEY_REPLAY_STATE)
|
||||
)
|
||||
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
||||
self.control = runtime.control if isinstance(runtime, Runtime) else None
|
||||
|
||||
@@ -455,6 +464,7 @@ class PregelLoop:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
if not writes:
|
||||
return
|
||||
self._written_task_ids.add(task_id)
|
||||
# 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())
|
||||
@@ -468,6 +478,11 @@ class PregelLoop:
|
||||
for w in self.checkpoint_pending_writes
|
||||
if w[0] != task_id or w[1] not in WRITES_IDX_MAP
|
||||
]
|
||||
# Savers key a regular write by its position in the list, and the
|
||||
# special ones dropped above won't be there next time, so they go
|
||||
# last: a resume saved ahead of an update would move it to a new
|
||||
# slot once the list is saved again, and it would apply twice.
|
||||
writes = sorted(writes, key=lambda w: w[0] in WRITES_IDX_MAP)
|
||||
writes_to_save: WritesT = [
|
||||
w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
|
||||
] + list(writes)
|
||||
@@ -556,10 +571,12 @@ class PregelLoop:
|
||||
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
|
||||
},
|
||||
)
|
||||
# group by task id
|
||||
# Writes loaded with the checkpoint are stored on it already. Saving them
|
||||
# again can file them under new indexes, and then they apply twice.
|
||||
by_task = defaultdict(list)
|
||||
for task_id, channel, value in self.checkpoint_pending_writes:
|
||||
by_task[task_id].append((channel, value))
|
||||
if task_id in self._written_task_ids:
|
||||
by_task[task_id].append((channel, value))
|
||||
# submit writes to checkpointer
|
||||
for task_id, writes in by_task.items():
|
||||
if self.checkpointer_put_writes_accepts_task_path and hasattr(
|
||||
@@ -904,6 +921,7 @@ class PregelLoop:
|
||||
if tid != NULL_TASK_ID and isinstance(self.specs.get(ch), DeltaChannel)
|
||||
}
|
||||
self._exit_command_writes = []
|
||||
self._written_task_ids = set()
|
||||
# Resuming from a previous checkpoint requires two things:
|
||||
# 1. A prior checkpoint exists (channel_versions is non-empty)
|
||||
# 2. The input signals continuation (not a fresh run with new input)
|
||||
@@ -931,7 +949,7 @@ class PregelLoop:
|
||||
# writes so that interrupt() calls re-fire instead of returning
|
||||
# stale values. But if we're actively resuming, keep them —
|
||||
# multi-interrupt scenarios need previously resolved values preserved.
|
||||
is_time_traveling = self.is_replaying and (
|
||||
is_time_traveling = self._addressed and (
|
||||
# Time-travel to a subgraph checkpoint: the parent sets
|
||||
# RESUMING=True (it can't distinguish time-travel from resume),
|
||||
# so we check if this subgraph's own ns is in checkpoint_map.
|
||||
@@ -950,6 +968,10 @@ class PregelLoop:
|
||||
or configurable.get(CONFIG_KEY_RESUMING, False)
|
||||
)
|
||||
)
|
||||
# Finished tasks rerun only on time travel, or on a resume at a
|
||||
# checkpoint the thread has moved past. A resume addressed at the
|
||||
# thread's latest checkpoint hands their writes back, as without an id.
|
||||
self.is_replaying = is_time_traveling or self._addressed_checkpoint_superseded
|
||||
if is_time_traveling:
|
||||
self.checkpoint_pending_writes = [
|
||||
w for w in self.checkpoint_pending_writes if w[1] != RESUME
|
||||
@@ -1165,11 +1187,11 @@ class PregelLoop:
|
||||
# checkpoint ID since the fork was created after the subgraph's
|
||||
# checkpoints from the original execution.
|
||||
#
|
||||
# Only gate on is_time_traveling (not is_replaying). When the
|
||||
# Only gate on is_time_traveling (not _addressed). When the
|
||||
# client resumes with an explicit checkpoint_id that happens to
|
||||
# point at the current head (e.g. LangGraph Studio sending
|
||||
# `checkpoint: {checkpoint_id}` alongside Command(resume=...)),
|
||||
# is_replaying is True but is_time_traveling is False. In that
|
||||
# the run is addressed but not time traveling. In that
|
||||
# case subgraphs should load their latest checkpoint normally,
|
||||
# not go through ReplayState's before-bound lookup which would
|
||||
# miss subgraph checkpoints created during processing of the
|
||||
@@ -1789,7 +1811,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# This covers both normal replay and subgraphs resolved via
|
||||
# checkpoint_map during time-travel.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
if saved is not None and saved.metadata.get("source") in ("update", "fork"):
|
||||
if saved is not None:
|
||||
self._addressed_checkpoint_superseded = checkpoint_superseded(
|
||||
self.checkpointer, self.checkpoint_config, saved
|
||||
)
|
||||
@@ -2045,7 +2067,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# This covers both normal replay and subgraphs resolved via
|
||||
# checkpoint_map during time-travel.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
if saved is not None and saved.metadata.get("source") in ("update", "fork"):
|
||||
if saved is not None:
|
||||
self._addressed_checkpoint_superseded = await acheckpoint_superseded(
|
||||
self.checkpointer, self.checkpoint_config, saved
|
||||
)
|
||||
|
||||
@@ -521,10 +521,11 @@ def test_addressed_resume_keeps_a_rerun_tasks_write_to_a_new_channel(
|
||||
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke(_both("in"), config, durability=durability)
|
||||
head = graph.get_state(config).config
|
||||
paused = graph.get_state(config).config
|
||||
graph.update_state(config, None)
|
||||
|
||||
live = graph.invoke(
|
||||
Command(resume="yes", update={"flag": True}), head, durability=durability
|
||||
Command(resume="yes", update={"flag": True}), paused, durability=durability
|
||||
)
|
||||
|
||||
state = graph.get_state(config)
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
import operator
|
||||
from typing import Annotated, Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
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.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel._checkpoint import channels_from_checkpoint, empty_checkpoint
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
class DeltaState(TypedDict):
|
||||
messages: Annotated[
|
||||
list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3)
|
||||
]
|
||||
|
||||
|
||||
class PlainState(TypedDict):
|
||||
messages: Annotated[list, add_messages]
|
||||
|
||||
|
||||
def _app(schema: type, checkpointer: BaseCheckpointSaver) -> Any:
|
||||
builder = StateGraph(schema)
|
||||
builder.add_node(
|
||||
"chat",
|
||||
lambda state: {"messages": [AIMessage(f"reply {len(state['messages'])}")]},
|
||||
)
|
||||
builder.add_edge(START, "chat")
|
||||
builder.add_edge("chat", END)
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _delta_thread(checkpointer: BaseCheckpointSaver) -> RunnableConfig:
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
delta = _app(DeltaState, checkpointer)
|
||||
for i in range(4):
|
||||
delta.invoke({"messages": [HumanMessage(f"q{i}")]}, config)
|
||||
assert any(
|
||||
isinstance(t.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
|
||||
for t in checkpointer.list(config)
|
||||
), "the thread should have a snapshot to read through"
|
||||
return config
|
||||
|
||||
|
||||
def test_plain_reducer_reads_every_delta_written_checkpoint(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
config = _delta_thread(sync_checkpointer)
|
||||
delta = _app(DeltaState, sync_checkpointer)
|
||||
plain = _app(PlainState, sync_checkpointer)
|
||||
|
||||
assert [s.values for s in plain.get_state_history(config)] == [
|
||||
s.values for s in delta.get_state_history(config)
|
||||
]
|
||||
|
||||
|
||||
def test_plain_reducer_continues_delta_written_thread(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
config = _delta_thread(sync_checkpointer)
|
||||
delta = _app(DeltaState, sync_checkpointer)
|
||||
before = delta.get_state(config).values["messages"]
|
||||
|
||||
after = _app(PlainState, sync_checkpointer).invoke(
|
||||
{"messages": [HumanMessage("q4")]}, config
|
||||
)["messages"]
|
||||
|
||||
assert after[:-2] == before
|
||||
assert [m.content for m in after[-2:]] == ["q4", f"reply {len(before) + 1}"]
|
||||
head = sync_checkpointer.get_tuple(config)
|
||||
assert head is not None
|
||||
assert head.checkpoint["channel_values"]["messages"] == after
|
||||
assert delta.get_state(config).values["messages"] == after
|
||||
|
||||
|
||||
async def test_plain_reducer_reads_and_continues_delta_written_thread_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
delta = _app(DeltaState, async_checkpointer)
|
||||
for i in range(4):
|
||||
await delta.ainvoke({"messages": [HumanMessage(f"q{i}")]}, config)
|
||||
plain = _app(PlainState, async_checkpointer)
|
||||
|
||||
expected = [s.values async for s in delta.aget_state_history(config)]
|
||||
assert [s.values async for s in plain.aget_state_history(config)] == expected
|
||||
|
||||
after = (await plain.ainvoke({"messages": [HumanMessage("q4")]}, config))[
|
||||
"messages"
|
||||
]
|
||||
assert after[:-2] == expected[0]["messages"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("seed", "writes", "expected"),
|
||||
[
|
||||
(_DeltaSnapshot(["a"]), [["b"], ["c"]], ["a", "b", "c"]),
|
||||
(["a"], [["b"], Overwrite(["c"]), ["d"]], ["c", "d"]),
|
||||
(None, [["b"], Overwrite(["c"]), Overwrite(["d"]), ["e"]], ["d", "e"]),
|
||||
],
|
||||
)
|
||||
def test_plain_reducer_replays_delta_history_one_write_at_a_time(
|
||||
seed: Any, writes: list[Any], expected: list[str]
|
||||
) -> None:
|
||||
history: dict[str, Any] = {"writes": [("t", "items", w) for w in writes]}
|
||||
if seed is not None:
|
||||
history["seed"] = seed
|
||||
checkpoint = empty_checkpoint()
|
||||
checkpoint["channel_versions"]["items"] = 1
|
||||
saver = InMemorySaver()
|
||||
|
||||
with patch.object(
|
||||
saver, "get_delta_channel_history", return_value={"items": history}
|
||||
):
|
||||
channels, _ = channels_from_checkpoint(
|
||||
{"items": BinaryOperatorAggregate(list, operator.add)},
|
||||
checkpoint,
|
||||
saver=saver,
|
||||
config={"configurable": {"thread_id": "1"}},
|
||||
)
|
||||
|
||||
assert channels["items"].get() == expected
|
||||
|
||||
|
||||
def test_plain_thread_reads_never_walk_history() -> None:
|
||||
saver = InMemorySaver()
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
plain = _app(PlainState, saver)
|
||||
|
||||
with patch.object(saver, "get_delta_channel_history") as walk:
|
||||
plain.invoke({"messages": [HumanMessage("q0")]}, config)
|
||||
plain.update_state(config, {"messages": [HumanMessage("note")]})
|
||||
plain.invoke({"messages": [HumanMessage("q1")]}, config)
|
||||
plain.get_state(config)
|
||||
list(plain.get_state_history(config))
|
||||
|
||||
walk.assert_not_called()
|
||||
@@ -152,6 +152,83 @@ def test_finished_task_does_not_report_answered_interrupt(
|
||||
assert asked == [["A1", "B"]]
|
||||
|
||||
|
||||
def test_finished_task_output_applies_once(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("a", lambda s: {"log": [f"a:{interrupt('A')}"]})
|
||||
builder.add_node("b", lambda s: {"log": [f"b:{interrupt('B')}"]})
|
||||
builder.add_node("c", lambda s: {"log": [f"c:{interrupt('C')}"]})
|
||||
builder.add_edge(START, "a")
|
||||
builder.add_edge(START, "b")
|
||||
builder.add_edge(START, "c")
|
||||
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||
config = _config()
|
||||
|
||||
result = graph.invoke({"log": []}, config, durability=durability)
|
||||
while pending := result.get("__interrupt__"):
|
||||
result = graph.invoke(
|
||||
Command(resume={pending[0].id: "yes"}), config, durability=durability
|
||||
)
|
||||
|
||||
assert sorted(result["log"]) == ["a:yes", "b:yes", "c:yes"]
|
||||
|
||||
|
||||
def test_resume_with_update_applies_the_update_once(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
def ask_three_times(state: State) -> State:
|
||||
answers = [interrupt(str(question)) for question in range(1, 4)]
|
||||
return {"log": ["a:" + "".join(answers)]}
|
||||
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("a", ask_three_times)
|
||||
.add_edge(START, "a")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
config = _config()
|
||||
|
||||
result = graph.invoke({"log": []}, config, durability=durability)
|
||||
for i in range(1, 4):
|
||||
result = graph.invoke(
|
||||
Command(resume=str(i), update={"log": [f"u{i}"]}),
|
||||
config,
|
||||
durability=durability,
|
||||
)
|
||||
|
||||
assert sorted(result["log"]) == ["a:123", "u1", "u2", "u3"]
|
||||
assert sorted(graph.get_state(config).values["log"]) == sorted(result["log"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpoint_id", ["none", "head"])
|
||||
def test_resume_at_the_latest_checkpoint_does_not_rerun_finished_task(
|
||||
sync_checkpointer: BaseCheckpointSaver, checkpoint_id: str
|
||||
) -> None:
|
||||
calls: Counter[str] = Counter()
|
||||
graph = _build_parallel(sync_checkpointer, calls)
|
||||
config = _config()
|
||||
graph.invoke({"log": []}, config)
|
||||
|
||||
def resume(question: str, answer: str) -> Any:
|
||||
snapshot = graph.get_state(config)
|
||||
target = (
|
||||
snapshot.config
|
||||
if checkpoint_id == "head"
|
||||
else {"configurable": {**config["configurable"], "checkpoint_id": None}}
|
||||
)
|
||||
return graph.invoke(
|
||||
Command(resume={_interrupt_by_value(snapshot, question).id: answer}),
|
||||
target,
|
||||
)
|
||||
|
||||
resume("A1", "yes")
|
||||
result = resume("B", "ok")
|
||||
|
||||
assert sorted(result["log"]) == ["a:yes", "b:ok"]
|
||||
assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again"
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_finished_task_does_not_report_answered_interrupt_async(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
@@ -474,6 +551,64 @@ def test_functional_task_finished_with_none_is_not_rerun(
|
||||
assert calls == {"a": 2, "b": 3}
|
||||
|
||||
|
||||
def test_subgraph_finished_task_is_not_rerun(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
calls: Counter[str] = Counter()
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("child", _build_parallel(None, calls))
|
||||
builder.add_edge(START, "child")
|
||||
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||
config = _config()
|
||||
|
||||
graph.invoke({"log": []}, config, durability=durability)
|
||||
snapshot = graph.get_state(config)
|
||||
graph.invoke(
|
||||
Command(resume={_interrupt_by_value(snapshot, "A1").id: "yes"}),
|
||||
config,
|
||||
durability=durability,
|
||||
)
|
||||
snapshot = graph.get_state(config)
|
||||
result = graph.invoke(
|
||||
Command(resume={_interrupt_by_value(snapshot, "B").id: "ok"}),
|
||||
config,
|
||||
durability=durability,
|
||||
)
|
||||
|
||||
assert sorted(result["log"]) == ["a:yes", "b:ok"]
|
||||
assert calls == {"a": 2, "b": 3}, "resuming B ran the finished A again"
|
||||
|
||||
|
||||
def test_nested_entrypoint_finished_task_is_not_rerun(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
calls: Counter[str] = Counter()
|
||||
|
||||
@task
|
||||
def step_1(query: str) -> str:
|
||||
calls["step_1"] += 1
|
||||
return f"{query} bar"
|
||||
|
||||
@task
|
||||
def human_feedback(query: str) -> str:
|
||||
calls["human_feedback"] += 1
|
||||
return f"{query} {interrupt('feedback?')}"
|
||||
|
||||
@entrypoint()
|
||||
def child(query: str) -> str:
|
||||
return human_feedback(step_1(query).result()).result()
|
||||
|
||||
@entrypoint(checkpointer=sync_checkpointer)
|
||||
def parent(query: str) -> str:
|
||||
return child.invoke(query)
|
||||
|
||||
config = _config()
|
||||
parent.invoke("foo", config)
|
||||
|
||||
assert parent.invoke(Command(resume="baz"), config) == "foo bar baz"
|
||||
assert calls == {"step_1": 1, "human_feedback": 2}
|
||||
|
||||
|
||||
# --- Reading task status from recorded writes ---
|
||||
|
||||
|
||||
|
||||
@@ -7737,6 +7737,8 @@ def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
# --- CHILD GRAPH ---
|
||||
|
||||
answered: Counter[str] = Counter()
|
||||
|
||||
class ChildState(BaseModel):
|
||||
prompt: str = Field(..., description="What is going to be asked to the user?")
|
||||
human_input: str | None = Field(None, description="What the human said")
|
||||
@@ -7746,6 +7748,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
|
||||
|
||||
def get_human_input(state: ChildState):
|
||||
human_input = interrupt(state.prompt)
|
||||
answered["human"] += 1
|
||||
|
||||
return dict(
|
||||
human_inputs=[human_input], # update parent state
|
||||
@@ -7753,6 +7756,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
|
||||
|
||||
def get_dolphin_input(state: ChildState):
|
||||
human_input = interrupt(state.prompt)
|
||||
answered["dolphin"] += 1
|
||||
|
||||
return dict(
|
||||
human_inputs=[human_input], # update parent state
|
||||
@@ -7846,6 +7850,7 @@ def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> N
|
||||
|
||||
assert invokes == 5
|
||||
assert len(events) == 5
|
||||
assert answered == {"human": 2, "dolphin": 2}
|
||||
|
||||
|
||||
def test_pregel_loop_refcount():
|
||||
|
||||
@@ -0,0 +1,467 @@
|
||||
import operator
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import PRUNE_DELTA_THREADS_SQL, PRUNE_HEADS_SQL
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer, add_messages
|
||||
from langgraph.types import Command, StateSnapshot, interrupt
|
||||
from tests.conftest import NO_DOCKER
|
||||
from tests.conftest_checkpointer import (
|
||||
DEFAULT_POSTGRES_URI,
|
||||
_checkpointer_postgres,
|
||||
_checkpointer_postgres_aio,
|
||||
_checkpointer_postgres_aio_pipe,
|
||||
_checkpointer_postgres_aio_pool,
|
||||
_checkpointer_postgres_pipe,
|
||||
_checkpointer_postgres_pool,
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.anyio,
|
||||
pytest.mark.skipif(NO_DOCKER, reason="needs Postgres"),
|
||||
]
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
SYNC = [
|
||||
_checkpointer_postgres,
|
||||
_checkpointer_postgres_pipe,
|
||||
_checkpointer_postgres_pool,
|
||||
]
|
||||
ASYNC = [
|
||||
_checkpointer_postgres_aio,
|
||||
_checkpointer_postgres_aio_pipe,
|
||||
_checkpointer_postgres_aio_pool,
|
||||
]
|
||||
|
||||
|
||||
class State(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
notes: Annotated[list[str], operator.add]
|
||||
|
||||
|
||||
class DeltaState(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
|
||||
def _reply(ask: bool):
|
||||
def reply(state: dict) -> dict:
|
||||
n = len(state["messages"])
|
||||
answer = interrupt("approve?") if ask else "ok"
|
||||
return {"messages": [AIMessage(f"{answer} {n}", id=f"reply-{n}")]}
|
||||
|
||||
return reply
|
||||
|
||||
|
||||
def _graph(
|
||||
checkpointer: Any,
|
||||
*,
|
||||
ask: bool = False,
|
||||
fail: set[str] | None = None,
|
||||
calls: dict[str, int] | None = None,
|
||||
) -> Any:
|
||||
def note(name: str):
|
||||
def node(state: State) -> dict:
|
||||
if calls is not None:
|
||||
calls[name] = calls.get(name, 0) + 1
|
||||
if fail is not None and name in fail:
|
||||
fail.discard(name)
|
||||
raise ValueError(f"{name} failed")
|
||||
return {"notes": [name]}
|
||||
|
||||
return node
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("left", note("left"))
|
||||
builder.add_node("right", note("right"))
|
||||
builder.add_node("reply", _reply(ask))
|
||||
builder.add_edge(START, "left")
|
||||
builder.add_edge(START, "right")
|
||||
builder.add_edge(["left", "right"], "reply")
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _counting_subgraph(name: str, *, checkpointer: bool | None) -> Any:
|
||||
class Calls(TypedDict):
|
||||
messages: Annotated[list[AnyMessage], add_messages]
|
||||
calls: Annotated[list[str], operator.add]
|
||||
|
||||
def count(state: Calls) -> dict:
|
||||
n = len(state["messages"])
|
||||
reply = AIMessage(f"{name} call {len(state['calls']) + 1}", id=f"{name}-{n}")
|
||||
return {"calls": [name], "messages": [reply]}
|
||||
|
||||
builder = StateGraph(Calls)
|
||||
builder.add_node("count", count)
|
||||
builder.add_edge(START, "count")
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _delta_graph(checkpointer: Any) -> Any:
|
||||
builder = StateGraph(DeltaState)
|
||||
builder.add_node("reply", _reply(ask=False))
|
||||
builder.add_edge(START, "reply")
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _config(thread_id: str) -> RunnableConfig:
|
||||
return {"configurable": {"thread_id": thread_id}}
|
||||
|
||||
|
||||
def _say(i: int) -> dict:
|
||||
return {"messages": [HumanMessage(f"hi {i}", id=f"hi-{i}")]}
|
||||
|
||||
|
||||
def _checkpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
|
||||
return [c.checkpoint["id"] for c in checkpointer.list(_config(thread_id))]
|
||||
|
||||
|
||||
def _namespaces(checkpointer: Any, thread_id: str) -> list[str]:
|
||||
return sorted(
|
||||
c.config["configurable"]["checkpoint_ns"]
|
||||
for c in checkpointer.list(_config(thread_id))
|
||||
)
|
||||
|
||||
|
||||
async def _acheckpoint_ids(checkpointer: Any, thread_id: str) -> list[str]:
|
||||
return [c.checkpoint["id"] async for c in checkpointer.alist(_config(thread_id))]
|
||||
|
||||
|
||||
def _parentless(snapshot: StateSnapshot) -> StateSnapshot:
|
||||
return snapshot._replace(
|
||||
parent_config=None,
|
||||
tasks=tuple(
|
||||
t._replace(state=_parentless(t.state))
|
||||
if isinstance(t.state, StateSnapshot)
|
||||
else t
|
||||
for t in snapshot.tasks
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _run_once_after(
|
||||
saver: PostgresSaver, sql: str, action: Callable[[], Any], *, nth: int = 1
|
||||
) -> None:
|
||||
seen = 0
|
||||
cursor = saver._cursor
|
||||
|
||||
class Hooked:
|
||||
def __init__(self, cur: Any) -> None:
|
||||
self.cur = cur
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self.cur, name)
|
||||
|
||||
def execute(self, query: Any, params: Any = None, **kwargs: Any) -> None:
|
||||
nonlocal seen
|
||||
self.cur.execute(query, params, **kwargs)
|
||||
if query is sql:
|
||||
seen += 1
|
||||
if seen == nth:
|
||||
action()
|
||||
|
||||
@contextmanager
|
||||
def hooked(**kwargs: Any):
|
||||
with cursor(**kwargs) as cur:
|
||||
yield Hooked(cur)
|
||||
|
||||
saver._cursor = hooked
|
||||
|
||||
|
||||
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_keeps_the_state_and_the_thread_runs_on_like_its_twin(
|
||||
make_checkpointer: Any, durability: str
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
graph = _graph(checkpointer)
|
||||
for i in range(3):
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(_say(i), _config(thread_id), durability=durability)
|
||||
before = graph.get_state(_config("pruned"))
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert _checkpoint_ids(checkpointer, "pruned") == [
|
||||
before.config["configurable"]["checkpoint_id"]
|
||||
]
|
||||
assert graph.get_state(_config("pruned")) == _parentless(before)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(_say(3), _config(thread_id), durability=durability)
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_keeps_a_pending_interrupt_resumable(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
graph = _graph(checkpointer, ask=True)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
for i in range(2):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
|
||||
graph.invoke(_say(2), _config(thread_id))
|
||||
before = graph.get_state(_config("pruned"))
|
||||
assert before.interrupts
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert graph.get_state(_config("pruned")) == _parentless(before)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(Command(resume="yes 2"), _config(thread_id))
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_keeps_writes_of_tasks_that_finished_before_a_failure(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
calls: dict[str, int] = {}
|
||||
fail: set[str] = set()
|
||||
graph = _graph(checkpointer, fail=fail, calls=calls)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
for i in range(2):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
fail.add("right")
|
||||
with pytest.raises(ValueError, match="right failed"):
|
||||
graph.invoke(_say(2), _config(thread_id))
|
||||
before = graph.get_state(_config("pruned"))
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert graph.get_state(_config("pruned")) == _parentless(before)
|
||||
left_runs = calls["left"]
|
||||
graph.invoke(None, _config("pruned"))
|
||||
assert calls["left"] == left_runs, "a task that had finished ran again"
|
||||
graph.invoke(None, _config("twin"))
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_keeps_an_interrupted_subgraph_resumable(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
sub = StateGraph(State)
|
||||
sub.add_node("reply", _reply(ask=True))
|
||||
sub.add_edge(START, "reply")
|
||||
parent = StateGraph(State)
|
||||
parent.add_node("sub", sub.compile())
|
||||
parent.add_edge(START, "sub")
|
||||
graph = parent.compile(checkpointer=checkpointer)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
for i in range(2):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
graph.invoke(Command(resume=f"yes {i}"), _config(thread_id))
|
||||
graph.invoke(_say(2), _config(thread_id))
|
||||
before = graph.get_state(_config("pruned"), subgraphs=True)
|
||||
interrupted = before.tasks[0].state.config["configurable"]["checkpoint_ns"]
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert _namespaces(checkpointer, "pruned") == ["", interrupted]
|
||||
assert graph.get_state(_config("pruned"), subgraphs=True) == _parentless(before)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(Command(resume="yes 2"), _config(thread_id))
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("durability", ["sync", "async", "exit"])
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_keeps_a_checkpointer_true_subgraph_but_not_finished_calls(
|
||||
make_checkpointer: Any, durability: str
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
parent = StateGraph(State)
|
||||
parent.add_node("memory", _counting_subgraph("memory", checkpointer=True))
|
||||
parent.add_node("call", _counting_subgraph("call", checkpointer=None))
|
||||
parent.add_edge(START, "memory")
|
||||
parent.add_edge("memory", "call")
|
||||
graph = parent.compile(checkpointer=checkpointer)
|
||||
for i in range(3):
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(_say(i), _config(thread_id), durability=durability)
|
||||
|
||||
checkpointer.prune(["pruned"])
|
||||
|
||||
assert _namespaces(checkpointer, "pruned") == ["", "memory"]
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(_say(3), _config(thread_id), durability=durability)
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
def test_keep_latest_keeps_a_turn_saved_after_it_read_the_heads() -> None:
|
||||
with (
|
||||
_checkpointer_postgres() as checkpointer,
|
||||
PostgresSaver.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
|
||||
) as pruner,
|
||||
):
|
||||
graph = _graph(checkpointer)
|
||||
for i in range(3):
|
||||
for thread_id in ("pruned", "twin"):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
_run_once_after(
|
||||
pruner, PRUNE_HEADS_SQL, lambda: graph.invoke(_say(3), _config("pruned"))
|
||||
)
|
||||
|
||||
pruner.prune(["pruned"])
|
||||
|
||||
graph.invoke(_say(3), _config("twin"))
|
||||
assert (
|
||||
graph.get_state(_config("pruned")).values
|
||||
== graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
def test_keep_latest_refuses_a_thread_that_switches_to_a_delta_channel_mid_batch() -> (
|
||||
None
|
||||
):
|
||||
with (
|
||||
_checkpointer_postgres() as checkpointer,
|
||||
PostgresSaver.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
|
||||
) as pruner,
|
||||
):
|
||||
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
|
||||
for i in range(3):
|
||||
for thread_id in ("switched", "twin"):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
_run_once_after(
|
||||
pruner,
|
||||
PRUNE_DELTA_THREADS_SQL,
|
||||
lambda: delta_graph.invoke(_say(3), _config("switched")),
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="switched"):
|
||||
pruner.prune(["switched"])
|
||||
|
||||
delta_graph.invoke(_say(3), _config("twin"))
|
||||
assert (
|
||||
delta_graph.get_state(_config("switched")).values
|
||||
== delta_graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
def test_keep_latest_keeps_a_delta_channel_turn_saved_after_the_thread_check() -> None:
|
||||
with (
|
||||
_checkpointer_postgres() as checkpointer,
|
||||
PostgresSaver.from_conn_string(
|
||||
DEFAULT_POSTGRES_URI + checkpointer.conn.info.dbname
|
||||
) as pruner,
|
||||
):
|
||||
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
|
||||
for i in range(3):
|
||||
for thread_id in ("switched", "twin"):
|
||||
graph.invoke(_say(i), _config(thread_id))
|
||||
head = _checkpoint_ids(checkpointer, "switched")[0]
|
||||
_run_once_after(
|
||||
pruner,
|
||||
PRUNE_DELTA_THREADS_SQL,
|
||||
lambda: delta_graph.invoke(_say(3), _config("switched")),
|
||||
nth=2,
|
||||
)
|
||||
|
||||
pruner.prune(["switched"])
|
||||
|
||||
assert _checkpoint_ids(checkpointer, "switched")[-1] == head
|
||||
delta_graph.invoke(_say(3), _config("twin"))
|
||||
assert (
|
||||
delta_graph.get_state(_config("switched")).values
|
||||
== delta_graph.get_state(_config("twin")).values
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_checkpointer", SYNC)
|
||||
def test_keep_latest_refuses_a_batch_with_a_delta_channel_thread(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
with make_checkpointer() as checkpointer:
|
||||
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
|
||||
for i in range(3):
|
||||
graph.invoke(_say(i), _config("ordinary"))
|
||||
delta_graph.invoke(_say(i), _config("delta"))
|
||||
before = {t: _checkpoint_ids(checkpointer, t) for t in ("ordinary", "delta")}
|
||||
delta_before = delta_graph.get_state(_config("delta"))
|
||||
|
||||
with pytest.raises(NotImplementedError, match="delta"):
|
||||
checkpointer.prune(["ordinary", "delta"])
|
||||
|
||||
assert {t: _checkpoint_ids(checkpointer, t) for t in before} == before
|
||||
assert delta_graph.get_state(_config("delta")) == delta_before
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@pytest.mark.parametrize("make_checkpointer", ASYNC)
|
||||
async def test_akeep_latest_keeps_a_pending_interrupt_resumable(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
async with make_checkpointer() as checkpointer:
|
||||
graph = _graph(checkpointer, ask=True)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
for i in range(2):
|
||||
await graph.ainvoke(_say(i), _config(thread_id))
|
||||
await graph.ainvoke(Command(resume=f"yes {i}"), _config(thread_id))
|
||||
await graph.ainvoke(_say(2), _config(thread_id))
|
||||
before = await graph.aget_state(_config("pruned"))
|
||||
|
||||
await checkpointer.aprune(["pruned"])
|
||||
|
||||
assert await _acheckpoint_ids(checkpointer, "pruned") == [
|
||||
before.config["configurable"]["checkpoint_id"]
|
||||
]
|
||||
assert await graph.aget_state(_config("pruned")) == _parentless(before)
|
||||
for thread_id in ("pruned", "twin"):
|
||||
await graph.ainvoke(Command(resume="yes 2"), _config(thread_id))
|
||||
assert (await graph.aget_state(_config("pruned"))).values == (
|
||||
await graph.aget_state(_config("twin"))
|
||||
).values
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_checkpointer", ASYNC)
|
||||
async def test_akeep_latest_refuses_a_batch_with_a_delta_channel_thread(
|
||||
make_checkpointer: Any,
|
||||
) -> None:
|
||||
async with make_checkpointer() as checkpointer:
|
||||
graph, delta_graph = _graph(checkpointer), _delta_graph(checkpointer)
|
||||
for i in range(3):
|
||||
await graph.ainvoke(_say(i), _config("ordinary"))
|
||||
await delta_graph.ainvoke(_say(i), _config("delta"))
|
||||
before = {
|
||||
t: await _acheckpoint_ids(checkpointer, t) for t in ("ordinary", "delta")
|
||||
}
|
||||
|
||||
with pytest.raises(NotImplementedError, match="delta"):
|
||||
await checkpointer.aprune(["ordinary", "delta"])
|
||||
|
||||
assert {t: await _acheckpoint_ids(checkpointer, t) for t in before} == before
|
||||
@@ -1177,6 +1177,43 @@ def test_subgraph_interrupt_resume_with_explicit_head_checkpoint_id(
|
||||
assert result["value"] == ["sub_a", "human:answer", "sub_b"]
|
||||
|
||||
|
||||
def test_subgraph_replay_reasks_answered_parallel_interrupt(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""Replaying a subgraph checkpoint asks its step's questions again, including
|
||||
one already answered. An update checkpoint replays in place instead of
|
||||
forking, so its tasks keep their ids and its saved answers could be reused."""
|
||||
|
||||
subgraph = (
|
||||
StateGraph(State)
|
||||
.add_node("a", lambda s: {"value": [f"a:{interrupt('a')}"]})
|
||||
.add_node("b", lambda s: {"value": [f"b:{interrupt('b')}"]})
|
||||
.add_edge(START, "a")
|
||||
.add_edge(START, "b")
|
||||
.compile(checkpointer=True)
|
||||
)
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("subgraph_node", subgraph)
|
||||
.add_edge(START, "subgraph_node")
|
||||
.compile(checkpointer=sync_checkpointer)
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
def subgraph_head() -> RunnableConfig:
|
||||
return graph.get_state(config, subgraphs=True).tasks[0].state.config
|
||||
|
||||
graph.invoke({"value": []}, config)
|
||||
graph.update_state(subgraph_head(), {"value": ["u"]}, as_node="__start__")
|
||||
result = graph.invoke(None, config)
|
||||
question_a = next(i for i in result["__interrupt__"] if i.value == "a")
|
||||
graph.invoke(Command(resume={question_a.id: "yes"}), config)
|
||||
|
||||
result = graph.invoke(None, subgraph_head())
|
||||
|
||||
assert sorted(i.value for i in result["__interrupt__"]) == ["a", "b"]
|
||||
|
||||
|
||||
def test_subgraph_replay_loads_accumulated_state_then_resume(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
|
||||
Generated
+3
-3
@@ -1685,7 +1685,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.2"
|
||||
version = "3.2.0"
|
||||
source = { editable = "../checkpoint-postgres" }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
@@ -1734,7 +1734,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.1"
|
||||
version = "3.1.2"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
@@ -1835,7 +1835,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
description = "Library with high-level APIs for creating and executing LangGraph agents and tools."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
Generated
+3
-3
@@ -418,7 +418,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "3.1.2"
|
||||
version = "3.2.0"
|
||||
source = { editable = "../checkpoint-postgres" }
|
||||
dependencies = [
|
||||
{ name = "langgraph-checkpoint" },
|
||||
@@ -467,7 +467,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.1"
|
||||
version = "3.1.2"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
@@ -512,7 +512,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
@@ -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.6"
|
||||
__version__ = "0.4.7"
|
||||
|
||||
__all__ = [
|
||||
"Auth",
|
||||
|
||||
Generated
+1
-1
@@ -433,7 +433,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-prebuilt"
|
||||
version = "1.1.0"
|
||||
version = "1.1.1"
|
||||
source = { editable = "../prebuilt" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Reference in New Issue
Block a user