diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 5ffcaaf5a..a5e7604cd 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -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,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con "supports_pipeline", "index_config", "embeddings", - "supports_ttl", + "ttl_config", + "_ttl_sweeper_task", + "_ttl_stop_event", ) supports_ttl: bool = True @@ -128,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( @@ -143,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 @@ -169,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. @@ -200,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. @@ -276,30 +290,100 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con 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", - ) + ) -> asyncio.Task[None]: + """Periodically delete expired store items based on TTL. - while True: - await asyncio.sleep(sweep_interval_minutes_ * 60) + 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: - 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) + 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, diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 954b8e88d..d88ae19ed 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -1,8 +1,8 @@ import asyncio +import concurrent.futures import json import logging import threading -import time from collections import defaultdict from collections.abc import Iterable, Iterator, Sequence from contextlib import contextmanager @@ -81,7 +81,8 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS store_prefix_idx ON store USING btree (p 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; @@ -253,7 +254,7 @@ class BasePostgresStore(Generic[C]): results = [] for namespace, items in namespace_groups.items(): - _, keys = zip(*items, strict=True) + _, keys = zip(*items) this_refresh_ttls = refresh_ttls[namespace] query = """ @@ -292,7 +293,7 @@ class BasePostgresStore(Generic[C]): put_ops: Sequence[tuple[int, PutOp]], ) -> tuple[ list[tuple[str, Sequence]], - tuple[str, Sequence[tuple[str, str, str, str]]] | None, + Optional[tuple[str, Sequence[tuple[str, str, str, str]]]], ]: dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {} for _, op in put_ops: @@ -319,7 +320,9 @@ class BasePostgresStore(Generic[C]): ) params = (_namespace_to_text(namespace), *keys) queries.append((query, params)) - embedding_request: tuple[str, Sequence[tuple[str, str, str, str]]] | None = None + embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = ( + None + ) if inserts: values = [] insertion_params = [] @@ -400,7 +403,7 @@ class BasePostgresStore(Generic[C]): self, search_ops: Sequence[tuple[int, SearchOp]], ) -> tuple[ - list[tuple[str, list[None | str | list[float]]]], # queries, params + list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params list[tuple[int, str]], # idx, query_text pairs to embed ]: """ @@ -412,7 +415,6 @@ class BasePostgresStore(Generic[C]): queries = [] embedding_requests = [] - for idx, (_, op) in enumerate(search_ops): filter_params = [] filter_clauses = [] @@ -430,7 +432,7 @@ class BasePostgresStore(Generic[C]): filter_params.extend([key, orjson.dumps(value).decode("utf-8")]) ns_condition = "TRUE" - ns_param: Sequence[str] | None = None + 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)}%",) @@ -512,7 +514,7 @@ class BasePostgresStore(Generic[C]): else: base_query = f""" - SELECT store.prefix, store.key, store.value, store.created_at, store.updated_at, 0 AS score + 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 @@ -625,7 +627,7 @@ class BasePostgresStore(Generic[C]): class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): - """Postgres-backed store with optional vector search using pgvector. + """Postgres-backed store with oktional vector search using pgvector. !!! example "Examples" Basic setup and usage: @@ -694,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__ = ( @@ -703,7 +710,8 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): "supports_pipeline", "index_config", "embeddings", - "supports_ttl", + "_ttl_sweeper_thread", + "_ttl_stop_event", ) supports_ttl: bool = True @@ -730,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 @@ -740,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. @@ -771,16 +782,16 @@ 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. @@ -798,30 +809,96 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): 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", - ) + def start_ttl_sweeper( + self, sweep_interval_minutes: Optional[int] = None + ) -> concurrent.futures.Future[None]: + """Periodically delete expired store items based on TTL. - while True: - time.sleep(sweep_interval_minutes_ * 60) + 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: - expired_items = self.sweep_ttl() - if expired_items > 0: - logger.info(f"Store swept {expired_items} expired items") + 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: - logger.exception("Store TTL sweep iteration failed", exc_info=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]]: @@ -1020,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") diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 068ec1502..09502403d 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -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 diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index 50d697962..4ee37484e 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -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 diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index fff90cc56..de914a177 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -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): diff --git a/libs/cli/langgraph_cli/config.py b/libs/cli/langgraph_cli/config.py index 6200b1ad7..011432147 100644 --- a/libs/cli/langgraph_cli/config.py +++ b/libs/cli/langgraph_cli/config.py @@ -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):