wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
361 lines
12 KiB
Python
361 lines
12 KiB
Python
# type: ignore
|
|
|
|
import re
|
|
from contextlib import contextmanager
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from langchain_core.runnables import RunnableConfig
|
|
from langgraph.checkpoint.base import (
|
|
EXCLUDED_METADATA_KEYS,
|
|
Checkpoint,
|
|
CheckpointMetadata,
|
|
create_checkpoint,
|
|
empty_checkpoint,
|
|
)
|
|
from langgraph.checkpoint.serde.types import TASKS
|
|
from psycopg import Connection
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
|
|
from tests.conftest import DEFAULT_POSTGRES_URI
|
|
|
|
|
|
def _exclude_keys(config: dict[str, Any]) -> dict[str, Any]:
|
|
return {k: v for k, v in config.items() if k not in EXCLUDED_METADATA_KEYS}
|
|
|
|
|
|
@contextmanager
|
|
def _pool_saver():
|
|
"""Fixture for pool mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
# yield checkpointer
|
|
with ConnectionPool(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
max_size=10,
|
|
kwargs={"autocommit": True, "row_factory": dict_row},
|
|
) as pool:
|
|
checkpointer = PostgresSaver(pool)
|
|
checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@contextmanager
|
|
def _pipe_saver():
|
|
"""Fixture for pipeline mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
with Connection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = PostgresSaver(conn)
|
|
checkpointer.setup()
|
|
with conn.pipeline() as pipe:
|
|
checkpointer = PostgresSaver(conn, pipe=pipe)
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@contextmanager
|
|
def _base_saver():
|
|
"""Fixture for regular connection mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
with Connection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = PostgresSaver(conn)
|
|
checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@contextmanager
|
|
def _shallow_saver():
|
|
"""Fixture for regular connection mode testing with a shallow checkpointer."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
with Connection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = ShallowPostgresSaver(conn)
|
|
checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
|
|
conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@contextmanager
|
|
def _saver(name: str):
|
|
if name == "base":
|
|
with _base_saver() as saver:
|
|
yield saver
|
|
elif name == "shallow":
|
|
with _shallow_saver() as saver:
|
|
yield saver
|
|
elif name == "pool":
|
|
with _pool_saver() as saver:
|
|
yield saver
|
|
elif name == "pipe":
|
|
with _pipe_saver() as saver:
|
|
yield saver
|
|
|
|
|
|
@pytest.fixture
|
|
def test_data():
|
|
"""Fixture providing test data for checkpoint tests."""
|
|
config_1: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_id": "1",
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
config_2: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_id": "2",
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
config_3: RunnableConfig = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_id": "2-inner",
|
|
"checkpoint_ns": "inner",
|
|
}
|
|
}
|
|
|
|
chkpnt_1: Checkpoint = empty_checkpoint()
|
|
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
|
|
chkpnt_3: Checkpoint = empty_checkpoint()
|
|
|
|
metadata_1: CheckpointMetadata = {
|
|
"source": "input",
|
|
"step": 2,
|
|
"score": 1,
|
|
}
|
|
metadata_2: CheckpointMetadata = {
|
|
"source": "loop",
|
|
"step": 1,
|
|
"score": None,
|
|
}
|
|
metadata_3: CheckpointMetadata = {}
|
|
|
|
return {
|
|
"configs": [config_1, config_2, config_3],
|
|
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
|
|
"metadata": [metadata_1, metadata_2, metadata_3],
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
|
def test_combined_metadata(saver_name: str, test_data) -> None:
|
|
with _saver(saver_name) as saver:
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_ns": "",
|
|
"__super_private_key": "super_private_value",
|
|
},
|
|
"metadata": {"run_id": "my_run_id"},
|
|
}
|
|
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
|
metadata: CheckpointMetadata = {
|
|
"source": "loop",
|
|
"step": 1,
|
|
"score": None,
|
|
}
|
|
saver.put(config, chkpnt, metadata, {})
|
|
checkpoint = saver.get_tuple(config)
|
|
assert checkpoint.metadata == {
|
|
**metadata,
|
|
"run_id": "my_run_id",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
|
def test_search(saver_name: str, test_data) -> None:
|
|
with _saver(saver_name) as saver:
|
|
configs = test_data["configs"]
|
|
checkpoints = test_data["checkpoints"]
|
|
metadata = test_data["metadata"]
|
|
|
|
saver.put(configs[0], checkpoints[0], metadata[0], {})
|
|
saver.put(configs[1], checkpoints[1], metadata[1], {})
|
|
saver.put(configs[2], checkpoints[2], metadata[2], {})
|
|
|
|
# call method / assertions
|
|
query_1 = {"source": "input"} # search by 1 key
|
|
query_2 = {
|
|
"step": 1,
|
|
} # search by multiple keys
|
|
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
|
|
query_4 = {"source": "update", "step": 1} # no match
|
|
|
|
search_results_1 = list(saver.list(None, filter=query_1))
|
|
assert len(search_results_1) == 1
|
|
assert search_results_1[0].metadata == {
|
|
**_exclude_keys(configs[0]["configurable"]),
|
|
**metadata[0],
|
|
}
|
|
|
|
search_results_2 = list(saver.list(None, filter=query_2))
|
|
assert len(search_results_2) == 1
|
|
assert search_results_2[0].metadata == {
|
|
**_exclude_keys(configs[1]["configurable"]),
|
|
**metadata[1],
|
|
}
|
|
|
|
search_results_3 = list(saver.list(None, filter=query_3))
|
|
assert len(search_results_3) == 3
|
|
|
|
search_results_4 = list(saver.list(None, filter=query_4))
|
|
assert len(search_results_4) == 0
|
|
|
|
# search by config (defaults to checkpoints across all namespaces)
|
|
search_results_5 = list(saver.list({"configurable": {"thread_id": "thread-2"}}))
|
|
assert len(search_results_5) == 2
|
|
assert {
|
|
search_results_5[0].config["configurable"]["checkpoint_ns"],
|
|
search_results_5[1].config["configurable"]["checkpoint_ns"],
|
|
} == {"", "inner"}
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
|
def test_null_chars(saver_name: str, test_data) -> None:
|
|
with _saver(saver_name) as saver:
|
|
config = saver.put(
|
|
test_data["configs"][0],
|
|
test_data["checkpoints"][0],
|
|
{"my_key": "\x00abc"},
|
|
{},
|
|
)
|
|
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
|
|
assert (
|
|
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]
|
|
== "abc"
|
|
)
|
|
|
|
|
|
def test_nonnull_migrations() -> None:
|
|
_leading_comment_remover = re.compile(r"^/\*.*?\*/")
|
|
for migration in PostgresSaver.MIGRATIONS:
|
|
statement = _leading_comment_remover.sub("", migration).split()[0]
|
|
assert statement.strip()
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
|
def test_pending_sends_migration(saver_name: str) -> None:
|
|
with _saver(saver_name) as saver:
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": "thread-1",
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
|
|
# create the first checkpoint
|
|
# and put some pending sends
|
|
checkpoint_0 = empty_checkpoint()
|
|
config = saver.put(config, checkpoint_0, {}, {})
|
|
saver.put_writes(
|
|
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
|
|
)
|
|
saver.put_writes(config, [(TASKS, "send-3")], task_id="task-2")
|
|
|
|
# check that fetching checkpoint_0 doesn't attach pending sends
|
|
# (they should be attached to the next checkpoint)
|
|
tuple_0 = saver.get_tuple(config)
|
|
assert tuple_0.checkpoint["channel_values"] == {}
|
|
assert tuple_0.checkpoint["channel_versions"] == {}
|
|
|
|
# create the second checkpoint
|
|
checkpoint_1 = create_checkpoint(checkpoint_0, {}, 1)
|
|
config = saver.put(config, checkpoint_1, {}, {})
|
|
|
|
# check that pending sends are attached to checkpoint_1
|
|
checkpoint_1 = saver.get_tuple(config)
|
|
assert checkpoint_1.checkpoint["channel_values"] == {
|
|
TASKS: ["send-1", "send-2", "send-3"]
|
|
}
|
|
assert TASKS in checkpoint_1.checkpoint["channel_versions"]
|
|
|
|
# check that list also applies the migration
|
|
search_results = [
|
|
c for c in saver.list({"configurable": {"thread_id": "thread-1"}})
|
|
]
|
|
assert len(search_results) == 2
|
|
assert search_results[-1].checkpoint["channel_values"] == {}
|
|
assert search_results[-1].checkpoint["channel_versions"] == {}
|
|
assert search_results[0].checkpoint["channel_values"] == {
|
|
TASKS: ["send-1", "send-2", "send-3"]
|
|
}
|
|
assert TASKS in search_results[0].checkpoint["channel_versions"]
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
|
def test_get_checkpoint_no_channel_values(
|
|
monkeypatch, saver_name: str, test_data
|
|
) -> None:
|
|
"""Backwards compatibility test that verifies a checkpoint with no channel_values key can be retrieved without throwing an error."""
|
|
with _saver(saver_name) as saver:
|
|
config = {
|
|
"configurable": {
|
|
"thread_id": "thread-2",
|
|
"checkpoint_ns": "",
|
|
"__super_private_key": "super_private_value",
|
|
},
|
|
}
|
|
chkpnt: Checkpoint = create_checkpoint(empty_checkpoint(), {}, 1)
|
|
saver.put(config, chkpnt, {}, {})
|
|
|
|
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
|
|
|
def patched_load_checkpoint_tuple(value):
|
|
value["checkpoint"].pop("channel_values", None)
|
|
return load_checkpoint_tuple(value)
|
|
|
|
monkeypatch.setattr(
|
|
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
|
)
|
|
|
|
checkpoint = saver.get_tuple(config)
|
|
assert checkpoint.checkpoint["channel_values"] == {}
|