WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 09:02:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add sweeper

Browse Source
This commit is contained in:
William Fu-Hinthorn
2025-03-13 18:55:16 -07:00
parent 98ebc45f31
commit 06ca07432d
2 changed files with 264 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/store/postgres/aio.py
+45
View File
@@ -115,7 +115,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
"supports_pipeline",
"index_config",
"embeddings",
"supports_ttl",
)
supports_ttl: bool = True
def __init__(
self,
@@ -256,6 +258,49 @@ 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
) -> None:
"""Periodically delete expired store items based on TTL."""
if not self.ttl_config:
return
sweep_interval_minutes_ = float(
cast(
float,
sweep_interval_minutes
or self.ttl_config.get("sweep_interval_minutes")
or 5,
)
)
logger.info(
f"Starting store TTL sweeper with interval {sweep_interval_minutes_} minutes",
)
while True:
await asyncio.sleep(sweep_interval_minutes_ * 60)
try:
expired_items = await 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)
async def _execute_batch(
self,
grouped_ops: dict,
libs/checkpoint-postgres/langgraph/store/postgres/base.py
+219 -85
View File
@@ -2,6 +2,7 @@ import asyncio
import json
import logging
import threading
import time
from collections import defaultdict
from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
@@ -74,6 +75,16 @@ 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 +236,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})
_, keys = zip(*items, strict=True)
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(
@@ -246,9 +292,8 @@ class BasePostgresStore(Generic[C]):
put_ops: Sequence[tuple[int, PutOp]],
) -> tuple[
list[tuple[str, Sequence]],
Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
tuple[str, Sequence[tuple[str, str, str, str]]] | None,
]:
# Last-write wins
dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
for _, op in put_ops:
dedupped_ops[(op.namespace, op.key)] = op
@@ -274,23 +319,32 @@ class BasePostgresStore(Generic[C]):
)
params = (_namespace_to_text(namespace), *keys)
queries.append((query, params))
embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
None
)
embedding_request: tuple[str, Sequence[tuple[str, str, str, str]]] | None = None
if inserts:
values = []
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 +358,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 +373,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))
@@ -344,95 +400,109 @@ class BasePostgresStore(Generic[C]):
self,
search_ops: Sequence[tuple[int, SearchOp]],
) -> tuple[
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
list[tuple[str, list[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: Sequence[str] | None = 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 +510,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, 0 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
@@ -612,7 +703,9 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
"supports_pipeline",
"index_config",
"embeddings",
"supports_ttl",
)
supports_ttl: bool = True
def __init__(
self,
@@ -689,6 +782,47 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
else:
yield cls(conn, index=index)
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) -> None:
"""Periodically delete expired store items based on TTL."""
if not self.ttl_config:
return
sweep_interval_minutes_ = float(
cast(
float,
sweep_interval_minutes
or self.ttl_config.get("sweep_interval_minutes")
or 5,
)
)
logger.info(
f"Starting store TTL sweeper with interval {sweep_interval_minutes_} minutes",
)
while True:
time.sleep(sweep_interval_minutes_ * 60)
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)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
"""Create a database cursor as a context manager.