Compare commits

..
5 changed files with 33 additions and 146 deletions
@@ -441,12 +441,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
for i, ch in enumerate(channels): for i, ch in enumerate(channels):
if ch in seeded: if ch in seeded:
continue continue
# Pages start at the thread head, so the target may not have # First-time entry: cursor starts at the target's parent.
# loaded yet; a `None` cursor would read as "target is a root".
if ch not in walk_cursor_by_ch: if ch not in walk_cursor_by_ch:
if target_id not in parent_of: walk_cursor_by_ch[ch] = parent_of.get(target_id)
continue
walk_cursor_by_ch[ch] = parent_of[target_id]
cur_cid = walk_cursor_by_ch[ch] cur_cid = walk_cursor_by_ch[ch]
ch_chain = chain_by_ch[ch] ch_chain = chain_by_ch[ch]
hb_i = hb_by_i_by_cid[i] hb_i = hb_by_i_by_cid[i]
@@ -1,132 +0,0 @@
from __future__ import annotations
from typing import Any
from uuid import uuid4
import pytest
from langgraph.checkpoint.base import (
Checkpoint,
DeltaChannelHistory,
empty_checkpoint,
)
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.types import _DeltaSnapshot
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE
from tests.conftest import DEFAULT_URI
CHANNEL = "items"
STEPS = 8
SEED_STEP = 1
SEED_VALUE = [10, 20]
TARGET_STEP = 4
# The real page size is the control; the rest leave the target off the first
# page (three checkpoints are newer than it).
PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1]
def _step_args(
thread_id: str, step: int, parent: dict | None
) -> tuple[dict, Checkpoint, dict[str, Any]]:
config: dict = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
if parent is not None:
config["configurable"]["checkpoint_id"] = parent["configurable"][
"checkpoint_id"
]
checkpoint: Checkpoint = empty_checkpoint()
checkpoint["id"] = str(uuid6(clock_seq=step))
checkpoint["channel_versions"][CHANNEL] = f"v{step}"
if step == SEED_STEP:
checkpoint["channel_values"][CHANNEL] = _DeltaSnapshot(list(SEED_VALUE))
return config, checkpoint, {CHANNEL: f"v{step}"}
return config, checkpoint, {}
async def _abuild_chain(saver: AsyncPostgresSaver) -> list[dict]:
thread_id = str(uuid4())
parent: dict | None = None
configs: list[dict] = []
for step in range(STEPS):
config, checkpoint, new_versions = _step_args(thread_id, step, parent)
parent = await saver.aput(
config,
checkpoint,
{"source": "loop", "step": step, "parents": {}},
new_versions,
)
await saver.aput_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
configs.append(parent)
return configs
def _build_chain(saver: PostgresSaver) -> list[dict]:
thread_id = str(uuid4())
parent: dict | None = None
configs: list[dict] = []
for step in range(STEPS):
config, checkpoint, new_versions = _step_args(thread_id, step, parent)
parent = saver.put(
config,
checkpoint,
{"source": "loop", "step": step, "parents": {}},
new_versions,
)
saver.put_writes(parent, [(CHANNEL, f"w{step}")], str(uuid4()))
configs.append(parent)
return configs
def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None:
seed = entry.get("seed")
assert isinstance(seed, _DeltaSnapshot), (
f"page_size={page_size}: expected a snapshot seed, "
f"got {entry.get('seed', '<missing>')!r}"
)
assert seed.value == SEED_VALUE
assert [w[2] for w in entry["writes"]] == ["w1", "w2", "w3"], (
f"page_size={page_size}: got {[w[2] for w in entry['writes']]}"
)
@pytest.mark.parametrize("page_size", PAGE_SIZES)
async def test_async_target_older_than_the_first_page(
page_size: int, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size)
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
configs = await _abuild_chain(saver)
result = await saver.aget_delta_channel_history(
config=configs[TARGET_STEP], channels=[CHANNEL]
)
_assert_history(result[CHANNEL], page_size)
@pytest.mark.parametrize("page_size", PAGE_SIZES)
def test_sync_target_older_than_the_first_page(
page_size: int, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size)
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
configs = _build_chain(saver)
result = saver.get_delta_channel_history(
config=configs[TARGET_STEP], channels=[CHANNEL]
)
_assert_history(result[CHANNEL], page_size)
async def test_root_target_has_no_history_and_still_terminates(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", 1)
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
configs = await _abuild_chain(saver)
result = await saver.aget_delta_channel_history(
config=configs[0], channels=[CHANNEL]
)
assert result[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",