From bcb2351298915565f4abce5e07dd800f2e4a7717 Mon Sep 17 00:00:00 2001 From: Vishal Bala Date: Tue, 13 Oct 2026 17:45:57 +0200 Subject: [PATCH] perf: run SemanticCache.check() as KNN with a client-side threshold check() and acheck() searched with a VectorRangeQuery, so Redis discarded every candidate beyond distance_threshold and a miss carried no information about how close it came. They now run a KNN VectorQuery for num_results candidates and keep those at or within the threshold in Python, inclusive as the range radius was. On the FLAT index the cache always uses, both queries are exhaustive, so hit and miss decisions match. Measured on Redis 8.4.7 they agreed on 3,000 of 3,000 samples per configuration and on 200 of 200 queries within 2e-6 of the threshold. The one theoretical difference is that the client compares the distance as Redis returns it, rounded to 12 significant digits, so a distance within about 1e-12 of the threshold can now fall on the other side of it. Unfiltered lookups take 2-3% longer. With a tag filter matching 10% of entries they take 13% less time at 1,000 entries and 36% less at 10,000, because KNN compares vectors only within the filtered subset where the range query scanned the whole index. The nearest candidate's distance is now in hand on every lookup, which the planned OpenTelemetry instrumentation reports. A per-call distance_threshold is still validated as the range query validated it: TypeError for a non-number, ValueError for a negative value. An id-only row (the Redis 8.8+ expiry race) still reaches _process_cache_results and is skipped there rather than raising. Query construction and the threshold filter are shared by the sync and async paths. --- redisvl/extensions/cache/llm/schema.py | 2 + redisvl/extensions/cache/llm/semantic.py | 90 +++++++++++++------ tests/integration/test_llmcache.py | 44 +++++++++ tests/unit/test_extension_result_hardening.py | 9 ++ 4 files changed, 117 insertions(+), 28 deletions(-) diff --git a/redisvl/extensions/cache/llm/schema.py b/redisvl/extensions/cache/llm/schema.py index c0cdd4e24..4442b7ac8 100644 --- a/redisvl/extensions/cache/llm/schema.py +++ b/redisvl/extensions/cache/llm/schema.py @@ -128,6 +128,8 @@ def from_params(cls, name: str, prefix: str, vector_dims: int, dtype: str): "dims": vector_dims, "datatype": dtype, "distance_metric": "cosine", + # SemanticCache.check() filters an exact KNN by + # threshold; HNSW would make its hits approximate. "algorithm": "flat", }, }, diff --git a/redisvl/extensions/cache/llm/semantic.py b/redisvl/extensions/cache/llm/semantic.py index 439423cc2..cbe7eaa20 100644 --- a/redisvl/extensions/cache/llm/semantic.py +++ b/redisvl/extensions/cache/llm/semantic.py @@ -23,7 +23,7 @@ UPDATED_AT_FIELD_NAME, ) from redisvl.index import AsyncSearchIndex, SearchIndex -from redisvl.query import VectorRangeQuery +from redisvl.query import VectorQuery from redisvl.query.filter import FilterExpression from redisvl.redis.utils import hashify from redisvl.utils.log import get_logger @@ -486,6 +486,8 @@ def check( # (exact match only), so only fall back on None. if distance_threshold is None: distance_threshold = self._distance_threshold + else: + self._validate_distance_threshold(distance_threshold) # Vectorize prompt if not provided if vector is None and prompt is not None: @@ -497,23 +499,12 @@ def check( else: raise ValueError("Failed to generate a valid vector for the query.") - # Create the vector search query - query = VectorRangeQuery( - vector=vector, - vector_field_name=CACHE_VECTOR_FIELD_NAME, - return_fields=self.return_fields, - distance_threshold=distance_threshold, - num_results=num_results, - return_score=True, - filter_expression=filter_expression, - dtype=self._vectorizer.dtype, - ) - # Search the cache! - cache_search_results = self._index.query(query) + query = self._build_check_query(vector, num_results, filter_expression) + candidates = self._index.query(query) redis_keys, cache_hits = self._process_cache_results( - cache_search_results, - return_fields, # type: ignore + self._within_threshold(candidates, distance_threshold), + return_fields, ) # Refresh TTL on all found keys @@ -581,6 +572,8 @@ async def acheck( # (exact match only), so only fall back on None. if distance_threshold is None: distance_threshold = self._distance_threshold + else: + self._validate_distance_threshold(distance_threshold) # Vectorize prompt if not provided if vector is None and prompt is not None: @@ -592,29 +585,70 @@ async def acheck( else: raise ValueError("Failed to generate a valid vector for the query.") - # Create the vector search query - query = VectorRangeQuery( + # Search the cache! + query = self._build_check_query(vector, num_results, filter_expression) + candidates = await aindex.query(query) + redis_keys, cache_hits = self._process_cache_results( + self._within_threshold(candidates, distance_threshold), + return_fields, + ) + + # Refresh TTL on all found keys async + await asyncio.gather(*[self.aexpire(key) for key in redis_keys]) + + return cache_hits + + def _build_check_query( + self, + vector: list[float], + num_results: int, + filter_expression: FilterExpression | None, + ) -> VectorQuery: + """Build the KNN lookup shared by :meth:`check` and :meth:`acheck`. + + :meth:`_within_threshold` then applies the distance threshold, so the + nearest distance is known even on a miss. + """ + return VectorQuery( vector=vector, vector_field_name=CACHE_VECTOR_FIELD_NAME, return_fields=self.return_fields, - distance_threshold=distance_threshold, num_results=num_results, return_score=True, filter_expression=filter_expression, dtype=self._vectorizer.dtype, ) - # Search the cache! - cache_search_results = await aindex.query(query) - redis_keys, cache_hits = self._process_cache_results( - cache_search_results, - return_fields, # type: ignore - ) + @staticmethod + def _within_threshold( + candidates: list[dict[str, Any]], distance_threshold: float + ) -> list[dict[str, Any]]: + """Keep the candidates whose distance is at most the threshold. - # Refresh TTL on all found keys async - await asyncio.gather(*[self.aexpire(key) for key in redis_keys]) + Inclusive, as a range query's radius is, at the precision of the + distance Redis returns (12 significant digits). A candidate with no distance + (an id-only row from the Redis 8.8+ expiry race) is kept so that + :meth:`_process_cache_results` skips it with its usual warning. + """ + return [ + candidate + for candidate in candidates + if (distance := candidate.get(VectorQuery.DISTANCE_ID)) is None + or float(distance) <= distance_threshold + ] - return cache_hits + @staticmethod + def _validate_distance_threshold(distance_threshold: float) -> None: + """Validate a per-call ``distance_threshold`` override. + + Raises: + TypeError: If the threshold is not a float or int. + ValueError: If the threshold is negative. + """ + if not isinstance(distance_threshold, (float, int)): + raise TypeError("distance_threshold must be of type float or int") + if distance_threshold < 0: + raise ValueError("distance_threshold must be non-negative") def _process_cache_results( self, diff --git a/tests/integration/test_llmcache.py b/tests/integration/test_llmcache.py index 4f5c59184..a0fa1e47e 100644 --- a/tests/integration/test_llmcache.py +++ b/tests/integration/test_llmcache.py @@ -596,6 +596,38 @@ async def test_acheck_with_zero_distance_threshold(cache): assert (await cache.acheck("what's the capital of france?")) != [] +@pytest.mark.parametrize("use_async", [False, True]) +async def test_check_threshold_is_inclusive(cache, vectorizer, use_async): + """A candidate at the threshold, as Redis reports its distance, hits; one + just beyond it does not.""" + prompts = [ + "what is the capital of france?", + "what's the capital of france?", + "how do I bake sourdough bread?", + ] + for prompt in prompts: + cache.store(prompt, "response", vector=vectorizer.embed(prompt)) + query = vectorizer.embed(prompts[0]) + + async def check(distance_threshold): + kwargs = dict( + vector=query, num_results=3, distance_threshold=distance_threshold + ) + return await cache.acheck(**kwargs) if use_async else cache.check(**kwargs) + + everything = await check(2) + distances = [hit["vector_distance"] for hit in everything] + assert len(distances) == 3 and distances == sorted(distances) + + at_threshold = await check(distances[1]) + assert [hit["key"] for hit in at_threshold] == [ + hit["key"] for hit in everything[:2] + ] + + beyond_threshold = await check(distances[1] - 1e-6) + assert [hit["key"] for hit in beyond_threshold] == [everything[0]["key"]] + + def test_check_invalid_input(cache): with pytest.raises(ValueError): cache.check() @@ -603,6 +635,12 @@ def test_check_invalid_input(cache): with pytest.raises(TypeError): cache.check(prompt="test", return_fields="bad value") + with pytest.raises(ValueError): + cache.check(prompt="test", distance_threshold=-0.1) + + with pytest.raises(TypeError): + cache.check(prompt="test", distance_threshold="0.1") + @pytest.mark.asyncio async def test_async_check_invalid_input(cache): @@ -612,6 +650,12 @@ async def test_async_check_invalid_input(cache): with pytest.raises(TypeError): await cache.acheck(prompt="test", return_fields="bad value") + with pytest.raises(ValueError): + await cache.acheck(prompt="test", distance_threshold=-0.1) + + with pytest.raises(TypeError): + await cache.acheck(prompt="test", distance_threshold="0.1") + def test_bad_connection_info(vectorizer, worker_id): with pytest.raises(ConnectionError): diff --git a/tests/unit/test_extension_result_hardening.py b/tests/unit/test_extension_result_hardening.py index a8b10a56b..717738643 100644 --- a/tests/unit/test_extension_result_hardening.py +++ b/tests/unit/test_extension_result_hardening.py @@ -61,6 +61,15 @@ def test_process_cache_results_bad_then_valid_ordering(): assert redis_keys == ["llmcache:key1"] # only the valid hit's key +def test_within_threshold_leaves_id_only_hit_for_the_skip_path(): + """The threshold filter must not raise on a hit with no distance.""" + far_hit = {**_valid_cache_hit(), "vector_distance": "0.3"} + candidates = [{"id": "llmcache:expired"}, far_hit] + assert SemanticCache._within_threshold(candidates, 0.2) == [ + {"id": "llmcache:expired"} + ] + + @pytest.mark.parametrize("as_text", [True, False]) def test_format_context_drops_id_only_message(as_text, caplog): """A message with 'id' but no entry_id/session_tag hits the KeyError arm."""