From 852a129881705f382263b9b77a91d9ef3c5fc584 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Tue, 11 Mar 2025 15:53:13 -0700 Subject: [PATCH] Default store ttl config --- .../langgraph/store/postgres/base.py | 3 + .../langgraph/store/base/__init__.py | 84 ++++++++++++++++--- libs/checkpoint/langgraph/store/base/batch.py | 30 +++++-- 3 files changed, 100 insertions(+), 17 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 4ee98926e..c8445dd42 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -39,6 +39,7 @@ from langgraph.store.base import ( Result, SearchItem, SearchOp, + TTLConfig, ensure_embeddings, get_text_at_path, tokenize_path, @@ -622,6 +623,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, index: Optional[PostgresIndexConfig] = None, + ttl: Optional[TTLConfig] = None, ) -> None: super().__init__() self._deserializer = deserializer @@ -634,6 +636,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): self.embeddings, self.index_config = _ensure_index_config(self.index_config) else: self.embeddings = None + self.ttl_config = ttl @classmethod @contextmanager diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index cb0235d5a..190f61689 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -496,6 +496,18 @@ class InvalidNamespaceError(ValueError): """Provided namespace is invalid.""" +class TTLConfig(TypedDict, total=False): + """Configuration for TTL (time-to-live) behavior in the store.""" + + refresh_on_read: bool + """Default behavior for refreshing TTLs on read operations (GET and SEARCH). + + If True, TTLs will be refreshed on read operations (get/search) by default. + This can be overridden per-operation by explicitly setting refresh_ttl. + Defaults to True if not configured. + """ + + class IndexConfig(TypedDict, total=False): """Configuration for indexing documents for semantic search in the store. @@ -640,7 +652,8 @@ class BaseStore(ABC): Subclasses must explicitly set `supports_ttl = True` to enable this feature. """ - supports_ttl = False + supports_ttl: bool = False + ttl_config: Optional[TTLConfig] = None __slots__ = ("__weakref__",) @@ -669,7 +682,11 @@ class BaseStore(ABC): """ def get( - self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + self, + namespace: tuple[str, ...], + key: str, + *, + refresh_ttl: Optional[bool] = None, ) -> Optional[Item]: """Retrieve a single item. @@ -677,12 +694,15 @@ class BaseStore(ABC): namespace: Hierarchical path for the item. key: Unique identifier within the namespace. refresh_ttl: Whether to refresh TTLs for the returned item. + If None (default), uses the store's default refresh_ttl setting. If no TTL is specified, this argument is ignored. Returns: The retrieved item or None if not found. """ - return self.batch([GetOp(namespace, str(key), refresh_ttl)])[0] + return self.batch( + [GetOp(namespace, str(key), _ensure_refresh(self.ttl_config, refresh_ttl))] + )[0] def search( self, @@ -693,7 +713,7 @@ class BaseStore(ABC): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - refresh_ttl: bool = True, + refresh_ttl: Optional[bool] = None, ) -> list[SearchItem]: """Search for items within a namespace prefix. @@ -743,7 +763,16 @@ class BaseStore(ABC): and requires proper embedding configuration. """ return self.batch( - [SearchOp(namespace_prefix, filter, limit, offset, query, refresh_ttl)] + [ + SearchOp( + namespace_prefix, + filter, + limit, + offset, + query, + _ensure_refresh(self.ttl_config, refresh_ttl), + ) + ] )[0] def put( @@ -876,7 +905,11 @@ class BaseStore(ABC): return self.batch([op])[0] async def aget( - self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + self, + namespace: tuple[str, ...], + key: str, + *, + refresh_ttl: Optional[bool] = None, ) -> Optional[Item]: """Asynchronously retrieve a single item. @@ -887,7 +920,17 @@ class BaseStore(ABC): Returns: The retrieved item or None if not found. """ - return (await self.abatch([GetOp(namespace, str(key), refresh_ttl)]))[0] + return ( + await self.abatch( + [ + GetOp( + namespace, + str(key), + _ensure_refresh(self.ttl_config, refresh_ttl), + ) + ] + ) + )[0] async def asearch( self, @@ -898,7 +941,7 @@ class BaseStore(ABC): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - refresh_ttl: bool = True, + refresh_ttl: Optional[bool] = None, ) -> list[SearchItem]: """Asynchronously search for items within a namespace prefix. @@ -909,8 +952,8 @@ class BaseStore(ABC): 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. + If None (default), uses the store's TTLConfig.refresh_default setting. + If TTLConfig is not provided or no TTL is specified, this argument is ignored. Returns: List of items matching the search criteria. @@ -950,7 +993,16 @@ class BaseStore(ABC): """ return ( await self.abatch( - [SearchOp(namespace_prefix, filter, limit, offset, query, refresh_ttl)] + [ + SearchOp( + namespace_prefix, + filter, + limit, + offset, + query, + _ensure_refresh(self.ttl_config, refresh_ttl), + ) + ] ) )[0] @@ -1116,6 +1168,16 @@ def _validate_namespace(namespace: tuple[str, ...]) -> None: ) +def _ensure_refresh( + ttl_config: Optional[TTLConfig], refresh_ttl: Optional[bool] = None +) -> bool: + if refresh_ttl is not None: + return refresh_ttl + if ttl_config is not None: + return ttl_config.get("refresh_on_read", True) + return True + + __all__ = [ "BaseStore", "Item", diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index d0bba1efe..16d0b1f36 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -16,6 +16,7 @@ from langgraph.store.base import ( Result, SearchItem, SearchOp, + _ensure_refresh, _validate_namespace, ) @@ -65,11 +66,24 @@ class AsyncBatchedBaseStore(BaseStore): pass async def aget( - self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + self, + namespace: tuple[str, ...], + key: str, + *, + refresh_ttl: Optional[bool] = None, ) -> Optional[Item]: assert not self._task.done() fut = self._loop.create_future() - self._aqueue.put_nowait((fut, GetOp(namespace, key, refresh_ttl=refresh_ttl))) + self._aqueue.put_nowait( + ( + fut, + GetOp( + namespace, + key, + refresh_ttl=_ensure_refresh(self.ttl_config, refresh_ttl), + ), + ) + ) return await fut async def asearch( @@ -81,7 +95,7 @@ class AsyncBatchedBaseStore(BaseStore): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - refresh_ttl: bool = True, + refresh_ttl: Optional[bool] = None, ) -> list[SearchItem]: assert not self._task.done() fut = self._loop.create_future() @@ -94,7 +108,7 @@ class AsyncBatchedBaseStore(BaseStore): limit, offset, query, - refresh_ttl=refresh_ttl, + refresh_ttl=_ensure_refresh(self.ttl_config, refresh_ttl), ), ) ) @@ -157,7 +171,11 @@ class AsyncBatchedBaseStore(BaseStore): @_check_loop def get( - self, namespace: tuple[str, ...], key: str, *, refresh_ttl: bool = True + self, + namespace: tuple[str, ...], + key: str, + *, + refresh_ttl: Optional[bool] = None, ) -> Optional[Item]: return asyncio.run_coroutine_threadsafe( self.aget(namespace, key=key, refresh_ttl=refresh_ttl), self._loop @@ -173,7 +191,7 @@ class AsyncBatchedBaseStore(BaseStore): filter: Optional[dict[str, Any]] = None, limit: int = 10, offset: int = 0, - refresh_ttl: bool = True, + refresh_ttl: Optional[bool] = None, ) -> list[SearchItem]: return asyncio.run_coroutine_threadsafe( self.asearch(