diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ae29958c..33e79f75 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -171,6 +171,9 @@ make test-all # Run tests on a specific file uv run pytest tests/unit/test_fields.py -v +# Run the Redis Cluster tests, which are skipped by default +uv run pytest --run-cluster-tests -m requires_cluster + # Run tests with coverage uv run pytest --cov=redisvl --cov-report=html ``` @@ -184,6 +187,8 @@ belonging to other workers and other pull requests' CI runs. See the module docstring in `tests/integration/test_langcache_semantic_cache_integration.py` for a worked example. +**Note:** Tests marked `requires_cluster` only run when you pass `--run-cluster-tests`. The cluster itself is provisioned for you by the `redis_cluster_container` fixture in `tests/conftest.py`, so Docker needs to be running. + ## Documentation Documentation is served from the `docs/` directory and built using Sphinx. diff --git a/docs/user_guide/10_embeddings_cache.ipynb b/docs/user_guide/10_embeddings_cache.ipynb index e74afc02..ddf9b9e1 100644 --- a/docs/user_guide/10_embeddings_cache.ipynb +++ b/docs/user_guide/10_embeddings_cache.ipynb @@ -123,7 +123,7 @@ "### Storing Embeddings\n", "\n", "Let's store some text with its embedding in the cache. The `set` method takes the following parameters:\n", - "- `text`: The input text that was embedded\n", + "- `content`: The input text that was embedded\n", "- `model_name`: The name of the embedding model used\n", "- `embedding`: The embedding vector\n", "- `metadata`: Optional metadata associated with the embedding\n", diff --git a/redisvl/extensions/cache/base.py b/redisvl/extensions/cache/base.py index f14abb86..c3e02989 100644 --- a/redisvl/extensions/cache/base.py +++ b/redisvl/extensions/cache/base.py @@ -148,6 +148,17 @@ async def _get_async_redis_client(self) -> AsyncRedisClient: ) return self._async_redis_client + def _resolve_ttl(self, ttl: int | None = None) -> int | None: + """Resolve an explicit TTL against this cache's default TTL. + + Args: + ttl (Optional[int], optional): An explicit time-to-live in seconds. + + Returns: + Optional[int]: The TTL to apply, or None if no expiration is set. + """ + return ttl if ttl is not None else self._ttl + def expire(self, key: str, ttl: int | None = None) -> None: """Set or refresh the expiration time for a key in the cache. @@ -161,7 +172,7 @@ def expire(self, key: str, ttl: int | None = None) -> None: If neither the provided TTL nor the default TTL is set (both are None), this method will have no effect. """ - _ttl = ttl if ttl is not None else self._ttl + _ttl = self._resolve_ttl(ttl) if _ttl: client = self._get_redis_client() client.expire(key, _ttl) @@ -179,7 +190,7 @@ async def aexpire(self, key: str, ttl: int | None = None) -> None: If neither the provided TTL nor the default TTL is set (both are None), this method will have no effect. """ - _ttl = ttl if ttl is not None else self._ttl + _ttl = self._resolve_ttl(ttl) if _ttl: client = await self._get_async_redis_client() await client.expire(key, _ttl) diff --git a/redisvl/extensions/cache/embeddings/embeddings.py b/redisvl/extensions/cache/embeddings/embeddings.py index 20ebab7e..53454d3e 100644 --- a/redisvl/extensions/cache/embeddings/embeddings.py +++ b/redisvl/extensions/cache/embeddings/embeddings.py @@ -325,7 +325,7 @@ def set( # Store in Redis client = self._get_redis_client() - client.hset(name=key, mapping=cache_entry) # type: ignore + client.hset(name=key, mapping=cache_entry) # type: ignore[arg-type] # Set TTL if specified self.expire(key, ttl) @@ -352,6 +352,11 @@ def mset( Returns: List[str]: List of Redis keys where the embeddings were stored. + Note: + The batch is pipelined, not transactional, so on a Redis Cluster it + fans out across shards. If it fails partway, some entries will have + been written; the operation is idempotent, so simply retry it. + .. code-block:: python # Store multiple embeddings @@ -383,6 +388,7 @@ def mset( client = self._get_redis_client() keys = [] + _ttl = self._resolve_ttl(ttl) with client.pipeline(transaction=False) as pipeline: # Process all entries @@ -390,14 +396,14 @@ def mset( # Prepare and store key, cache_entry = self._prepare_entry_data(**item) keys.append(key) - pipeline.hset(name=key, mapping=cache_entry) # type: ignore + pipeline.hset(name=key, mapping=cache_entry) # type: ignore[arg-type] + # Queue the expiry with its write so no entry is ever left + # unexpiring. HSET on its own leaves a key's TTL untouched. + if _ttl: + pipeline.expire(key, _ttl) pipeline.execute() - # Set TTLs - for key in keys: - self.expire(key, ttl) - return keys def exists(self, content: bytes | str, model_name: str) -> bool: @@ -768,6 +774,11 @@ async def amset( Returns: List[str]: List of Redis keys where the embeddings were stored. + Note: + The batch is pipelined, not transactional, so on a Redis Cluster it + fans out across shards. If it fails partway, some entries will have + been written; the operation is idempotent, so simply retry it. + .. code-block:: python # Store multiple embeddings asynchronously @@ -791,6 +802,7 @@ async def amset( client = await self._get_async_redis_client() keys = [] + _ttl = self._resolve_ttl(ttl) async with client.pipeline(transaction=False) as pipeline: # Process all entries @@ -798,14 +810,17 @@ async def amset( # Prepare and store key, cache_entry = self._prepare_entry_data(**item) keys.append(key) - await pipeline.hset(name=key, mapping=cache_entry) # type: ignore + # Never await a queued command: queueing is synchronous, and + # awaiting a cluster pipeline calls initialize(), which clears + # the queue. Only execute() is awaited. + pipeline.hset(name=key, mapping=cache_entry) # type: ignore[arg-type] + # Queue the expiry with its write so no entry is ever left + # unexpiring. HSET on its own leaves a key's TTL untouched. + if _ttl: + pipeline.expire(key, _ttl) await pipeline.execute() - # Set TTLs - for key in keys: - await self.aexpire(key, ttl) - return keys async def amexists_by_keys(self, keys: list[str]) -> list[bool]: @@ -833,7 +848,7 @@ async def amexists_by_keys(self, keys: list[str]) -> list[bool]: async with client.pipeline(transaction=False) as pipeline: # Queue all exists operations for key in keys: - await pipeline.exists(key) + pipeline.exists(key) results = await pipeline.execute() # Convert to boolean values diff --git a/tests/integration/test_redis_cluster_support.py b/tests/integration/test_redis_cluster_support.py index fb5c1ea9..76c0fcb0 100644 --- a/tests/integration/test_redis_cluster_support.py +++ b/tests/integration/test_redis_cluster_support.py @@ -156,7 +156,12 @@ async def test_async_search_index_client(redis_cluster_url, redis_test_name): @pytest.mark.requires_cluster @pytest.mark.asyncio async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name): - """Test that EmbeddingsCache correctly handles AsyncRedisCluster clients.""" + """Test that EmbeddingsCache correctly handles AsyncRedisCluster clients. + + Batch writes are the subtle case. Queueing a command on an async pipeline is + synchronous, so awaiting the returned pipeline used to clear a cluster + pipeline's queue -- ``amset`` returned every key having written nothing. + """ cluster_client = RedisConnectionFactory.get_async_redis_cluster_connection( redis_cluster_url ) @@ -164,6 +169,12 @@ async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name name=redis_test_name("embedcache"), async_redis_client=cluster_client ) + contents = [f"hey_{i}" for i in range(10)] + items = [ + {"content": content, "model_name": "test", "embedding": [1.0, 2.0, float(i)]} + for i, content in enumerate(contents) + ] + try: await cache.aset( content="hey", @@ -175,6 +186,23 @@ async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name assert result["embedding"] == [1, 2, 3] await cache.aclear() assert await cache.aget("hey", "test") is None + + await cache.amset(items) + + # Count with scan_iter, which fans out to every primary. KEYS and DBSIZE + # are routed to a single node and would only report one shard's worth. + # SCAN only promises each key at least once, so de-duplicate. + prefix = cache._get_prefix() + scanned = {key async for key in cluster_client.scan_iter(match=f"{prefix}*")} + assert len(scanned) == len(items) + + results = await cache.amget(contents, "test") + assert all(result is not None for result in results) + assert results[3]["embedding"] == [1.0, 2.0, 3.0] + + # amexists_by_keys queued through the same pipeline and returned []. + assert await cache.amexists(contents, "test") == [True] * len(items) + await cache.aclear() finally: # Manually close the cluster client to prevent connection leaks await cluster_client.aclose() diff --git a/tests/unit/test_embedcache_pipeline.py b/tests/unit/test_embedcache_pipeline.py new file mode 100644 index 00000000..75122387 --- /dev/null +++ b/tests/unit/test_embedcache_pipeline.py @@ -0,0 +1,115 @@ +"""Unit tests for how EmbeddingsCache drives async 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. + +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. +""" + +import pytest + +from redisvl.extensions.cache.embeddings.embeddings import EmbeddingsCache + + +class FakeAsyncClusterPipeline: + """Async pipeline that drops its queue when awaited, like ClusterPipeline.""" + + def __init__(self): + self.queued: list[tuple[str, str]] = [] + self.executed: list[list[tuple[str, str]]] = [] + + def __await__(self): + async def initialize(): + self.queued.clear() + return self + + return initialize().__await__() + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + return False + + def hset(self, name, mapping): + self.queued.append(("hset", name)) + return self + + def expire(self, key, ttl): + self.queued.append(("expire", key)) + return self + + def exists(self, key): + self.queued.append(("exists", key)) + return self + + async def execute(self): + self.executed.append(list(self.queued)) + results = [1] * len(self.queued) + self.queued.clear() + return results + + +class FakeAsyncClusterClient: + """Minimal async client that hands out FakeAsyncClusterPipeline instances.""" + + def __init__(self): + self.pipelines: list[FakeAsyncClusterPipeline] = [] + + def pipeline(self, transaction=False): + pipeline = FakeAsyncClusterPipeline() + self.pipelines.append(pipeline) + return pipeline + + +def make_cache(client, ttl=None): + return EmbeddingsCache(name="embedcache", ttl=ttl, async_redis_client=client) + + +@pytest.mark.parametrize( + "ttl, expected", + [ + (None, [("hset", "embedcache:a"), ("hset", "embedcache:b")]), + ( + 60, + [ + ("hset", "embedcache:a"), + ("expire", "embedcache:a"), + ("hset", "embedcache:b"), + ("expire", "embedcache:b"), + ], + ), + ], + ids=["no_ttl", "with_ttl"], +) +@pytest.mark.asyncio +async def test_amset_sends_every_command_in_one_pipeline(monkeypatch, ttl, expected): + """amset must queue each write, and its expiry, into a single execute().""" + monkeypatch.setattr( + EmbeddingsCache, "_make_entry_id", lambda self, content, model_name: content + ) + client = FakeAsyncClusterClient() + cache = make_cache(client, ttl=ttl) + items = [ + {"content": name, "model_name": "m", "embedding": [0.1, 0.2]} + for name in ("a", "b") + ] + + keys = await cache.amset(items) + + assert keys == ["embedcache:a", "embedcache:b"] + # A dropped queue would show up as a single empty execute(). + assert [pipeline.executed for pipeline in client.pipelines] == [[expected]] + + +@pytest.mark.asyncio +async def test_amexists_by_keys_returns_one_result_per_key(): + """A dropped queue made this return [] rather than a bool per key.""" + client = FakeAsyncClusterClient() + cache = make_cache(client) + + assert await cache.amexists_by_keys(["k1", "k2", "k3"]) == [True, True, True]