Compare commits

..
7 changed files with 71 additions and 146 deletions
@@ -507,12 +507,13 @@ class SqliteSaver(BaseCheckpointSaver[str]):
Two-stage query: Two-stage query:
* Stage 1 (streamed): recursive CTE over `checkpoints` following * Stage 1 (paged): newest-first slice of `checkpoints` returning
`parent_checkpoint_id` from the target, returning `(checkpoint_id, parent_checkpoint_id, type, checkpoint)` per
`(checkpoint_id, type, checkpoint)` per ancestor. Sqlite has no ancestor. Sqlite has no JSONB, so we ship the full serialized
JSONB, so we ship the full serialized checkpoint blob and inspect checkpoint blob and inspect `channel_values` in Python. Pages
`channel_values` in Python. Stops reading when every channel has newest-first by `checkpoint_id` with a `< cursor` predicate;
found its seed or the chain is exhausted. page size is `DELTA_PAGE_SIZE`. Stops paging when every channel
has found its seed or the chain is exhausted.
* Stage 2 (per-channel UNION ALL): one branch per channel reading * Stage 2 (per-channel UNION ALL): one branch per channel reading
`writes` filtered to that channel's specific `chain_cids`. No `writes` filtered to that channel's specific `chain_cids`. No
@@ -537,14 +538,12 @@ class SqliteSaver(BaseCheckpointSaver[str]):
seeded: set[str] = set() seeded: set[str] = set()
with self.cursor(transaction=False) as cur: with self.cursor(transaction=False) as cur:
cur.execute( cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id))
DELTA_STAGE1_SQL,
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
)
for row in cur: for row in cur:
cid, type_tag, blob = row cid, parent_cid, type_tag, blob = row
if step_walk_with_row( if step_walk_with_row(
cid=cid, cid=cid,
parent_cid=parent_cid,
type_tag=type_tag, type_tag=type_tag,
blob=blob, blob=blob,
target_id=checkpoint_id, target_id=checkpoint_id,
@@ -26,29 +26,16 @@ from typing import Any
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite
# Stage 1 streams target, then its ancestors nearest-first, by following # Stage 1 streams ancestors of `target_cid` newest-first. The `<=`
# `parent_checkpoint_id`. Ids carry no ordering guarantee, so a range scan by # predicate keeps target itself in the stream so we can read its
# id can miss a parent whose id sorts above its child's. Target is the anchor # `parent_checkpoint_id` from the first row without a separate lookup;
# row; its own writes/seed are skipped (matches the `BaseCheckpointSaver` # the caller skips target's own writes/seed (matches the
# contract). # `BaseCheckpointSaver` contract).
#
# `put` is `INSERT OR REPLACE`, so re-putting an existing id under a
# descendant's config makes the chain a loop. `step_walk_with_row` stops on a
# repeated id; sqlite yields recursive rows lazily, so abandoning the cursor
# ends the recursion.
DELTA_STAGE1_SQL = ( DELTA_STAGE1_SQL = (
"WITH RECURSIVE ancestors(checkpoint_id, parent_checkpoint_id, type, "
"checkpoint) AS ("
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint " "SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
"FROM checkpoints " "FROM checkpoints "
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? " "WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
"UNION ALL " "ORDER BY checkpoint_id DESC"
"SELECT c.checkpoint_id, c.parent_checkpoint_id, c.type, c.checkpoint "
"FROM checkpoints c JOIN ancestors a "
"ON c.checkpoint_id = a.parent_checkpoint_id "
"WHERE c.thread_id = ? AND c.checkpoint_ns = ?"
") "
"SELECT checkpoint_id, type, checkpoint FROM ancestors"
) )
@@ -81,6 +68,7 @@ def build_delta_stage2_sql(*, chain_lens: Sequence[int]) -> str:
def step_walk_with_row( def step_walk_with_row(
*, *,
cid: str, cid: str,
parent_cid: str | None,
type_tag: str, type_tag: str,
blob: bytes, blob: bytes,
target_id: str, target_id: str,
@@ -93,32 +81,36 @@ def step_walk_with_row(
) -> bool: ) -> bool:
"""Process one streamed stage-1 row in the merged ancestor walk. """Process one streamed stage-1 row in the merged ancestor walk.
The cursor returns (cid, type, blob) rows in walk order starting at The cursor returns (cid, parent_cid, type, blob) rows in
target. The first row is target itself and is skipped (target's own `checkpoint_id` DESC order starting at target. The first row is
writes/seed are not part of the contract). target itself; we read its parent_cid to seed the walk and otherwise
skip it (target's own writes/seed are not part of the contract).
For each subsequent row we deserialize the blob, append the cid to For each subsequent row, if `cid` matches the walk's current
every not-yet-seeded channel's chain, and check `channel_values` for position, we deserialize the blob, append the cid to every
not-yet-seeded channel's chain, and check `channel_values` for
seeds. The deserialized checkpoint is dropped before advancing — no seeds. The deserialized checkpoint is dropped before advancing — no
cross-row cache, so peak in-flight is one deserialized checkpoint. cross-row cache, so peak in-flight is one deserialized checkpoint.
Returns True when the caller can stop iterating and close the cursor: Off-path rows (different branch on the same thread) advance the
every requested channel is seeded, or the chain revisited a checkpoint. cursor without doing any work.
Returns True when every requested channel is seeded — the caller
can stop iterating and close the cursor.
""" """
if "started" not in walk_state: if "started" not in walk_state:
if cid == target_id: if cid == target_id:
walk_state["started"] = True walk_state["started"] = True
walk_state["cur_cid"] = parent_cid
walk_state["active"] = {ch for ch in channels if ch not in seeded} walk_state["active"] = {ch for ch in channels if ch not in seeded}
walk_state["walked"] = {cid}
# Not target yet (or target not present): keep streaming. # Not target yet (or target not present): keep streaming.
return False return False
active: set[str] = walk_state["active"] active: set[str] = walk_state["active"]
if not active: if not active:
return True return True
walked: set[str] = walk_state["walked"] if cid != walk_state["cur_cid"]:
if cid in walked: # Off-path row from a sibling branch — skip without deserializing.
return True return False
walked.add(cid)
for ch in active: for ch in active:
chain_by_ch[ch].append(cid) chain_by_ch[ch].append(cid)
ckpt = serde.loads_typed((type_tag, blob)) ckpt = serde.loads_typed((type_tag, blob))
@@ -128,6 +120,7 @@ def step_walk_with_row(
seeded.add(ch) seeded.add(ch)
active.discard(ch) active.discard(ch)
del ckpt, channel_values del ckpt, channel_values
walk_state["cur_cid"] = parent_cid
return not active return not active
@@ -625,8 +625,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
"""Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`. """Fast-path override of `BaseCheckpointSaver.aget_delta_channel_history`.
See `SqliteSaver.get_delta_channel_history` for design notes; this See `SqliteSaver.get_delta_channel_history` for design notes; this
is the async equivalent using `aiosqlite` cursors. Stage 1 streams is the async equivalent using `aiosqlite` cursors. Stage 1 pages
the parent chain from the target and Python-deserializes each the parent chain newest-first and Python-deserializes each
checkpoint blob to find per-channel snapshots; stage 2 fetches checkpoint blob to find per-channel snapshots; stage 2 fetches
only the relevant writes via per-channel UNION ALL. only the relevant writes via per-channel UNION ALL.
""" """
@@ -650,13 +650,13 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
async with self.lock, self.conn.cursor() as cur: async with self.lock, self.conn.cursor() as cur:
await cur.execute( await cur.execute(
DELTA_STAGE1_SQL, DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
) )
async for row in cur: async for row in cur:
cid, type_tag, blob = row cid, parent_cid, type_tag, blob = row
if step_walk_with_row( if step_walk_with_row(
cid=cid, cid=cid,
parent_cid=parent_cid,
type_tag=type_tag, type_tag=type_tag,
blob=blob, blob=blob,
target_id=checkpoint_id, target_id=checkpoint_id,
@@ -1,89 +0,0 @@
from __future__ import annotations
from typing import Any
import pytest
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
Checkpoint,
DeltaChannelHistory,
empty_checkpoint,
)
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
CHANNEL = "ch"
CONFIG: dict[str, Any] = {"configurable": {"thread_id": "t", "checkpoint_ns": ""}}
EXPECTED: DeltaChannelHistory = {
"writes": [("task", CHANNEL, "write-root")],
"seed": "seed",
}
def _checkpoint(checkpoint_id: str, values: dict[str, Any]) -> Checkpoint:
value = empty_checkpoint()
value["id"] = checkpoint_id
value["channel_values"] = values
return value
PARENT_ID_ORDERS = [
pytest.param("z-older", "a-newer", id="parent_id_sorts_above_child"),
pytest.param("a-older", "z-newer", id="parent_id_sorts_below_child"),
]
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
def test_sync_walk_reaches_parent_whatever_the_id_order(
root_id: str, child_id: str
) -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
root = saver.put(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
child = saver.put(root, _checkpoint(child_id, {}), {}, {})
got = saver.get_delta_channel_history(config=child, channels=[CHANNEL])
reference = BaseCheckpointSaver.get_delta_channel_history(
saver, config=child, channels=[CHANNEL]
)
assert got[CHANNEL] == EXPECTED
assert got[CHANNEL] == reference[CHANNEL], "fast path disagrees with base"
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
async def test_async_walk_reaches_parent_whatever_the_id_order(
root_id: str, child_id: str
) -> None:
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
root = await saver.aput(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
child = await saver.aput(root, _checkpoint(child_id, {}), {}, {})
got = await saver.aget_delta_channel_history(config=child, channels=[CHANNEL])
assert got[CHANNEL] == EXPECTED
def test_walk_reaches_root_of_long_chain_with_descending_ids() -> None:
steps = 40
with SqliteSaver.from_conn_string(":memory:") as saver:
parent = saver.put(
CONFIG, _checkpoint(f"id-{steps:03d}", {CHANNEL: "seed"}), {}, {}
)
saver.put_writes(parent, [(CHANNEL, "write-root")], "task")
for step in range(steps - 1, 0, -1):
parent = saver.put(parent, _checkpoint(f"id-{step:03d}", {}), {}, {})
got = saver.get_delta_channel_history(config=parent, channels=[CHANNEL])
assert got[CHANNEL] == EXPECTED
def test_walk_terminates_when_put_makes_the_parent_chain_cycle() -> None:
with SqliteSaver.from_conn_string(":memory:") as saver:
a = saver.put(CONFIG, _checkpoint("cid-a", {}), {}, {})
b = saver.put(a, _checkpoint("cid-b", {}), {}, {})
repoint_a_under_b = _checkpoint("cid-a", {})
saver.put(b, repoint_a_under_b, {}, {})
got = saver.get_delta_channel_history(config=b, channels=[CHANNEL])
assert got[CHANNEL] == {"writes": []}
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.4.31" __version__ = "0.4.32"
+29 -7
View File
@@ -12,6 +12,7 @@ from collections.abc import Callable, Mapping, Sequence
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import asdict, dataclass, field from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone from datetime import datetime, timezone
from functools import partial
from typing import Protocol, TypeVar from typing import Protocol, TypeVar
import click import click
@@ -1690,11 +1691,18 @@ OPT_HOST_URL = click.option(
) )
OPT_AGENT_ID = click.option( OPT_AGENT_ID = click.option(
"--agent-id", help="Logical agent ID (requires agent mode enabled for the tenant)." "--agent-id",
envvar="LANGSMITH_AGENT_ID",
show_envvar=True,
help="Logical agent ID (requires agent mode enabled for the tenant).",
) )
OPT_AGENT_ENVIRONMENT = click.option( OPT_AGENT_ENVIRONMENT = partial(
"--environment", click.option,
"--agent-environment",
"environment",
envvar="LANGSMITH_AGENT_ENVIRONMENT",
show_envvar=True,
type=click.Choice(["development", "staging", "production"]), type=click.Choice(["development", "staging", "production"]),
help="Agent environment (requires agent mode enabled for the tenant).", help="Agent environment (requires agent mode enabled for the tenant).",
) )
@@ -1798,7 +1806,9 @@ def _deploy_base_options(
OPT_HOST_API_KEY, OPT_HOST_API_KEY,
OPT_HOST_DEPLOYMENT_NAME, OPT_HOST_DEPLOYMENT_NAME,
OPT_AGENT_ID, OPT_AGENT_ID,
OPT_AGENT_ENVIRONMENT, OPT_AGENT_ENVIRONMENT()
if include_docker_args
else OPT_AGENT_ENVIRONMENT(type=str),
click.option( click.option(
"--deployment-id", "--deployment-id",
help=( help=(
@@ -1930,6 +1940,12 @@ def deploy(ctx: click.Context, **_: object):
# otherwise, we return None here and click will proceed to actually run the subcommand (list or delete) # otherwise, we return None here and click will proceed to actually run the subcommand (list or delete)
if ctx.invoked_subcommand is not None: if ctx.invoked_subcommand is not None:
return return
environment_param = next(
param for param in _deploy_cmd.params if param.name == "environment"
)
ctx.params["environment"] = environment_param.type_cast_value(
ctx, ctx.params["environment"]
)
if ( if (
ctx.params.get("agent_id") is not None ctx.params.get("agent_id") is not None
or ctx.params.get("environment") is not None or ctx.params.get("environment") is not None
@@ -1982,13 +1998,14 @@ def _deploy_cmd(
validate_deploy_commands(install_command, build_command) validate_deploy_commands(install_command, build_command)
agent = None agent = None
if agent_id is not None or environment is not None: if agent_id is not None or environment is not None:
em.note("Note: --agent-id and --agent-environment flags are in private beta")
if not agent_id or not agent_id.strip() or not environment: if not agent_id or not agent_id.strip() or not environment:
raise click.UsageError( raise click.UsageError(
"--agent-id and --environment are required together." "--agent-id and --agent-environment are required together."
) )
if name is not None or deployment_id is not None: if name is not None or deployment_id is not None:
raise click.UsageError( raise click.UsageError(
"--agent-id and --environment cannot be combined with --name or --deployment-id." "--agent-id and --agent-environment cannot be combined with --name or --deployment-id."
) )
agent = {"agent_id": agent_id, "environment": environment} agent = {"agent_id": agent_id, "environment": environment}
if not config.exists(): if not config.exists():
@@ -2124,7 +2141,7 @@ def _deploy_cmd(
@OPT_HOST_API_KEY @OPT_HOST_API_KEY
@OPT_HOST_URL @OPT_HOST_URL
@OPT_AGENT_ID @OPT_AGENT_ID
@OPT_AGENT_ENVIRONMENT @OPT_AGENT_ENVIRONMENT()
@click.option( @click.option(
"--name-contains", "--name-contains",
default="", default="",
@@ -2138,6 +2155,11 @@ def deploy_list(
agent_id: str | None, agent_id: str | None,
environment: str | None, environment: str | None,
) -> None: ) -> None:
if agent_id is not None or environment is not None:
click.secho(
"Note: --agent-id and --agent-environment flags are in private beta",
fg="yellow",
)
if agent_id is not None and not agent_id.strip(): if agent_id is not None and not agent_id.strip():
raise click.UsageError("--agent-id must not be empty.") raise click.UsageError("--agent-id must not be empty.")
filters = {} filters = {}
@@ -58,7 +58,7 @@ AGENT_ARGS = [
"deploy", "deploy",
"--agent-id", "--agent-id",
"customer-support", "customer-support",
"--environment", "--agent-environment",
"staging", "staging",
"--remote", "--remote",
"--no-wait", "--no-wait",