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
418 lines
14 KiB
Python
418 lines
14 KiB
Python
# type: ignore
|
|
|
|
from contextlib import asynccontextmanager
|
|
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 AsyncConnection
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import AsyncConnectionPool
|
|
|
|
from langgraph.checkpoint.postgres.aio import (
|
|
AsyncPostgresSaver,
|
|
AsyncShallowPostgresSaver,
|
|
)
|
|
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}
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _pool_saver():
|
|
"""Fixture for pool mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
# yield checkpointer
|
|
async with AsyncConnectionPool(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
max_size=10,
|
|
kwargs={"autocommit": True, "row_factory": dict_row},
|
|
) as pool:
|
|
checkpointer = AsyncPostgresSaver(pool)
|
|
await checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _pipe_saver():
|
|
"""Fixture for pipeline mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = AsyncPostgresSaver(conn)
|
|
await checkpointer.setup()
|
|
async with conn.pipeline() as pipe:
|
|
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _base_saver():
|
|
"""Fixture for regular connection mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = AsyncPostgresSaver(conn)
|
|
await checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _shallow_saver():
|
|
"""Fixture for shallow connection mode testing."""
|
|
database = f"test_{uuid4().hex[:16]}"
|
|
# create unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"CREATE DATABASE {database}")
|
|
try:
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI + database,
|
|
autocommit=True,
|
|
prepare_threshold=0,
|
|
row_factory=dict_row,
|
|
) as conn:
|
|
checkpointer = AsyncShallowPostgresSaver(conn)
|
|
await checkpointer.setup()
|
|
yield checkpointer
|
|
finally:
|
|
# drop unique db
|
|
async with await AsyncConnection.connect(
|
|
DEFAULT_POSTGRES_URI, autocommit=True
|
|
) as conn:
|
|
await conn.execute(f"DROP DATABASE {database}")
|
|
|
|
|
|
@asynccontextmanager
|
|
async def _saver(name: str):
|
|
if name == "base":
|
|
async with _base_saver() as saver:
|
|
yield saver
|
|
elif name == "shallow":
|
|
async with _shallow_saver() as saver:
|
|
yield saver
|
|
elif name == "pool":
|
|
async with _pool_saver() as saver:
|
|
yield saver
|
|
elif name == "pipe":
|
|
async 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"])
|
|
async def test_combined_metadata(saver_name: str, test_data) -> None:
|
|
async 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,
|
|
}
|
|
await saver.aput(config, chkpnt, metadata, {})
|
|
checkpoint = await saver.aget_tuple(config)
|
|
assert checkpoint.metadata == {
|
|
**metadata,
|
|
"run_id": "my_run_id",
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
|
|
async def test_asearch(saver_name: str, test_data) -> None:
|
|
async with _saver(saver_name) as saver:
|
|
configs = test_data["configs"]
|
|
checkpoints = test_data["checkpoints"]
|
|
metadata = test_data["metadata"]
|
|
|
|
await saver.aput(configs[0], checkpoints[0], metadata[0], {})
|
|
await saver.aput(configs[1], checkpoints[1], metadata[1], {})
|
|
await saver.aput(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 = [c async for c in saver.alist(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 = [c async for c in saver.alist(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 = [c async for c in saver.alist(None, filter=query_3)]
|
|
assert len(search_results_3) == 3
|
|
|
|
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
|
|
assert len(search_results_4) == 0
|
|
|
|
# search by config (defaults to checkpoints across all namespaces)
|
|
search_results_5 = [
|
|
c async for c in saver.alist({"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"])
|
|
async def test_null_chars(saver_name: str, test_data) -> None:
|
|
async with _saver(saver_name) as saver:
|
|
config = await saver.aput(
|
|
test_data["configs"][0],
|
|
test_data["checkpoints"][0],
|
|
{"my_key": "\x00abc"},
|
|
{},
|
|
)
|
|
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
|
|
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
|
|
0
|
|
].metadata["my_key"] == "abc"
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
|
async def test_pending_sends_migration(saver_name: str) -> None:
|
|
async 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 = await saver.aput(config, checkpoint_0, {}, {})
|
|
await saver.aput_writes(
|
|
config, [(TASKS, "send-1"), (TASKS, "send-2")], task_id="task-1"
|
|
)
|
|
await saver.aput_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 = await saver.aget_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 = await saver.aput(config, checkpoint_1, {}, {})
|
|
|
|
# check that pending sends are attached to checkpoint_1
|
|
tuple_1 = await saver.aget_tuple(config)
|
|
assert tuple_1.checkpoint["channel_values"] == {
|
|
TASKS: ["send-1", "send-2", "send-3"]
|
|
}
|
|
assert TASKS in tuple_1.checkpoint["channel_versions"]
|
|
|
|
# check that list also applies the migration
|
|
search_results = [
|
|
c async for c in saver.alist({"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"])
|
|
async 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."""
|
|
async 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)
|
|
await saver.aput(config, chkpnt, {}, {})
|
|
|
|
load_checkpoint_tuple = saver._load_checkpoint_tuple
|
|
|
|
async def patched_load_checkpoint_tuple(value):
|
|
value["checkpoint"].pop("channel_values", None)
|
|
return await load_checkpoint_tuple(value)
|
|
|
|
monkeypatch.setattr(
|
|
saver, "_load_checkpoint_tuple", patched_load_checkpoint_tuple
|
|
)
|
|
|
|
checkpoint = await saver.aget_tuple(config)
|
|
assert checkpoint.checkpoint["channel_values"] == {}
|
|
|
|
|
|
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
|
|
async def test_delta_channel_chain_reconstruction(saver_name: str) -> None:
|
|
"""AsyncPostgresSaver reconstructs DeltaChannel chain via point-lookup traversal."""
|
|
pytest.importorskip(
|
|
"langgraph.channels.delta", reason="langgraph core not installed"
|
|
)
|
|
|
|
from typing import Annotated
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
from langgraph.channels.delta import DeltaChannel
|
|
from langgraph.graph import START, StateGraph
|
|
from langgraph.graph.message import _messages_delta_reducer
|
|
from typing_extensions import TypedDict
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
|
|
|
def respond(state: State) -> dict:
|
|
n = len(state["messages"])
|
|
return {"messages": [AIMessage(content=f"reply-{n}", id=f"ai-{n}")]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
|
|
async with _saver(saver_name) as saver:
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "diff-channel-test-1"}}
|
|
|
|
await graph.ainvoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
|
await graph.ainvoke(
|
|
{"messages": [HumanMessage(content="there", id="h2")]}, config
|
|
)
|
|
|
|
state = await graph.aget_state(config)
|
|
msgs = state.values["messages"]
|
|
assert len(msgs) == 4, f"expected 4, got {len(msgs)}: {msgs}"
|
|
assert msgs[0].content == "hi"
|
|
assert msgs[1].content == "reply-1"
|
|
assert msgs[2].content == "there"
|
|
assert msgs[3].content == "reply-3"
|