Add TTL Sweeper (#3849)

This commit is contained in:
William FH
2025-03-14 14:29:40 -07:00
committed by GitHub
12 changed files with 724 additions and 223 deletions
@@ -2,6 +2,7 @@ import asyncio
import logging
from collections.abc import AsyncIterator, Iterable, Sequence
from contextlib import asynccontextmanager
from types import TracebackType
from typing import Any, Callable, Optional, Union, cast
import orjson
@@ -25,6 +26,7 @@ from langgraph.store.postgres.base import (
PoolConfig,
PostgresIndexConfig,
Row,
TTLConfig,
_decode_ns_bytes,
_ensure_index_config,
_group_ops,
@@ -106,6 +108,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
Semantic search is disabled by default. You can enable it by providing an `index` configuration
when creating the store. Without this configuration, all `index` arguments passed to
`put` or `aput` will have no effect.
Note:
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
the background task that removes expired items. Call `stop_ttl_sweeper()` to properly
clean up resources when you're done with the store.
"""
__slots__ = (
@@ -115,7 +122,11 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
"supports_pipeline",
"index_config",
"embeddings",
"ttl_config",
"_ttl_sweeper_task",
"_ttl_stop_event",
)
supports_ttl: bool = True
def __init__(
self,
@@ -126,6 +137,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> None:
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
raise ValueError(
@@ -141,10 +153,13 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
self.index_config = index
if self.index_config:
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
else:
self.embeddings = None
self.ttl_config = ttl
self._ttl_sweeper_task: Optional[asyncio.Task[None]] = None
self._ttl_stop_event = asyncio.Event()
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
results: list[Result] = [None] * num_ops
@@ -167,6 +182,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> AsyncIterator["AsyncPostgresStore"]:
"""Create a new AsyncPostgresStore instance from a connection string.
@@ -198,16 +214,16 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool, index=index)
yield cls(conn=pool, index=index, ttl=ttl)
else:
async with await AsyncConnection.connect(
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
if pipeline:
async with conn.pipeline() as pipe:
yield cls(conn=conn, pipe=pipe, index=index)
yield cls(conn=conn, pipe=pipe, index=index, ttl=ttl)
else:
yield cls(conn=conn, index=index)
yield cls(conn=conn, index=index, ttl=ttl)
async def setup(self) -> None:
"""Set up the store database asynchronously.
@@ -256,6 +272,119 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
)
async def sweep_ttl(self) -> int:
"""Delete expired store items based on TTL.
Returns:
int: The number of deleted items.
"""
async with self._cursor() as cur:
await cur.execute(
"""
DELETE FROM store
WHERE expires_at IS NOT NULL AND expires_at < NOW()
"""
)
deleted_count = cur.rowcount
return deleted_count
async def start_ttl_sweeper(
self, sweep_interval_minutes: Optional[int] = None
) -> asyncio.Task[None]:
"""Periodically delete expired store items based on TTL.
Returns:
Task that can be awaited or cancelled.
"""
if not self.ttl_config:
return asyncio.create_task(asyncio.sleep(0))
if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
return self._ttl_sweeper_task
self._ttl_stop_event.clear()
interval = float(
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
)
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
async def _sweep_loop() -> None:
while not self._ttl_stop_event.is_set():
try:
try:
await asyncio.wait_for(
self._ttl_stop_event.wait(),
timeout=interval * 60,
)
break
except asyncio.TimeoutError:
pass
expired_items = await self.sweep_ttl()
if expired_items > 0:
logger.info(f"Store swept {expired_items} expired items")
except asyncio.CancelledError:
break
except Exception as exc:
logger.exception("Store TTL sweep iteration failed", exc_info=exc)
task = asyncio.create_task(_sweep_loop())
task.set_name("ttl_sweeper")
self._ttl_sweeper_task = task
return task
async def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
"""Stop the TTL sweeper task if it's running.
Args:
timeout: Maximum time to wait for the task to stop, in seconds.
If None, wait indefinitely.
Returns:
bool: True if the task was successfully stopped or wasn't running,
False if the timeout was reached before the task stopped.
"""
if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
return True
logger.info("Stopping TTL sweeper task")
self._ttl_stop_event.set()
if timeout is not None:
try:
await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
success = True
except asyncio.TimeoutError:
success = False
else:
await self._ttl_sweeper_task
success = True
if success:
self._ttl_sweeper_task = None
logger.info("TTL sweeper task stopped")
else:
logger.warning("Timed out waiting for TTL sweeper task to stop")
return success
async def __aenter__(self) -> "AsyncPostgresStore":
return self
async def __aexit__(
self,
exc_type: Optional[type[BaseException]],
exc_val: Optional[BaseException],
exc_tb: Optional["TracebackType"],
) -> None:
# Ensure the TTL sweeper task is stopped when exiting the context
if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
# Set the event to signal the task to stop
self._ttl_stop_event.set()
# We don't wait for the task to complete here to avoid blocking
# The task will clean up itself gracefully
async def _execute_batch(
self,
grouped_ops: dict,
@@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import json
import logging
import threading
@@ -74,6 +75,17 @@ CREATE TABLE IF NOT EXISTS store (
"""
-- For faster lookups by prefix
CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pattern_ops);
""",
"""
-- Add expires_at column to store table
ALTER TABLE store
ADD COLUMN expires_at TIMESTAMP WITH TIME ZONE,
ADD COLUMN ttl_minutes INT;
""",
"""
-- Add indexes for efficient TTL sweeping
CREATE INDEX idx_store_expires_at ON store (expires_at)
WHERE expires_at IS NOT NULL;
""",
]
@@ -225,20 +237,55 @@ class BasePostgresStore(Generic[C]):
self,
get_ops: Sequence[tuple[int, GetOp]],
) -> list[tuple[str, tuple, tuple[str, ...], list]]:
"""
Build queries to fetch (and optionally refresh the TTL of) multiple keys per namespace.
Each returned element is a tuple of:
(sql_query_string, sql_params, namespace, items_for_this_namespace)
where items_for_this_namespace is the original list of (idx, key, refresh_ttl).
"""
namespace_groups = defaultdict(list)
refresh_ttls = defaultdict(list)
for idx, op in get_ops:
namespace_groups[op.namespace].append((idx, op.key))
refresh_ttls[op.namespace].append(op.refresh_ttl)
results = []
for namespace, items in namespace_groups.items():
_, keys = zip(*items)
keys_to_query = ",".join(["%s"] * len(keys))
query = f"""
SELECT key, value, created_at, updated_at
FROM store
WHERE prefix = %s AND key IN ({keys_to_query})
this_refresh_ttls = refresh_ttls[namespace]
query = """
WITH passed_in AS (
SELECT unnest(%s::text[]) AS key,
unnest(%s::bool[]) AS do_refresh
),
updated AS (
UPDATE store s
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
FROM passed_in p
WHERE s.prefix = %s
AND s.key = p.key
AND p.do_refresh = TRUE
AND s.ttl_minutes IS NOT NULL
RETURNING s.key
)
SELECT s.key, s.value, s.created_at, s.updated_at
FROM store s
JOIN passed_in p ON s.key = p.key
WHERE s.prefix = %s
"""
params = (_namespace_to_text(namespace), *keys)
ns_text = _namespace_to_text(namespace)
params = (
list(keys), # -> unnest(%s::text[])
list(this_refresh_ttls), # -> unnest(%s::bool[])
ns_text, # -> prefix = %s (for UPDATE)
ns_text, # -> prefix = %s (for final SELECT)
)
results.append((query, params, namespace, items))
return results
def _prepare_batch_PUT_queries(
@@ -248,7 +295,6 @@ class BasePostgresStore(Generic[C]):
list[tuple[str, Sequence]],
Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
]:
# Last-write wins
dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
for _, op in put_ops:
dedupped_ops[(op.namespace, op.key)] = op
@@ -282,15 +328,26 @@ class BasePostgresStore(Generic[C]):
insertion_params = []
vector_values = []
embedding_request_params = []
# Handle TTL expiration
# First handle main store insertions
for op in inserts:
values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
if op.ttl is not None:
expires_at_str = f"NOW() + INTERVAL '{op.ttl*60} seconds'"
ttl_minutes = op.ttl
else:
expires_at_str = "NULL"
ttl_minutes = None
values.append(
f"(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, {expires_at_str}, %s)"
)
insertion_params.extend(
[
_namespace_to_text(op.namespace),
op.key,
Jsonb(cast(dict, op.value)),
ttl_minutes,
]
)
@@ -304,7 +361,7 @@ class BasePostgresStore(Generic[C]):
k = op.key
if op.index is None:
paths = self.index_config["__tokenized_fields"]
paths = cast(dict, self.index_config)["__tokenized_fields"]
else:
paths = [(ix, tokenize_path(ix)) for ix in op.index]
@@ -319,11 +376,13 @@ class BasePostgresStore(Generic[C]):
values_str = ",".join(values)
query = f"""
INSERT INTO store (prefix, key, value, created_at, updated_at)
INSERT INTO store (prefix, key, value, created_at, updated_at, expires_at, ttl_minutes)
VALUES {values_str}
ON CONFLICT (prefix, key) DO UPDATE
SET value = EXCLUDED.value,
updated_at = CURRENT_TIMESTAMP
updated_at = CURRENT_TIMESTAMP,
expires_at = EXCLUDED.expires_at,
ttl_minutes = EXCLUDED.ttl_minutes
"""
queries.append((query, insertion_params))
@@ -347,92 +406,105 @@ class BasePostgresStore(Generic[C]):
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
list[tuple[int, str]], # idx, query_text pairs to embed
]:
"""
Build per-SearchOp SQL queries (with optional TTL refresh) plus embedding requests.
Returns:
- queries: list of (SQL, param_list)
- embedding_requests: list of (original_index_in_search_ops, text_query)
"""
queries = []
embedding_requests = []
for idx, (_, op) in enumerate(search_ops):
# Build filter conditions first
filter_params = []
filter_conditions = []
filter_clauses = []
if op.filter:
for key, value in op.filter.items():
if isinstance(value, dict):
for op_name, val in value.items():
condition, filter_params_ = self._get_filter_condition(
condition, params_ = self._get_filter_condition(
key, op_name, val
)
filter_conditions.append(condition)
filter_params.extend(filter_params_)
filter_clauses.append(condition)
filter_params.extend(params_)
else:
filter_conditions.append("value->%s = %s::jsonb")
filter_params.extend([key, json.dumps(value)])
filter_clauses.append("value->%s = %s::jsonb")
filter_params.extend([key, orjson.dumps(value).decode("utf-8")])
ns_condition = "TRUE"
ns_param: Optional[Sequence[Union[str]]] = None
if op.namespace_prefix:
ns_condition = "store.prefix LIKE %s"
ns_param = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_param = ()
extra_filters = (
" AND " + " AND ".join(filter_clauses) if filter_clauses else ""
)
# Vector search branch
if op.query and self.index_config:
# We'll embed the text later, so record the request.
embedding_requests.append((idx, op.query))
score_operator, post_operator = get_distance_operator(self)
post_operator = post_operator.replace("scored", "uniq")
vector_type = (
cast(PostgresIndexConfig, self.index_config)
.get("ann_index_config", {})
.get("vector_type", "vector")
)
# For hamming bit vectors, or “regular” vectors
if (
vector_type == "bit"
and self.index_config.get("distance_type") == "hamming"
and cast(dict, self.index_config).get("distance_type") == "hamming"
):
score_operator = score_operator % (
"%s",
self.index_config["dims"],
cast(dict, self.index_config)["dims"],
)
else:
score_operator = score_operator % (
"%s",
vector_type,
)
score_operator = score_operator % ("%s", vector_type)
vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
vectors_per_doc_estimate = cast(dict, self.index_config)[
"__estimated_num_vectors"
]
expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
# Vector search with CTE for proper score handling
filter_str = (
""
if not filter_conditions
else " AND " + " AND ".join(filter_conditions)
)
if op.namespace_prefix:
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_args = ()
if filter_str:
prefix_filter_str = f"WHERE {filter_str} "
else:
prefix_filter_str = ""
base_query = f"""
WITH scored AS (
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
FROM store s
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
{prefix_filter_str}
ORDER BY {score_operator} ASC
# “sub_scored” does the main vector search
# Then we do DISTINCT ON to drop duplicates if your store can have them
# Finally we limit & offset
vector_search_cte = f"""
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at,
{score_operator} AS neg_score
FROM store
JOIN store_vectors sv ON store.prefix = sv.prefix AND store.key = sv.key
WHERE {ns_condition} {extra_filters}
ORDER BY {score_operator} ASC
LIMIT %s
)
SELECT * FROM (
SELECT DISTINCT ON (prefix, key)
prefix, key, value, created_at, updated_at, {post_operator} as score
FROM scored
ORDER BY prefix, key, score DESC
) AS unique_docs
ORDER BY score DESC
LIMIT %s
OFFSET %s
"""
params = [
PLACEHOLDER, # Vector placeholder
*ns_args,
"""
search_results_sql = f"""
WITH scored AS (
{vector_search_cte}
)
SELECT uniq.prefix, uniq.key, uniq.value, uniq.created_at, uniq.updated_at,
{post_operator} AS score
FROM (
SELECT DISTINCT ON (scored.prefix, scored.key)
scored.prefix, scored.key, scored.value, scored.created_at, scored.updated_at, scored.neg_score
FROM scored
ORDER BY scored.prefix, scored.key, scored.neg_score ASC
) uniq
ORDER BY score DESC
LIMIT %s
OFFSET %s
"""
search_results_params = [
PLACEHOLDER,
*ns_param,
*filter_params,
PLACEHOLDER,
expanded_limit,
@@ -440,24 +512,45 @@ class BasePostgresStore(Generic[C]):
op.offset,
]
# Regular search branch
else:
base_query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
base_query = f"""
SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, NULL AS score
FROM store
WHERE {ns_condition} {extra_filters}
ORDER BY store.updated_at DESC
LIMIT %s
OFFSET %s
"""
search_results_sql = base_query
search_results_params = [
*ns_param,
*filter_params,
op.limit,
op.offset,
]
if filter_conditions:
params.extend(filter_params)
base_query += " AND " + " AND ".join(filter_conditions)
base_query += " ORDER BY updated_at DESC"
base_query += " LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
queries.append((base_query, params))
if op.refresh_ttl:
# Wrap entire primary query in a CTE, then perform "update_at"
final_sql = f"""
WITH search_results AS (
{search_results_sql}
),
updated AS (
UPDATE store s
SET expires_at = NOW() + (s.ttl_minutes || ' minutes')::interval
FROM search_results sr
WHERE s.prefix = sr.prefix
AND s.key = sr.key
AND s.ttl_minutes IS NOT NULL
)
SELECT sr.prefix, sr.key, sr.value, sr.created_at, sr.updated_at, sr.score
FROM search_results sr
"""
final_params = search_results_params[:] # copy
else:
final_sql = search_results_sql
final_params = search_results_params
queries.append((final_sql, final_params))
return queries, embedding_requests
@@ -603,6 +696,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
Make sure to call `setup()` before first use to create necessary tables and indexes.
The pgvector extension must be available to use vector search.
Note:
If you provide a TTL configuration, you must explicitly call `start_ttl_sweeper()` to begin
the background thread that removes expired items. Call `stop_ttl_sweeper()` to properly
clean up resources when you're done with the store.
"""
__slots__ = (
@@ -612,7 +710,10 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
"supports_pipeline",
"index_config",
"embeddings",
"_ttl_sweeper_thread",
"_ttl_stop_event",
)
supports_ttl: bool = True
def __init__(
self,
@@ -637,6 +738,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
else:
self.embeddings = None
self.ttl_config = ttl
self._ttl_sweeper_thread: Optional[threading.Thread] = None
self._ttl_stop_event = threading.Event()
@classmethod
@contextmanager
@@ -647,6 +750,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
ttl: Optional[TTLConfig] = None,
) -> Iterator["PostgresStore"]:
"""Create a new PostgresStore instance from a connection string.
@@ -678,16 +782,123 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool, index=index)
yield cls(conn=pool, index=index, ttl=ttl)
else:
with Connection.connect(
conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row
) as conn:
if pipeline:
with conn.pipeline() as pipe:
yield cls(conn, pipe=pipe, index=index)
yield cls(conn, pipe=pipe, index=index, ttl=ttl)
else:
yield cls(conn, index=index)
yield cls(conn, index=index, ttl=ttl)
def sweep_ttl(self) -> int:
"""Delete expired store items based on TTL.
Returns:
int: The number of deleted items.
"""
with self._cursor() as cur:
cur.execute(
"""
DELETE FROM store
WHERE expires_at IS NOT NULL AND expires_at < NOW()
"""
)
deleted_count = cur.rowcount
return deleted_count
def start_ttl_sweeper(
self, sweep_interval_minutes: Optional[int] = None
) -> concurrent.futures.Future[None]:
"""Periodically delete expired store items based on TTL.
Returns:
Future that can be waited on or cancelled.
"""
if not self.ttl_config:
future: concurrent.futures.Future[None] = concurrent.futures.Future()
future.set_result(None)
return future
if self._ttl_sweeper_thread and self._ttl_sweeper_thread.is_alive():
logger.info("TTL sweeper thread is already running")
# Return a future that can be used to cancel the existing thread
future = concurrent.futures.Future()
future.add_done_callback(
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
)
return future
self._ttl_stop_event.clear()
interval = float(
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
)
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
future = concurrent.futures.Future()
def _sweep_loop() -> None:
try:
while not self._ttl_stop_event.is_set():
if self._ttl_stop_event.wait(interval * 60):
break
try:
expired_items = self.sweep_ttl()
if expired_items > 0:
logger.info(f"Store swept {expired_items} expired items")
except Exception as exc:
logger.exception(
"Store TTL sweep iteration failed", exc_info=exc
)
future.set_result(None)
except Exception as exc:
future.set_exception(exc)
thread = threading.Thread(target=_sweep_loop, daemon=True, name="ttl-sweeper")
self._ttl_sweeper_thread = thread
thread.start()
future.add_done_callback(
lambda f: self._ttl_stop_event.set() if f.cancelled() else None
)
return future
def stop_ttl_sweeper(self, timeout: Optional[float] = None) -> bool:
"""Stop the TTL sweeper thread if it's running.
Args:
timeout: Maximum time to wait for the thread to stop, in seconds.
If None, wait indefinitely.
Returns:
bool: True if the thread was successfully stopped or wasn't running,
False if the timeout was reached before the thread stopped.
"""
if not self._ttl_sweeper_thread or not self._ttl_sweeper_thread.is_alive():
return True
logger.info("Stopping TTL sweeper thread")
self._ttl_stop_event.set()
self._ttl_sweeper_thread.join(timeout)
success = not self._ttl_sweeper_thread.is_alive()
if success:
self._ttl_sweeper_thread = None
logger.info("TTL sweeper thread stopped")
else:
logger.warning("Timed out waiting for TTL sweeper thread to stop")
return success
def __del__(self) -> None:
"""Ensure the TTL sweeper thread is stopped when the object is garbage collected."""
if hasattr(self, "_ttl_stop_event") and hasattr(self, "_ttl_sweeper_thread"):
self.stop_ttl_sweeper(timeout=0.1)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -886,8 +1097,14 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
with self._cursor() as cur:
version = _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
try:
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
except Exception as e:
logger.error(
f"Failed to apply migration {v}.\nSql={sql}\nError={e}"
)
raise
if self.index_config:
version = _get_version(cur, table="vector_migrations")
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint-postgres"
version = "2.0.16"
version = "2.0.17"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
@@ -26,6 +26,9 @@ from tests.conftest import (
CharacterEmbeddings,
)
TTL_SECONDS = 6
TTL_MINUTES = TTL_SECONDS / 60
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
async def store(request) -> AsyncIterator[AsyncPostgresStore]:
@@ -42,28 +45,52 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]:
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
ttl_config = {
"default_ttl": TTL_MINUTES,
"refresh_on_read": True,
"sweep_interval_minutes": TTL_MINUTES / 2,
}
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
async with AsyncPostgresStore.from_conn_string(
conn_string, ttl=ttl_config
) as store:
store.MIGRATIONS = [
(
mig.replace(
"ADD COLUMN ttl_minutes INT;", "ADD COLUMN ttl_minutes FLOAT;"
)
if isinstance(mig, str)
else mig
)
for mig in store.MIGRATIONS
]
await store.setup()
if request.param == "pipe":
async with AsyncPostgresStore.from_conn_string(
conn_string, pipeline=True
conn_string, pipeline=True, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
elif request.param == "pool":
async with AsyncPostgresStore.from_conn_string(
conn_string, pool_config={"min_size": 1, "max_size": 10}
conn_string, pool_config={"min_size": 1, "max_size": 10}, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
else: # default
async with AsyncPostgresStore.from_conn_string(conn_string) as store:
async with AsyncPostgresStore.from_conn_string(
conn_string, ttl=ttl_config
) as store:
await store.start_ttl_sweeper()
yield store
await store.stop_ttl_sweeper()
finally:
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
@@ -635,3 +662,28 @@ async def test_search_sorting(
assert len(set(r.key for r in results)) == 10
assert results[0].key == "M"
assert results[0].score > results[1].score
async def test_store_ttl(store):
# Assumes a TTL of 1 minute = 60 seconds
ns = ("foo",)
await store.start_ttl_sweeper()
await store.aput(
ns,
key="item1",
value={"foo": "bar"},
ttl=TTL_MINUTES, # type: ignore
)
await asyncio.sleep(TTL_SECONDS - 2)
res = await store.aget(ns, key="item1", refresh_ttl=True)
assert res is not None
await asyncio.sleep(TTL_SECONDS - 2)
results = await store.asearch(ns, query="foo", refresh_ttl=True)
assert len(results) == 1
await asyncio.sleep(TTL_SECONDS - 2)
res = await store.aget(ns, key="item1", refresh_ttl=False)
assert res is not None
await asyncio.sleep(TTL_SECONDS - 1)
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
results = await store.asearch(ns, query="bar", refresh_ttl=False)
assert len(results) == 0
+193 -125
View File
@@ -1,6 +1,7 @@
# type: ignore
import re
import time
from contextlib import contextmanager
from typing import Any, Optional
from uuid import uuid4
@@ -24,6 +25,9 @@ from tests.conftest import (
CharacterEmbeddings,
)
TTL_SECONDS = 6
TTL_MINUTES = TTL_SECONDS / 60
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
def store(request) -> PostgresStore:
@@ -32,29 +36,58 @@ def store(request) -> PostgresStore:
uri_base = "/".join(uri_parts[:-1])
query_params = ""
if "?" in uri_parts[-1]:
db_name, query_params = uri_parts[-1].split("?", 1)
_, query_params = uri_parts[-1].split("?", 1)
query_params = "?" + query_params
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
ttl_config = {
"default_ttl": TTL_MINUTES,
"refresh_on_read": True,
"sweep_interval_minutes": TTL_MINUTES / 2,
}
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with PostgresStore.from_conn_string(conn_string) as store:
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
store.MIGRATIONS = [
(
mig.replace(
"ADD COLUMN ttl_minutes INT;", "ADD COLUMN ttl_minutes FLOAT;"
)
if isinstance(mig, str)
else mig
)
for mig in store.MIGRATIONS
]
store.setup()
if request.param == "pipe":
with PostgresStore.from_conn_string(conn_string, pipeline=True) as store:
with PostgresStore.from_conn_string(
conn_string,
pipeline=True,
ttl=ttl_config,
) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
elif request.param == "pool":
with PostgresStore.from_conn_string(
conn_string, pool_config={"min_size": 1, "max_size": 10}
conn_string,
pool_config={"min_size": 1, "max_size": 10},
ttl=ttl_config,
) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
else: # default
with PostgresStore.from_conn_string(conn_string) as store:
with PostgresStore.from_conn_string(conn_string, ttl=ttl_config) as store:
store.start_ttl_sweeper()
yield store
store.stop_ttl_sweeper()
finally:
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@@ -220,134 +253,127 @@ def test_batch_list_namespaces_ops(store: PostgresStore) -> None:
assert all(ns[-1] == "public" for ns in results[2])
class TestPostgresStore:
@pytest.fixture(autouse=True)
def setup(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
store.setup()
def test_basic_store_ops(store) -> None:
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
def test_basic_store_ops(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
store.put(namespace, item_id, item_value)
item = store.get(namespace, item_id)
store.put(namespace, item_id, item_value)
item = store.get(namespace, item_id)
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
# Test update
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
store.put(namespace, item_id, updated_value)
updated_item = store.get(namespace, item_id)
# Test update
updated_value = {"title": "Updated Document", "content": "Hello, Updated!"}
store.put(namespace, item_id, updated_value)
updated_item = store.get(namespace, item_id)
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
# Test get from non-existent namespace
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
# Test get from non-existent namespace
different_namespace = ("test", "other_documents")
item_in_different_namespace = store.get(different_namespace, item_id)
assert item_in_different_namespace is None
# Test delete
store.delete(namespace, item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
# Test delete
store.delete(namespace, item_id)
deleted_item = store.get(namespace, item_id)
assert deleted_item is None
def test_list_namespaces(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
# Create test data with various namespaces
test_namespaces = [
("test", "documents", "public"),
("test", "documents", "private"),
("test", "images", "public"),
("test", "images", "private"),
("prod", "documents", "public"),
("prod", "documents", "private"),
]
def test_list_namespaces(store) -> None:
# Create test data with various namespaces
test_namespaces = [
("test", "documents", "public"),
("test", "documents", "private"),
("test", "images", "public"),
("test", "images", "private"),
("prod", "documents", "public"),
("prod", "documents", "private"),
]
# Insert test data
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
# Insert test data
for namespace in test_namespaces:
store.put(namespace, "dummy", {"content": "dummy"})
# Test listing with various filters
all_namespaces = store.list_namespaces()
assert len(all_namespaces) == len(test_namespaces)
# Test listing with various filters
all_namespaces = store.list_namespaces()
assert len(all_namespaces) == len(test_namespaces)
# Test prefix filtering
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
assert len(test_prefix_namespaces) == 4
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
# Test prefix filtering
test_prefix_namespaces = store.list_namespaces(prefix=["test"])
assert len(test_prefix_namespaces) == 4
assert all(ns[0] == "test" for ns in test_prefix_namespaces)
# Test suffix filtering
public_namespaces = store.list_namespaces(suffix=["public"])
assert len(public_namespaces) == 3
assert all(ns[-1] == "public" for ns in public_namespaces)
# Test suffix filtering
public_namespaces = store.list_namespaces(suffix=["public"])
assert len(public_namespaces) == 3
assert all(ns[-1] == "public" for ns in public_namespaces)
# Test max depth
depth_2_namespaces = store.list_namespaces(max_depth=2)
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
# Test max depth
depth_2_namespaces = store.list_namespaces(max_depth=2)
assert all(len(ns) <= 2 for ns in depth_2_namespaces)
# Test pagination
paginated_namespaces = store.list_namespaces(limit=3)
assert len(paginated_namespaces) == 3
# Test pagination
paginated_namespaces = store.list_namespaces(limit=3)
assert len(paginated_namespaces) == 3
# Cleanup
for namespace in test_namespaces:
store.delete(namespace, "dummy")
# Cleanup
for namespace in test_namespaces:
store.delete(namespace, "dummy")
def test_search(self) -> None:
with PostgresStore.from_conn_string(DEFAULT_URI) as store:
# Create test data
test_data = [
(
("test", "docs"),
"doc1",
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
),
(
("test", "docs"),
"doc2",
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
),
(
("test", "images"),
"img1",
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
),
]
for namespace, key, value in test_data:
store.put(namespace, key, value)
def test_search(store) -> None:
# Create test data
test_data = [
(
("test", "docs"),
"doc1",
{"title": "First Doc", "author": "Alice", "tags": ["important"]},
),
(
("test", "docs"),
"doc2",
{"title": "Second Doc", "author": "Bob", "tags": ["draft"]},
),
(
("test", "images"),
"img1",
{"title": "Image 1", "author": "Alice", "tags": ["final"]},
),
]
# Test basic search
all_items = store.search(["test"])
assert len(all_items) == 3
for namespace, key, value in test_data:
store.put(namespace, key, value)
# Test namespace filtering
docs_items = store.search(["test", "docs"])
assert len(docs_items) == 2
assert all(item.namespace == ("test", "docs") for item in docs_items)
# Test basic search
all_items = store.search(["test"])
assert len(all_items) == 3
# Test value filtering
alice_items = store.search(["test"], filter={"author": "Alice"})
assert len(alice_items) == 2
assert all(item.value["author"] == "Alice" for item in alice_items)
# Test namespace filtering
docs_items = store.search(["test", "docs"])
assert len(docs_items) == 2
assert all(item.namespace == ("test", "docs") for item in docs_items)
# Test pagination
paginated_items = store.search(["test"], limit=2)
assert len(paginated_items) == 2
# Test value filtering
alice_items = store.search(["test"], filter={"author": "Alice"})
assert len(alice_items) == 2
assert all(item.value["author"] == "Alice" for item in alice_items)
offset_items = store.search(["test"], offset=2)
assert len(offset_items) == 1
# Test pagination
paginated_items = store.search(["test"], limit=2)
assert len(paginated_items) == 2
# Cleanup
for namespace, key, _ in test_data:
store.delete(namespace, key)
offset_items = store.search(["test"], offset=2)
assert len(offset_items) == 1
# Cleanup
for namespace, key, _ in test_data:
store.delete(namespace, key)
@contextmanager
@@ -356,6 +382,7 @@ def _create_vector_store(
distance_type: str,
fake_embeddings: Embeddings,
text_fields: Optional[list[str]] = None,
enable_ttl: bool = True,
) -> PostgresStore:
"""Create a store with vector search enabled."""
database = f"test_{uuid4().hex[:16]}"
@@ -385,6 +412,7 @@ def _create_vector_store(
with PostgresStore.from_conn_string(
conn_string,
index=index_config,
ttl={"default_ttl": 2, "refresh_on_read": True} if enable_ttl else None,
) as store:
store.setup()
yield store
@@ -393,15 +421,19 @@ def _create_vector_store(
conn.execute(f"DROP DATABASE {database}")
_vector_params = [
(vector_type, distance_type, True)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
]
_vector_params += [(*_vector_params[-1][:2], False)]
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
],
params=_vector_params,
ids=lambda p: f"{p[0]}_{p[1]}",
)
def vector_store(
@@ -409,8 +441,10 @@ def vector_store(
fake_embeddings: Embeddings,
) -> PostgresStore:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
vector_type, distance_type, enable_ttl = request.param
with _create_vector_store(
vector_type, distance_type, fake_embeddings, enable_ttl=enable_ttl
) as store:
yield store
@@ -474,7 +508,10 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
assert not any(r.key == "doc4" for r in results_new)
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
@pytest.mark.parametrize("refresh_ttl", [True, False])
def test_vector_search_with_filters(
vector_store: PostgresStore, refresh_ttl: bool
) -> None:
"""Test combining vector search with filters."""
# Insert test documents
docs = [
@@ -487,16 +524,23 @@ def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
for key, value in docs:
vector_store.put(("test",), key, value)
results = vector_store.search(("test",), query="apple", filter={"color": "red"})
results = vector_store.search(
("test",), query="apple", filter={"color": "red"}, refresh_ttl=refresh_ttl
)
assert len(results) == 2
assert results[0].key == "doc1"
results = vector_store.search(("test",), query="car", filter={"color": "red"})
results = vector_store.search(
("test",), query="car", filter={"color": "red"}, refresh_ttl=refresh_ttl
)
assert len(results) == 2
assert results[0].key == "doc2"
results = vector_store.search(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
("test",),
query="bbbbluuu",
filter={"score": {"$gt": 3.2}},
refresh_ttl=refresh_ttl,
)
assert len(results) == 3
assert results[0].key == "doc4"
@@ -688,7 +732,7 @@ def test_embed_with_path_operation_config(
store.put(("test",), "doc5", doc5, index=False)
results = store.search(("test",))
assert len(results) == 3
assert all(r.score is None for r in results)
assert all(r.score is None for r in results), f"{results}"
assert any(r.key == "doc5" for r in results)
results = store.search(("test",), query="hhh")
@@ -790,3 +834,27 @@ def test_nonnull_migrations() -> None:
for migration in PostgresStore.MIGRATIONS:
statement = _leading_comment_remover.sub("", migration).split()[0]
assert statement.strip()
def test_store_ttl(store):
# Assumes a TTL of 1 minute = 60 seconds
ns = ("foo",)
store.put(
ns,
key="item1",
value={"foo": "bar"},
ttl=TTL_MINUTES, # type: ignore
)
time.sleep(TTL_SECONDS - 2)
res = store.get(ns, key="item1", refresh_ttl=True)
assert res is not None
time.sleep(TTL_SECONDS - 2)
results = store.search(ns, query="foo", refresh_ttl=True)
assert len(results) == 1
time.sleep(TTL_SECONDS - 2)
res = store.get(ns, key="item1", refresh_ttl=False)
assert res is not None
time.sleep(TTL_SECONDS - 1)
# Now has been (TTL_SECONDS-2)*2 > TTL_SECONDS + TTL_SECONDS/2
res = store.search(ns, query="bar", refresh_ttl=False)
assert len(res) == 0
@@ -537,6 +537,12 @@ class TTLConfig(TypedDict, total=False):
The expiration timer refreshes on both read and write operations.
Defaults to None (no expiration).
"""
sweep_interval_minutes: Optional[int]
"""Interval in minutes between TTL sweep operations.
If provided, the store will periodically delete expired items based on TTL.
Defaults to None (no sweeping).
"""
class IndexConfig(TypedDict, total=False):
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.19"
version = "2.0.20"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
+4 -1
View File
@@ -1,4 +1,4 @@
.PHONY: test lint format test-integration
.PHONY: test lint format test-integration update-schema
######################
# TESTING AND COVERAGE
@@ -31,3 +31,6 @@ lint lint_diff lint_package lint_tests:
format format_diff:
poetry run ruff format $(PYTHON_FILES)
poetry run ruff check --select I --fix $(PYTHON_FILES)
update-schema:
poetry run python generate_schema.py
+6
View File
@@ -27,6 +27,12 @@ class TTLConfig(TypedDict, total=False):
If provided, all new items will have this TTL unless explicitly overridden.
If omitted, items will have no TTL by default.
"""
sweep_interval_minutes: Optional[int]
"""Optional. Interval in minutes between TTL sweep iterations.
If provided, the store will periodically delete expired items based on the TTL.
If omitted, no automatic sweeping will occur.
"""
class IndexConfig(TypedDict, total=False):
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-cli"
version = "0.1.76"
version = "0.1.77"
description = "CLI for interacting with LangGraph API"
authors = []
license = "MIT"
+10
View File
@@ -459,6 +459,16 @@
},
"refresh_on_read": {
"type": "boolean"
},
"sweep_interval_minutes": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
}
},
"required": []
+10
View File
@@ -459,6 +459,16 @@
},
"refresh_on_read": {
"type": "boolean"
},
"sweep_interval_minutes": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
}
},
"required": []