diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 54138d77c..c8dd135d2 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -163,12 +163,17 @@ 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 ``` **Note:** Tests requiring external APIs need appropriate API keys set as environment variables. +**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 e74afc024..ddf9b9e1a 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 f6e965a07..645060439 100644 --- a/redisvl/extensions/cache/base.py +++ b/redisvl/extensions/cache/base.py @@ -144,6 +144,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. @@ -157,7 +168,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) @@ -175,7 +186,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 7218de5c4..2355dd02f 100644 --- a/redisvl/extensions/cache/embeddings/embeddings.py +++ b/redisvl/extensions/cache/embeddings/embeddings.py @@ -321,7 +321,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) # Set TTL if specified self.expire(key, ttl) @@ -348,6 +348,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 @@ -379,6 +384,7 @@ def mset( client = self._get_redis_client() keys = [] + _ttl = self._resolve_ttl(ttl) with client.pipeline(transaction=False) as pipeline: # Process all entries @@ -386,14 +392,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) + # 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: @@ -764,6 +770,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 @@ -787,6 +798,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 @@ -794,14 +806,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) + # 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]: @@ -829,7 +844,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 fc01922c0..1850808bf 100644 --- a/tests/integration/test_redis_cluster_support.py +++ b/tests/integration/test_redis_cluster_support.py @@ -154,7 +154,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 ) @@ -162,9 +167,15 @@ 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( - text="hey", + content="hey", model_name="test", embedding=[1, 2, 3], ) @@ -172,6 +183,23 @@ async def test_embeddings_cache_cluster_async(redis_cluster_url, redis_test_name assert result is not None assert result["embedding"] == [1, 2, 3] await cache.aclear() + + 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() @@ -187,7 +215,7 @@ def test_embeddings_cache_cluster_sync(redis_cluster_url, redis_test_name): for i in range(100): cache.set( - text=f"hey_{i}", + content=f"hey_{i}", model_name="test", embedding=[1, 2, 3], ) @@ -198,8 +226,8 @@ def test_embeddings_cache_cluster_sync(redis_cluster_url, redis_test_name): cache.mset( [ - {"text": "hey_0", "model_name": "test", "embedding": [1, 2, 3]}, - {"text": "hey_1", "model_name": "test", "embedding": [1, 2, 3]}, + {"content": "hey_0", "model_name": "test", "embedding": [1, 2, 3]}, + {"content": "hey_1", "model_name": "test", "embedding": [1, 2, 3]}, ] ) result = cache.mget(["hey_0", "hey_1"], "test") diff --git a/tests/unit/test_embedcache_pipeline.py b/tests/unit/test_embedcache_pipeline.py new file mode 100644 index 000000000..751223873 --- /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]