checkpoint postgres: add a shallow checkpointer (#2826)

This PR adds a "shallow" version of `PostgresSaver` checkpointer that
ONLY stores the most recent checkpoint and does NOT retain any history.
It is meant to be a light-weight drop-in replacement for the
PostgresSaver that supports most of the LangGraph persistence
functionality with the exception of time travel.
This commit is contained in:
Vadym Barda
2024-12-20 17:51:20 +00:00
committed by GitHub
parent 1de61f6fce
commit 44ee0199fd
14 changed files with 4322 additions and 1082 deletions
+36 -3
View File
@@ -16,7 +16,10 @@ from langgraph.checkpoint.base import (
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.checkpoint.postgres.aio import (
AsyncPostgresSaver,
AsyncShallowPostgresSaver,
)
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -103,11 +106,41 @@ async def _base_saver():
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
@@ -167,7 +200,7 @@ def test_data():
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
async def test_asearch(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
configs = test_data["configs"]
@@ -212,7 +245,7 @@ async def test_asearch(request, saver_name: str, test_data) -> None:
} == {"", "inner"}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe", "shallow"])
async def test_null_chars(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
config = await saver.aput(