Add TTL option for store items (#3704)

This commit is contained in:
William FH
2025-03-06 13:20:32 -08:00
committed by GitHub
parent 79595d43a5
commit 09bd5990d4
6 changed files with 124 additions and 34 deletions
@@ -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):
@@ -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()
@@ -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.
+29 -12
View File
@@ -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
+1 -1
View File
@@ -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"
+4 -4
View File
@@ -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()