From 0d85891c64be869f044118f661052d6a7a200d85 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Wed, 14 Oct 2026 11:05:22 +0200 Subject: [PATCH 1/3] perf(embeddings-cache): refresh hit TTLs in the lookup pipeline mget_by_keys() and amget_by_keys() fetched a batch in one pipeline, then refreshed each hit's TTL with its own EXPIRE call, so N hits cost 1 + N round trips whenever a TTL was configured. This sits on the hot path of every vectorizer with an EmbeddingsCache attached. Queue EXPIRE next to each HGETALL in the existing non-transactional pipeline and keep only the HGETALL replies. EXPIRE on a missing key returns 0 without creating it, so misses need no special case. The TTL is resolved once and the EXPIREs are skipped when none is set. The pipeline stays non-transactional so cluster keys may span slots. Ten hits on Redis 8.4 over localhost: median 3.99 ms before, 0.93 ms after (500 runs each). --- .../extensions/cache/embeddings/embeddings.py | 40 ++++---- tests/integration/test_embedcache.py | 54 +++++++++++ tests/unit/test_embedcache_pipeline.py | 95 +++++++++++++++++-- 3 files changed, 163 insertions(+), 26 deletions(-) diff --git a/redisvl/extensions/cache/embeddings/embeddings.py b/redisvl/extensions/cache/embeddings/embeddings.py index 53454d3e..136fe0e9 100644 --- a/redisvl/extensions/cache/embeddings/embeddings.py +++ b/redisvl/extensions/cache/embeddings/embeddings.py @@ -233,21 +233,22 @@ def mget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: EmbeddingsCache._warning_shown = True client = self._get_redis_client() + _ttl = self._resolve_ttl() + # Not transactional, so on a Redis Cluster the keys may span slots. with client.pipeline(transaction=False) as pipeline: - # Queue all hgetall operations + # Queue each lookup with its TTL refresh so the whole batch costs + # one round trip. EXPIRE on a missing key is a no-op. for key in keys: pipeline.hgetall(key) + if _ttl: + pipeline.expire(key, _ttl) results = pipeline.execute() - # Process results - processed_results = [] - for i, result in enumerate(results): - if result: # If cache hit, refresh TTL separately - self.expire(keys[i]) - processed_results.append(self._process_cache_data(result)) - - return processed_results + # With a TTL the replies alternate HGETALL, EXPIRE; keep the lookups. + if _ttl: + results = results[::2] + return [self._process_cache_data(result) for result in results] def mget( self, contents: Iterable[bytes | str], model_name: str @@ -660,22 +661,23 @@ async def amget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: return [] client = await self._get_async_redis_client() + _ttl = self._resolve_ttl() - # Use pipeline only for retrieval + # Not transactional, so on a Redis Cluster the keys may span slots. async with client.pipeline(transaction=False) as pipeline: - # Queue all hgetall operations + # Queue each lookup with its TTL refresh so the whole batch costs + # one round trip. EXPIRE on a missing key is a no-op. Queueing is + # synchronous: awaiting a cluster pipeline clears its queue. for key in keys: pipeline.hgetall(key) + if _ttl: + pipeline.expire(key, _ttl) results = await pipeline.execute() - # Process results and refresh TTLs separately - processed_results = [] - for i, result in enumerate(results): - if result: # If cache hit, refresh TTL - await self.aexpire(keys[i]) - processed_results.append(self._process_cache_data(result)) - - return processed_results + # With a TTL the replies alternate HGETALL, EXPIRE; keep the lookups. + if _ttl: + results = results[::2] + return [self._process_cache_data(result) for result in results] async def amget( self, contents: Iterable[bytes | str], model_name: str diff --git a/tests/integration/test_embedcache.py b/tests/integration/test_embedcache.py index 184efe40..fe194341 100644 --- a/tests/integration/test_embedcache.py +++ b/tests/integration/test_embedcache.py @@ -728,6 +728,60 @@ def test_batch_with_ttl(cache_with_ttl, sample_embedding_data): assert all(exist_results) +def _refuse_per_key_expire(monkeypatch, cache): + """Fail the test if a batch lookup refreshes TTLs one round trip at a time.""" + + def refuse(*args, **kwargs): + raise AssertionError("TTL refreshed outside the lookup pipeline") + + monkeypatch.setattr(cache, "expire", refuse) + monkeypatch.setattr(cache, "aexpire", refuse) + + +def _batch_items(sample_embedding_data): + return [ + { + "content": sample["content"], + "model_name": sample["model_name"], + "embedding": sample["embedding"], + } + for sample in sample_embedding_data + ] + + +def test_mget_by_keys_refreshes_every_hit_ttl( + cache_with_ttl, sample_embedding_data, monkeypatch +): + """Every hit gets the default TTL back; a miss is not created.""" + keys = cache_with_ttl.mset(_batch_items(sample_embedding_data), ttl=1000) + missing = cache_with_ttl._make_key("missing") + _refuse_per_key_expire(monkeypatch, cache_with_ttl) + + results = cache_with_ttl.mget_by_keys(keys + [missing]) + + assert [result is not None for result in results] == [True, True, True, False] + client = cache_with_ttl._get_redis_client() + assert all(0 < client.ttl(key) <= 2 for key in keys) + assert client.exists(missing) == 0 + + +@pytest.mark.asyncio +async def test_amget_by_keys_refreshes_every_hit_ttl( + cache_with_ttl, sample_embedding_data, monkeypatch +): + """Every hit gets the default TTL back; a miss is not created.""" + keys = await cache_with_ttl.amset(_batch_items(sample_embedding_data), ttl=1000) + missing = cache_with_ttl._make_key("missing") + _refuse_per_key_expire(monkeypatch, cache_with_ttl) + + results = await cache_with_ttl.amget_by_keys(keys + [missing]) + + assert [result is not None for result in results] == [True, True, True, False] + client = await cache_with_ttl._get_async_redis_client() + assert all([0 < await client.ttl(key) <= 2 for key in keys]) + assert await client.exists(missing) == 0 + + def test_large_batch_operations(cache): """Test operations with larger batches to ensure scalability.""" # Create a larger batch of items diff --git a/tests/unit/test_embedcache_pipeline.py b/tests/unit/test_embedcache_pipeline.py index 75122387..f2eeff00 100644 --- a/tests/unit/test_embedcache_pipeline.py +++ b/tests/unit/test_embedcache_pipeline.py @@ -13,12 +13,17 @@ import pytest from redisvl.extensions.cache.embeddings.embeddings import EmbeddingsCache +from redisvl.extensions.cache.embeddings.schema import CacheEntry class FakeAsyncClusterPipeline: - """Async pipeline that drops its queue when awaited, like ClusterPipeline.""" + """Async pipeline that drops its queue when awaited, like ClusterPipeline. - def __init__(self): + HGETALL replies come from ``store``; every other command replies 1. + """ + + def __init__(self, store=None): + self.store = store or {} self.queued: list[tuple[str, str]] = [] self.executed: list[list[tuple[str, str]]] = [] @@ -47,25 +52,57 @@ def exists(self, key): self.queued.append(("exists", key)) return self - async def execute(self): + def hgetall(self, key): + self.queued.append(("hgetall", key)) + return self + + def _run(self): self.executed.append(list(self.queued)) - results = [1] * len(self.queued) + results = [ + self.store.get(key, {}) if command == "hgetall" else 1 + for command, key in self.queued + ] self.queued.clear() return results + async def execute(self): + return self._run() + + +class FakeSyncPipeline(FakeAsyncClusterPipeline): + """The same recorder behind redis-py's sync pipeline interface.""" + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def execute(self): + return self._run() + class FakeAsyncClusterClient: """Minimal async client that hands out FakeAsyncClusterPipeline instances.""" - def __init__(self): + pipeline_class = FakeAsyncClusterPipeline + + def __init__(self, store=None): + self.store = store self.pipelines: list[FakeAsyncClusterPipeline] = [] + self.transactions: list[bool] = [] - def pipeline(self, transaction=False): - pipeline = FakeAsyncClusterPipeline() + def pipeline(self, transaction=True): + self.transactions.append(transaction) + pipeline = self.pipeline_class(self.store) self.pipelines.append(pipeline) return pipeline +class FakeSyncClient(FakeAsyncClusterClient): + pipeline_class = FakeSyncPipeline + + def make_cache(client, ttl=None): return EmbeddingsCache(name="embedcache", ttl=ttl, async_redis_client=client) @@ -113,3 +150,47 @@ async def test_amexists_by_keys_returns_one_result_per_key(): cache = make_cache(client) assert await cache.amexists_by_keys(["k1", "k2", "k3"]) == [True, True, True] + + +HIT = CacheEntry(entry_id="a", content="a", model_name="m", embedding=[0.1]) + + +@pytest.mark.parametrize( + "ttl, expected", + [ + (None, [("hgetall", "embedcache:a"), ("hgetall", "embedcache:b")]), + ( + 60, + [ + ("hgetall", "embedcache:a"), + ("expire", "embedcache:a"), + ("hgetall", "embedcache:b"), + ("expire", "embedcache:b"), + ], + ), + ], + ids=["no_ttl", "with_ttl"], +) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_mget_by_keys_refreshes_ttls_in_the_lookup_pipeline( + use_async, ttl, expected +): + """Lookups and TTL refreshes share one non-transactional execute(). + + A transactional pipeline would raise CROSSSLOT on a Redis Cluster, and a + refresh outside the pipeline costs one extra round trip per hit. + """ + store = {"embedcache:a": HIT.to_dict()} + if use_async: + client = FakeAsyncClusterClient(store) + cache = make_cache(client, ttl=ttl) + results = await cache.amget_by_keys(["embedcache:a", "embedcache:b"]) + else: + client = FakeSyncClient(store) + cache = EmbeddingsCache(name="embedcache", ttl=ttl, redis_client=client) + results = cache.mget_by_keys(["embedcache:a", "embedcache:b"]) + + assert results == [HIT.model_dump(exclude_none=True), None] + assert client.transactions == [False] + assert [pipeline.executed for pipeline in client.pipelines] == [[expected]] From 5d40a5430114a8069bd81206c37bbae7ba5e340a Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Wed, 14 Oct 2026 11:20:13 +0200 Subject: [PATCH 2/3] test(embeddings-cache): tighten TTL-refresh tests and trim comments Address review feedback on the pipelined TTL refresh: - The fake pipeline now records the EXPIRE TTL, so the unit tests catch a wrong TTL value as well as a refresh outside the lookup pipeline. - Merge the sync and async integration tests into one parametrised test and inline its single-use helper. - Describe the whole unit module in its docstring, and note why the fake client defaults to transaction=True. - Trim the pipeline comments and drop the async-only note on awaiting, which amset already explains. EXPIRE is now also sent for misses. It is a no-op on a missing key, but it means a TTL-configured lookup needs EXPIRE permission even when the whole batch misses; before, only a hit needed it. Such a cache already needs EXPIRE for every write and every hit. A non-hash key in the batch also gets the TTL before HGETALL's WRONGTYPE error is raised. --- .../extensions/cache/embeddings/embeddings.py | 11 ++-- tests/integration/test_embedcache.py | 60 ++++++------------- tests/unit/test_embedcache_pipeline.py | 36 +++++------ 3 files changed, 41 insertions(+), 66 deletions(-) diff --git a/redisvl/extensions/cache/embeddings/embeddings.py b/redisvl/extensions/cache/embeddings/embeddings.py index 136fe0e9..14a7fd88 100644 --- a/redisvl/extensions/cache/embeddings/embeddings.py +++ b/redisvl/extensions/cache/embeddings/embeddings.py @@ -237,15 +237,14 @@ def mget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: # Not transactional, so on a Redis Cluster the keys may span slots. with client.pipeline(transaction=False) as pipeline: - # Queue each lookup with its TTL refresh so the whole batch costs - # one round trip. EXPIRE on a missing key is a no-op. + # Refresh each TTL in the same pipeline; EXPIRE on a missing key is a no-op. for key in keys: pipeline.hgetall(key) if _ttl: pipeline.expire(key, _ttl) results = pipeline.execute() - # With a TTL the replies alternate HGETALL, EXPIRE; keep the lookups. + # Replies alternate HGETALL, EXPIRE; keep the lookups. if _ttl: results = results[::2] return [self._process_cache_data(result) for result in results] @@ -665,16 +664,14 @@ async def amget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: # Not transactional, so on a Redis Cluster the keys may span slots. async with client.pipeline(transaction=False) as pipeline: - # Queue each lookup with its TTL refresh so the whole batch costs - # one round trip. EXPIRE on a missing key is a no-op. Queueing is - # synchronous: awaiting a cluster pipeline clears its queue. + # Refresh each TTL in the same pipeline; EXPIRE on a missing key is a no-op. for key in keys: pipeline.hgetall(key) if _ttl: pipeline.expire(key, _ttl) results = await pipeline.execute() - # With a TTL the replies alternate HGETALL, EXPIRE; keep the lookups. + # Replies alternate HGETALL, EXPIRE; keep the lookups. if _ttl: results = results[::2] return [self._process_cache_data(result) for result in results] diff --git a/tests/integration/test_embedcache.py b/tests/integration/test_embedcache.py index fe194341..0e270a73 100644 --- a/tests/integration/test_embedcache.py +++ b/tests/integration/test_embedcache.py @@ -728,36 +728,29 @@ def test_batch_with_ttl(cache_with_ttl, sample_embedding_data): assert all(exist_results) -def _refuse_per_key_expire(monkeypatch, cache): - """Fail the test if a batch lookup refreshes TTLs one round trip at a time.""" - - def refuse(*args, **kwargs): - raise AssertionError("TTL refreshed outside the lookup pipeline") - - monkeypatch.setattr(cache, "expire", refuse) - monkeypatch.setattr(cache, "aexpire", refuse) - - -def _batch_items(sample_embedding_data): - return [ - { - "content": sample["content"], - "model_name": sample["model_name"], - "embedding": sample["embedding"], - } +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_mget_by_keys_refreshes_every_hit_ttl( + cache_with_ttl, sample_embedding_data, monkeypatch, use_async +): + """Every hit's TTL is reset within the lookup pipeline; a miss is not created.""" + items = [ + {key: sample[key] for key in ("content", "model_name", "embedding")} for sample in sample_embedding_data ] + keys = cache_with_ttl.mset(items, ttl=1000) + missing = cache_with_ttl._make_key("missing") + def refuse(*args, **kwargs): + raise AssertionError("TTL refreshed outside the lookup pipeline") -def test_mget_by_keys_refreshes_every_hit_ttl( - cache_with_ttl, sample_embedding_data, monkeypatch -): - """Every hit gets the default TTL back; a miss is not created.""" - keys = cache_with_ttl.mset(_batch_items(sample_embedding_data), ttl=1000) - missing = cache_with_ttl._make_key("missing") - _refuse_per_key_expire(monkeypatch, cache_with_ttl) + monkeypatch.setattr(cache_with_ttl, "expire", refuse) + monkeypatch.setattr(cache_with_ttl, "aexpire", refuse) - results = cache_with_ttl.mget_by_keys(keys + [missing]) + if use_async: + results = await cache_with_ttl.amget_by_keys(keys + [missing]) + else: + results = cache_with_ttl.mget_by_keys(keys + [missing]) assert [result is not None for result in results] == [True, True, True, False] client = cache_with_ttl._get_redis_client() @@ -765,23 +758,6 @@ def test_mget_by_keys_refreshes_every_hit_ttl( assert client.exists(missing) == 0 -@pytest.mark.asyncio -async def test_amget_by_keys_refreshes_every_hit_ttl( - cache_with_ttl, sample_embedding_data, monkeypatch -): - """Every hit gets the default TTL back; a miss is not created.""" - keys = await cache_with_ttl.amset(_batch_items(sample_embedding_data), ttl=1000) - missing = cache_with_ttl._make_key("missing") - _refuse_per_key_expire(monkeypatch, cache_with_ttl) - - results = await cache_with_ttl.amget_by_keys(keys + [missing]) - - assert [result is not None for result in results] == [True, True, True, False] - client = await cache_with_ttl._get_async_redis_client() - assert all([0 < await client.ttl(key) <= 2 for key in keys]) - assert await client.exists(missing) == 0 - - def test_large_batch_operations(cache): """Test operations with larger batches to ensure scalability.""" # Create a larger batch of items diff --git a/tests/unit/test_embedcache_pipeline.py b/tests/unit/test_embedcache_pipeline.py index f2eeff00..a9924855 100644 --- a/tests/unit/test_embedcache_pipeline.py +++ b/tests/unit/test_embedcache_pipeline.py @@ -1,13 +1,14 @@ -"""Unit tests for how EmbeddingsCache drives async Redis pipelines. +"""Unit tests for how EmbeddingsCache batches commands into Redis pipelines. -Queueing a command on a redis-py async pipeline is synchronous and returns the -pipeline itself, so a queued command must never be awaited. On -``redis.asyncio.cluster.ClusterPipeline`` awaiting it calls ``initialize()``, -which clears the queued commands -- batched writes then vanish without an error. +Each batch method must send all of its commands, TTL refreshes included, in one +non-transactional execute(). On the async side a queued command must never be +awaited: queueing is synchronous and returns the pipeline, and awaiting a +``redis.asyncio.cluster.ClusterPipeline`` calls ``initialize()``, which clears +the queue -- batched commands then vanish without an error. -The fake below reproduces exactly those two behaviours, so these tests fail if -an ``await`` is ever reintroduced. They need no Redis, so unlike the cluster -tests in ``tests/integration/test_redis_cluster_support.py`` they run by default. +The fakes record what each execute() receives and reproduce that queue-clearing. +They need no Redis, so unlike the cluster tests in +``tests/integration/test_redis_cluster_support.py`` these run by default. """ import pytest @@ -24,8 +25,8 @@ class FakeAsyncClusterPipeline: def __init__(self, store=None): self.store = store or {} - self.queued: list[tuple[str, str]] = [] - self.executed: list[list[tuple[str, str]]] = [] + self.queued: list[tuple] = [] + self.executed: list[list[tuple]] = [] def __await__(self): async def initialize(): @@ -45,7 +46,7 @@ def hset(self, name, mapping): return self def expire(self, key, ttl): - self.queued.append(("expire", key)) + self.queued.append(("expire", key, ttl)) return self def exists(self, key): @@ -60,7 +61,7 @@ def _run(self): self.executed.append(list(self.queued)) results = [ self.store.get(key, {}) if command == "hgetall" else 1 - for command, key in self.queued + for command, key, *_ in self.queued ] self.queued.clear() return results @@ -83,7 +84,7 @@ def execute(self): class FakeAsyncClusterClient: - """Minimal async client that hands out FakeAsyncClusterPipeline instances.""" + """Minimal client that hands out recording pipelines of ``pipeline_class``.""" pipeline_class = FakeAsyncClusterPipeline @@ -92,6 +93,7 @@ def __init__(self, store=None): self.pipelines: list[FakeAsyncClusterPipeline] = [] self.transactions: list[bool] = [] + # redis-py's default, so a missing transaction=False is caught. def pipeline(self, transaction=True): self.transactions.append(transaction) pipeline = self.pipeline_class(self.store) @@ -115,9 +117,9 @@ def make_cache(client, ttl=None): 60, [ ("hset", "embedcache:a"), - ("expire", "embedcache:a"), + ("expire", "embedcache:a", 60), ("hset", "embedcache:b"), - ("expire", "embedcache:b"), + ("expire", "embedcache:b", 60), ], ), ], @@ -163,9 +165,9 @@ async def test_amexists_by_keys_returns_one_result_per_key(): 60, [ ("hgetall", "embedcache:a"), - ("expire", "embedcache:a"), + ("expire", "embedcache:a", 60), ("hgetall", "embedcache:b"), - ("expire", "embedcache:b"), + ("expire", "embedcache:b", 60), ], ), ], From dbc317272a02e6b1df695b1a9de1b44c3b31d979 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Wed, 14 Oct 2026 11:26:11 +0200 Subject: [PATCH 3/3] perf(embeddings-cache): route single-key get and set through the pipelines With a TTL configured, a get_by_key()/aget_by_key() hit ran HGETALL and then a separate EXPIRE, and set()/aset() ran HSET and then a separate EXPIRE: two round trips each. set() also briefly left a fresh entry with no expiry, the window mset() already closes. Delegate them to mget_by_keys()/amget_by_keys() and mset()/amset(), which queue the EXPIRE in the same non-transactional pipeline. Each call now costs one round trip and four duplicated bodies go away. get_by_key() and set() keep their own async-only warning, so the hint still names the method the caller used. As with the batch read, get_by_key() and aget_by_key() now send EXPIRE on a miss too when a TTL is set. It leaves a missing key untouched, but a credential without write access to the key now fails on a miss as well as a hit; the installation guide's ACL section now says so. Sync timings on Redis 8.4 over Docker Desktop localhost: get() 381 us before, 249 us after; set() 505 us before, 248 us after. --- docs/user_guide/installation.md | 2 +- .../extensions/cache/embeddings/embeddings.py | 62 +++++-------------- tests/integration/test_embedcache_warnings.py | 8 +-- tests/unit/test_embedcache_pipeline.py | 28 +++++++++ 4 files changed, 49 insertions(+), 51 deletions(-) diff --git a/docs/user_guide/installation.md b/docs/user_guide/installation.md index aa0e3065..4e419c31 100644 --- a/docs/user_guide/installation.md +++ b/docs/user_guide/installation.md @@ -227,7 +227,7 @@ RedisSearchError: Error while fetching llmcache index info: User has no permissions to run the 'FT.INFO' command ``` -Reading a cache is not read-only either. When a cache has a TTL configured — at construction or later via `set_ttl()` — every read that hits refreshes the matched entries' TTL, so it issues `EXPIRE`. That includes `SemanticCache.check()`, every `EmbeddingsCache` getter, their async equivalents, and the reads a vectorizer performs for you when constructed with `cache=`. `EXPIRE` is in `@write`, not `@read` (measured with `ACL CAT write`), so a lookup-only credential fails on a cache *hit* rather than on the write that populated it. +Reading a cache is not read-only either. When a cache has a TTL configured — at construction or later via `set_ttl()` — every read that hits refreshes the matched entries' TTL, so it issues `EXPIRE`. That includes `SemanticCache.check()`, every `EmbeddingsCache` getter, their async equivalents, and the reads a vectorizer performs for you when constructed with `cache=`. `EXPIRE` is in `@write`, not `@read` (measured with `ACL CAT write`), so a lookup-only credential fails on a cache *hit* rather than on the write that populated it. `EmbeddingsCache` getters send `EXPIRE` in the same pipeline as the lookup, so for them it fails on a miss too. Granting the command is only half of it, because the key patterns must permit writes too. Measured on 8.4.5 with `ACL DRYRUN EXPIRE llmcache:abc 60`: diff --git a/redisvl/extensions/cache/embeddings/embeddings.py b/redisvl/extensions/cache/embeddings/embeddings.py index 14a7fd88..513bca9b 100644 --- a/redisvl/extensions/cache/embeddings/embeddings.py +++ b/redisvl/extensions/cache/embeddings/embeddings.py @@ -189,16 +189,7 @@ def get_by_key(self, key: str) -> dict[str, Any] | None: ) EmbeddingsCache._warning_shown = True - client = self._get_redis_client() - - # Get all fields - data = client.hgetall(key) - - # Refresh TTL if data exists - if data: - self.expire(key) - - return self._process_cache_data(data) # type: ignore + return self.mget_by_keys([key])[0] def mget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: """Get multiple embeddings by their Redis keys. @@ -310,11 +301,6 @@ def set( metadata={"source": "user_query"} ) """ - # Prepare data - key, cache_entry = self._prepare_entry_data( - content, model_name, embedding, metadata - ) - if self._should_warn_for_async_only(): if not EmbeddingsCache._warning_shown: logger.warning( @@ -323,14 +309,13 @@ def set( ) EmbeddingsCache._warning_shown = True - # Store in Redis - client = self._get_redis_client() - client.hset(name=key, mapping=cache_entry) # type: ignore[arg-type] - - # Set TTL if specified - self.expire(key, ttl) - - return key + item = { + "content": content, + "model_name": model_name, + "embedding": embedding, + "metadata": metadata, + } + return self.mset([item], ttl=ttl)[0] def mset( self, @@ -624,16 +609,7 @@ async def aget_by_key(self, key: str) -> dict[str, Any] | None: embedding_data = await cache.aget_by_key("embedcache:1234567890abcdef") """ - client = await self._get_async_redis_client() - - # Get all fields - data = await client.hgetall(key) # type: ignore - - # Refresh TTL if data exists - if data: - await self.aexpire(key) - - return self._process_cache_data(data) + return (await self.amget_by_keys([key]))[0] async def amget_by_keys(self, keys: list[str]) -> list[dict[str, Any] | None]: """Async get multiple embeddings by their Redis keys. @@ -739,19 +715,13 @@ async def aset( metadata={"source": "user_query"} ) """ - # Prepare data - key, cache_entry = self._prepare_entry_data( - content, model_name, embedding, metadata - ) - - # Store in Redis - client = await self._get_async_redis_client() - await client.hset(name=key, mapping=cache_entry) # type: ignore - - # Set TTL if specified - await self.aexpire(key, ttl) - - return key + item = { + "content": content, + "model_name": model_name, + "embedding": embedding, + "metadata": metadata, + } + return (await self.amset([item], ttl=ttl))[0] async def amset( self, diff --git a/tests/integration/test_embedcache_warnings.py b/tests/integration/test_embedcache_warnings.py index 13823dd5..2e467b67 100644 --- a/tests/integration/test_embedcache_warnings.py +++ b/tests/integration/test_embedcache_warnings.py @@ -30,10 +30,9 @@ async def test_sync_methods_warn_with_async_only_client( # Mock _get_redis_client to prevent actual connection attempt with patch.object(cache, "_get_redis_client") as mock_get_client: - # Mock the Redis client methods that would be called - mock_client = mock_get_client.return_value - mock_client.hgetall.return_value = {} # Empty result for get_by_key - mock_client.hset.return_value = 1 # Success for set + # Both calls run through a pipeline; with no TTL, one empty reply is a miss + pipeline = mock_get_client.return_value.pipeline.return_value.__enter__() + pipeline.execute.return_value = [{}] # Capture log warnings with caplog.at_level(logging.WARNING): @@ -46,6 +45,7 @@ async def test_sync_methods_warn_with_async_only_client( "initialized with async_redis_client only" in caplog.records[0].message ) assert "Use async methods" in caplog.records[0].message + assert "(get_by_key)" in caplog.records[0].message # Clear captured logs caplog.clear() diff --git a/tests/unit/test_embedcache_pipeline.py b/tests/unit/test_embedcache_pipeline.py index a9924855..b487b5c6 100644 --- a/tests/unit/test_embedcache_pipeline.py +++ b/tests/unit/test_embedcache_pipeline.py @@ -196,3 +196,31 @@ async def test_mget_by_keys_refreshes_ttls_in_the_lookup_pipeline( assert results == [HIT.model_dump(exclude_none=True), None] assert client.transactions == [False] assert [pipeline.executed for pipeline in client.pipelines] == [[expected]] + + +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_single_key_get_and_set_each_cost_one_pipeline(monkeypatch, use_async): + """Single-key reads and writes queue their EXPIRE in the same pipeline.""" + monkeypatch.setattr( + EmbeddingsCache, "_make_entry_id", lambda self, content, model_name: content + ) + store = {"embedcache:a": HIT.to_dict()} + if use_async: + client = FakeAsyncClusterClient(store) + cache = make_cache(client, ttl=60) + hit = await cache.aget_by_key("embedcache:a") + key = await cache.aset(content="b", model_name="m", embedding=[0.1], ttl=5) + else: + client = FakeSyncClient(store) + cache = EmbeddingsCache(name="embedcache", ttl=60, redis_client=client) + hit = cache.get_by_key("embedcache:a") + key = cache.set(content="b", model_name="m", embedding=[0.1], ttl=5) + + assert hit == HIT.model_dump(exclude_none=True) + assert key == "embedcache:b" + assert client.transactions == [False, False] + assert [pipeline.executed for pipeline in client.pipelines] == [ + [[("hgetall", "embedcache:a"), ("expire", "embedcache:a", 60)]], + [[("hset", "embedcache:b"), ("expire", "embedcache:b", 5)]], + ]