diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index a8d434360..0615f0ec7 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -20,7 +20,7 @@ from langgraph.store.base import ( ) from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.postgres.base import ( - _PLACEHOLDER, + PLACEHOLDER, BasePostgresStore, PoolConfig, PostgresIndexConfig, @@ -360,7 +360,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con for (idx, _), vector in zip(embedding_requests, vectors): _paramslist = queries[idx][1] for i in range(len(_paramslist)): - if _paramslist[i] is _PLACEHOLDER: + if _paramslist[i] is PLACEHOLDER: _paramslist[i] = vector for (idx, _), (query, params) in zip(search_ops, queries): diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index a3acc4744..4ee98926e 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -370,7 +370,7 @@ class BasePostgresStore(Generic[C]): if op.query and self.index_config: embedding_requests.append((idx, op.query)) - score_operator, post_operator = _get_distance_operator(self) + score_operator, post_operator = get_distance_operator(self) vector_type = ( cast(PostgresIndexConfig, self.index_config) .get("ann_index_config", {}) @@ -430,10 +430,10 @@ class BasePostgresStore(Generic[C]): OFFSET %s """ params = [ - _PLACEHOLDER, # Vector placeholder + PLACEHOLDER, # Vector placeholder *ns_args, *filter_params, - _PLACEHOLDER, + PLACEHOLDER, expanded_limit, op.limit, op.offset, @@ -828,7 +828,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): for (idx, _), embedding in zip(embedding_requests, embeddings): _paramslist = queries[idx][1] for i in range(len(_paramslist)): - if _paramslist[i] is _PLACEHOLDER: + if _paramslist[i] is PLACEHOLDER: _paramslist[i] = embedding for (idx, _), (query, params) in zip(search_ops, queries): @@ -1055,7 +1055,7 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]: return tuple(namespace.split(".")) -def _get_distance_operator(store: Any) -> tuple[str, str]: +def get_distance_operator(store: Any) -> tuple[str, str]: """Get the distance operator and score expression based on config.""" # Note: Today, we are not using ANN indices due to restrictions # on PGVector's support for mixing vector and non-vector filters @@ -1121,4 +1121,4 @@ def _ensure_index_config( return embeddings, index_config -_PLACEHOLDER = object() +PLACEHOLDER = object() diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index 19fc62d8d..35686caa0 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -166,6 +166,13 @@ class GetOp(NamedTuple): "doc456" # For a document ``` """ + refresh_ttl: bool = True + """Whether to refresh TTLs for the returned item. + + If no TTL was specified for the original item(s), + or if TTL support is not enabled for your adapter, + this argument is ignored. + """ class SearchOp(NamedTuple): @@ -260,6 +267,13 @@ class SearchOp(NamedTuple): - "technical documentation about REST APIs" - "machine learning papers from 2023" """ + refresh_ttl: bool = True + """Whether to refresh TTLs for the returned item. + + If no TTL was specified for the original item(s), + or if TTL support is not enabled for your adapter, + this argument is ignored. + """ # Type representing a namespace path that can include wildcards @@ -463,6 +477,15 @@ class PutOp(NamedTuple): ] ``` """ + ttl: Optional[float] = None + """Controls the TTL (time-to-live) for the item in minutes. + + If provided, and if the store you are using supports this feature, the item + will expire this many minutes after it was last accessed. The expiration timer + refreshes on both read operations (get/search) and write operations (put/update). + When the TTL expires, the item will be scheduled for deletion on a best-effort basis. + Defaults to None (no expiration). + """ Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp] @@ -612,8 +635,13 @@ class BaseStore(ABC): by providing an `index` configuration at creation time. Without this configuration, semantic search is disabled and any `index` arguments to storage operations will have no effect. + + Similarly, TTL (time-to-live) support is disabled by default. + Subclasses must explicitly set `supports_ttl = True` to enable this feature. """ + supports_ttl = False + __slots__ = ("__weakref__",) @abstractmethod @@ -640,17 +668,21 @@ class BaseStore(ABC): The order of results matches the order of input operations. """ - def get(self, namespace: tuple[str, ...], key: str) -> Optional[Item]: + def get( + self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + ) -> Optional[Item]: """Retrieve a single item. Args: namespace: Hierarchical path for the item. key: Unique identifier within the namespace. + refresh_ttl: Whether to refresh TTLs for the returned item. + If no TTL is specified, this argument is ignored. Returns: The retrieved item or None if not found. """ - return self.batch([GetOp(namespace, key)])[0] + return self.batch([GetOp(namespace, key, refresh_ttl)])[0] def search( self, @@ -661,6 +693,7 @@ class BaseStore(ABC): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, + refresh_ttl: bool = True, ) -> list[SearchItem]: """Search for items within a namespace prefix. @@ -670,6 +703,8 @@ class BaseStore(ABC): filter: Key-value pairs to filter results. limit: Maximum number of items to return. offset: Number of items to skip before returning results. + refresh_ttl: Whether to refresh TTLs for the returned items. + If no TTL is specified, this argument is ignored. Returns: List of items matching the search criteria. @@ -707,7 +742,9 @@ class BaseStore(ABC): Note: Natural language search support depends on your store implementation and requires proper embedding configuration. """ - return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0] + return self.batch( + [SearchOp(namespace_prefix, filter, limit, offset, query, refresh_ttl)] + )[0] def put( self, @@ -715,6 +752,8 @@ class BaseStore(ABC): key: str, value: dict[str, Any], index: Optional[Union[Literal[False], list[str]]] = None, + *, + ttl: Optional[float] = None, ) -> None: """Store or update an item in the store. @@ -735,12 +774,20 @@ class BaseStore(ABC): - Nested fields: "metadata.title" - Array access: "chapters[*].content" (each indexed separately) - Specific indices: "authors[0].name" + ttl: Time to live in minutes. Support for this argument depends on your store adapter. + If specified, the item will expire after this many minutes from when it was last accessed. + None means no expiration. Expired runs will be deleted opportunistically. + By default, the expiration timer refreshes on both read operations (get/search) + and write operations (put/update), whenever the item is included in the operation. Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, the `index` parameter will be ignored. + Similarly, TTL support depends on the specific store implementation. + Some implementations may not support expiration of items. + ???+ example "Examples" Store item. Indexing depends on how you configure the store. ```python @@ -759,7 +806,12 @@ class BaseStore(ABC): ``` """ _validate_namespace(namespace) - self.batch([PutOp(namespace, key, value, index=index)]) + if ttl is not None and not self.supports_ttl: + raise NotImplementedError( + f"TTL is not supported by {self.__class__.__name__}. " + f"Use a store implementation that supports TTL or set ttl=None." + ) + self.batch([PutOp(namespace, key, value, index=index, ttl=ttl)]) def delete(self, namespace: tuple[str, ...], key: str) -> None: """Delete an item. @@ -768,7 +820,7 @@ class BaseStore(ABC): namespace: Hierarchical path for the item. key: Unique identifier within the namespace. """ - self.batch([PutOp(namespace, key, None)]) + self.batch([PutOp(namespace, key, None, ttl=None)]) def list_namespaces( self, @@ -823,7 +875,9 @@ class BaseStore(ABC): ) return self.batch([op])[0] - async def aget(self, namespace: tuple[str, ...], key: str) -> Optional[Item]: + async def aget( + self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + ) -> Optional[Item]: """Asynchronously retrieve a single item. Args: @@ -833,7 +887,7 @@ class BaseStore(ABC): Returns: The retrieved item or None if not found. """ - return (await self.abatch([GetOp(namespace, key)]))[0] + return (await self.abatch([GetOp(namespace, key, refresh_ttl)]))[0] async def asearch( self, @@ -844,6 +898,7 @@ class BaseStore(ABC): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, + refresh_ttl: bool = True, ) -> list[SearchItem]: """Asynchronously search for items within a namespace prefix. @@ -853,6 +908,9 @@ class BaseStore(ABC): filter: Key-value pairs to filter results. limit: Maximum number of items to return. offset: Number of items to skip before returning results. + refresh_ttl: Whether to refresh TTLs for the returned items. + Defaults to True. If no TTL is specified, this argument + is ignored. Returns: List of items matching the search criteria. @@ -892,7 +950,7 @@ class BaseStore(ABC): """ return ( await self.abatch( - [SearchOp(namespace_prefix, filter, limit, offset, query)] + [SearchOp(namespace_prefix, filter, limit, offset, query, refresh_ttl)] ) )[0] @@ -902,6 +960,8 @@ class BaseStore(ABC): key: str, value: dict[str, Any], index: Optional[Union[Literal[False], list[str]]] = None, + *, + ttl: Optional[float] = None, ) -> None: """Asynchronously store or update an item in the store. @@ -922,12 +982,20 @@ class BaseStore(ABC): - Nested fields: "metadata.title" - Array access: "chapters[*].content" (each indexed separately) - Specific indices: "authors[0].name" + ttl: Time to live in minutes. Support for this argument depends on your store adapter. + If specified, the item will expire after this many minutes from when it was last accessed. + None means no expiration. Expired runs will be deleted opportunistically. + By default, the expiration timer refreshes on both read operations (get/search) + and write operations (put/update), whenever the item is included in the operation. Note: Indexing support depends on your store implementation. If you do not initialize the store with indexing capabilities, the `index` parameter will be ignored. + Similarly, TTL support depends on the specific store implementation. + Some implementations may not support expiration of items. + ???+ example "Examples" Store item. Indexing depends on how you configure the store. ```python @@ -954,7 +1022,12 @@ class BaseStore(ABC): ``` """ _validate_namespace(namespace) - await self.abatch([PutOp(namespace, key, value, index=index)]) + if ttl is not None and not self.supports_ttl: + raise NotImplementedError( + f"TTL is not supported by {self.__class__.__name__}. " + f"Use a store implementation that supports TTL or set ttl=None." + ) + await self.abatch([PutOp(namespace, key, value, index=index, ttl=ttl)]) async def adelete(self, namespace: tuple[str, ...], key: str) -> None: """Asynchronously delete an item. diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index debf187b3..d0bba1efe 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -65,13 +65,11 @@ class AsyncBatchedBaseStore(BaseStore): pass async def aget( - self, - namespace: tuple[str, ...], - key: str, + self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True ) -> Optional[Item]: assert not self._task.done() fut = self._loop.create_future() - self._aqueue.put_nowait((fut, GetOp(namespace, key))) + self._aqueue.put_nowait((fut, GetOp(namespace, key, refresh_ttl=refresh_ttl))) return await fut async def asearch( @@ -83,11 +81,22 @@ class AsyncBatchedBaseStore(BaseStore): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, + refresh_ttl: bool = True, ) -> list[SearchItem]: assert not self._task.done() fut = self._loop.create_future() self._aqueue.put_nowait( - (fut, SearchOp(namespace_prefix, filter, limit, offset, query)) + ( + fut, + SearchOp( + namespace_prefix, + filter, + limit, + offset, + query, + refresh_ttl=refresh_ttl, + ), + ) ) return await fut @@ -97,11 +106,13 @@ class AsyncBatchedBaseStore(BaseStore): key: str, value: dict[str, Any], index: Optional[Union[Literal[False], list[str]]] = None, + *, + ttl: Optional[float] = None, ) -> None: assert not self._task.done() _validate_namespace(namespace) fut = self._loop.create_future() - self._aqueue.put_nowait((fut, PutOp(namespace, key, value, index))) + self._aqueue.put_nowait((fut, PutOp(namespace, key, value, index, ttl=ttl))) return await fut async def adelete( @@ -146,12 +157,10 @@ class AsyncBatchedBaseStore(BaseStore): @_check_loop def get( - self, - namespace: tuple[str, ...], - key: str, + self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True ) -> Optional[Item]: return asyncio.run_coroutine_threadsafe( - self.aget(namespace, key=key), self._loop + self.aget(namespace, key=key, refresh_ttl=refresh_ttl), self._loop ).result() @_check_loop @@ -164,10 +173,16 @@ class AsyncBatchedBaseStore(BaseStore): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, + refresh_ttl: bool = True, ) -> list[SearchItem]: return asyncio.run_coroutine_threadsafe( self.asearch( - namespace_prefix, query=query, filter=filter, limit=limit, offset=offset + namespace_prefix, + query=query, + filter=filter, + limit=limit, + offset=offset, + refresh_ttl=refresh_ttl, ), self._loop, ).result() @@ -179,10 +194,12 @@ class AsyncBatchedBaseStore(BaseStore): key: str, value: dict[str, Any], index: Optional[Union[Literal[False], list[str]]] = None, + *, + ttl: Optional[float] = None, ) -> None: _validate_namespace(namespace) asyncio.run_coroutine_threadsafe( - self.aput(namespace, key=key, value=value, index=index), self._loop + self.aput(namespace, key=key, value=value, index=index, ttl=ttl), self._loop ).result() @_check_loop diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index 9610d8fc5..07a092f2c 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "2.0.16" +version = "2.0.17" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" diff --git a/libs/checkpoint/tests/test_store.py b/libs/checkpoint/tests/test_store.py index 053d9a981..9a306a086 100644 --- a/libs/checkpoint/tests/test_store.py +++ b/libs/checkpoint/tests/test_store.py @@ -148,8 +148,8 @@ async def test_async_batch_store(mocker: MockerFixture) -> None: assert abatch.call_count == 1 assert [tuple(c.args[0]) for c in abatch.call_args_list] == [ ( - GetOp(("a",), "b"), - GetOp(("c",), "d"), + GetOp(("a",), "b", refresh_ttl=True), + GetOp(("c",), "d", refresh_ttl=True), ), ] @@ -467,8 +467,8 @@ async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None: assert len(abatch.call_args_list) == 1 ops = list(abatch.call_args_list[0].args[1]) assert len(ops) == 2 - assert GetOp(("test",), "same") in ops - assert GetOp(("test",), "different") in ops + assert GetOp(("test",), "same", refresh_ttl=True) in ops + assert GetOp(("test",), "different", refresh_ttl=True) in ops abatch.reset_mock()