diff --git a/redisvl/extensions/cache/llm/schema.py b/redisvl/extensions/cache/llm/schema.py index c0cdd4e2..4442b7ac 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 d3253204..781731df 100644 --- a/redisvl/extensions/cache/llm/semantic.py +++ b/redisvl/extensions/cache/llm/semantic.py @@ -22,7 +22,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 @@ -485,6 +485,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: @@ -496,23 +498,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 @@ -579,6 +570,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: @@ -590,29 +583,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 + await self._aexpire_many(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 - await self._aexpire_many(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 cde70405..9aae294e 100644 --- a/tests/integration/test_llmcache.py +++ b/tests/integration/test_llmcache.py @@ -624,6 +624,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() @@ -631,6 +663,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): @@ -640,6 +678,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 a8b10a56..71773864 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."""