mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
Add TTL option for store items (#3704)
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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,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"
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user