perf: run SemanticCache.check() as KNN with a client-side threshold (FDE-85)

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",
                     },
                 },

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,

tests/integration/test_llmcache.py

@@ -624,13 +624,51 @@ 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()
 
     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):

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."""