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
+29 -3
View File
@@ -16,7 +16,7 @@ from langgraph.checkpoint.base import (
create_checkpoint,
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
@@ -91,11 +91,37 @@ def _base_saver():
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
@@ -155,7 +181,7 @@ def test_data():
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@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"]
@@ -198,7 +224,7 @@ def test_search(saver_name: str, test_data) -> None:
} == {"", "inner"}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
@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(