limiting the exact limit search window to 5000 records

wavefront/server/apps/floware/floware/config.ini

@@ -157,6 +157,9 @@ inactive_days_threshold=${INACTIVE_DAYS_THRESHOLD:60}
 [model]
 inference_service_url=${INFERENCE_SERVICE_URL}
 
+[knowledge_base]
+exact_match_max_candidates=${KB_EXACT_MATCH_MAX_CANDIDATES:1000}
+
 [embedding_url]
 embedding_service_url=${EMBEDDING_SERVICE_URL}
 

wavefront/server/modules/knowledge_base_module/knowledge_base_module/controllers/rag_retreival_controller.py

@@ -22,7 +22,11 @@
     NewInference,
 )
 from knowledge_base_module.services.kb_rag_retrieve import KBRagResponse
-from knowledge_base_module.services.image_rag_retrieve import ImageRagRetrieve
+from knowledge_base_module.services.image_rag_retrieve import (
+    DEFAULT_EXACT_MATCH_MAX_CANDIDATES,
+    EXACT_MATCH_HARD_CEILING,
+    ImageRagRetrieve,
+)
 from flo_cloud.cloud_storage import CloudStorageManager
 from pydantic import BaseModel, Field
 from datetime import datetime
@@ -209,6 +213,19 @@ async def _resolve_image_data(
     return (image_data_b64, None)
 
 
+def _resolve_exact_match_candidate_cap(config: dict) -> int:
+    """Resolve the exact-match candidate cap from config, clamped to `EXACT_MATCH_HARD_CEILING`."""
+    knowledge_base_config = (config or {}).get('knowledge_base') or {}
+    try:
+        configured_cap = int(
+            knowledge_base_config.get('exact_match_max_candidates')
+            or DEFAULT_EXACT_MATCH_MAX_CANDIDATES
+        )
+    except (TypeError, ValueError):
+        configured_cap = DEFAULT_EXACT_MATCH_MAX_CANDIDATES
+    return min(configured_cap, EXACT_MATCH_HARD_CEILING)
+
+
 @rag_retrieval_router.post('/v1/knowledge-base/{kb_id}/retrieve')
 @inject
 async def retrieve_query(
@@ -395,6 +412,7 @@ async def retrieve_query(
         if error_response is not None:
             return error_response
         inference_url = config['model']['inference_service_url']
+        exact_match_max_candidates = _resolve_exact_match_candidate_cap(config)
         retrieved_docs = await image_rag_retrieval.exact_match_dino(
             image_data,
             inference_url,
@@ -410,6 +428,7 @@ async def retrieve_query(
             filter6,
             created_at_start,
             created_at_end,
+            max_candidates=exact_match_max_candidates,
         )
         retrieved_docs = convert_uuids_to_str(retrieved_docs)
         match_count = len(retrieved_docs)

wavefront/server/modules/knowledge_base_module/knowledge_base_module/queries/generate_query.py

@@ -415,7 +415,7 @@ def get_image_embedding_dino_exact_match(
         filter1: str,
         document_date_start,
         document_date_end,
-        threshold: float,
+        max_candidates: int,
         filter2: Optional[str] = None,
         filter3: Optional[str] = None,
         filter4: Optional[str] = None,
@@ -432,20 +432,21 @@ def get_image_embedding_dino_exact_match(
         generic, caller-defined columns -- see `KnowledgeBaseDocuments` --
         this query has no notion of what they mean semantically.
 
-        Deliberately has no `ORDER BY`/`LIMIT` tied to the `<=>` distance
-        expression anywhere -- that is what would let Postgres route the
-        query through the HNSW index (`ix_kbe_embedding_vector_1_hnsw_cosine`)
-        for an *approximate* top-K search. Here we instead pre-filter to a
-        small candidate set via the real, indexed `filterN`/`document_date`
-        columns, then compute an exact cosine distance for every one of those
-        rows and only keep the ones above `threshold` -- so results are exact,
-        not approximate, and the count of matches is precise.
-
-        The threshold check is a plain scalar comparison on the computed
-        `dino_score`, so it has to live in an outer query over a subquery
-        (Postgres doesn't allow referencing a `SELECT`-list alias in a
-        same-level `WHERE`) -- it still runs after every distance in the
-        candidate set has already been computed exactly.
+        Candidates are capped via `ORDER BY d.id LIMIT :fetch_limit`, where
+        `fetch_limit` is computed here as `max_candidates + 1`. Ordering by
+        `d.id` -- a plain, non-vector column -- rather than the `<=>`
+        distance expression keeps this from ever engaging the HNSW index
+        (`ix_kbe_embedding_vector_1_hnsw_cosine`), which only gets used when
+        `ORDER BY`/`LIMIT` is tied directly to a `<=>` expression. That means
+        results here stay exact, not approximate, for whatever candidate set
+        the `LIMIT` lets through. The "+1" lets the caller detect when the
+        true candidate count exceeded `max_candidates` (i.e.
+        `len(rows) == max_candidates + 1`) without a separate `COUNT` query.
+
+        Deliberately has no `dino_score` threshold filter in this query --
+        that comparison is left to the caller so it can distinguish "no
+        candidates matched" from "the candidate set was truncated" using the
+        raw (pre-threshold) row count returned here.
         """
         filter_columns_clause, filter_columns_params = self.build_filter_columns_clause(
             filter1,
@@ -464,26 +465,25 @@ def get_image_embedding_dino_exact_match(
         params: Dict[str, Any] = {
             'query_embedding': query_embeddings,
             'kb_id': str(kb_id),
-            'threshold': threshold,
+            'fetch_limit': max_candidates + 1,
             **filter_columns_params,
         }
 
         sql_query = f"""
-        SELECT * FROM (
-            SELECT
-                e.id AS embedding_id,
-                d.id AS document_id,
-                d.file_path,
-                d.file_name,
-                d.knowledge_base_id,
-                d.metadata_value,
-                1 - ((e.embedding_vector_1::vector(1024)) <=> :query_embedding ::vector(1024)) AS dino_score
-            FROM {KnowledgeBaseEmbeddings.__tablename__} e
-            JOIN {KnowledgeBaseDocuments.__tablename__} d ON e.document_id = d.id
-            WHERE d.knowledge_base_id = :kb_id
-                {filter_columns_clause}
-        ) scored
-        WHERE dino_score > :threshold
+        SELECT
+            e.id AS embedding_id,
+            d.id AS document_id,
+            d.file_path,
+            d.file_name,
+            d.knowledge_base_id,
+            d.metadata_value,
+            1 - ((e.embedding_vector_1::vector(1024)) <=> :query_embedding ::vector(1024)) AS dino_score
+        FROM {KnowledgeBaseEmbeddings.__tablename__} e
+        JOIN {KnowledgeBaseDocuments.__tablename__} d ON e.document_id = d.id
+        WHERE d.knowledge_base_id = :kb_id
+            {filter_columns_clause}
+        ORDER BY d.id
+        LIMIT :fetch_limit
         """
 
         return sql_query, params

wavefront/server/modules/knowledge_base_module/knowledge_base_module/services/image_rag_retrieve.py

@@ -9,6 +9,14 @@
 from db_repo_module.repositories.sql_alchemy_repository import SQLAlchemyRepository
 from sqlalchemy.exc import SQLAlchemyError
 
+# Hard ceiling on exact_match_dino's candidate count -- no config/env var can
+# exceed this, so a misconfiguration can't fully disable the safety guard.
+EXACT_MATCH_HARD_CEILING = 5_000
+
+# Fallback candidate cap used when the caller doesn't supply one Deliberately conservative; 
+# tune based on real p95 latency measurements against the target KB size.
+DEFAULT_EXACT_MATCH_MAX_CANDIDATES = 1_000
+
 
 @dataclass
 class ImageMatch:
@@ -116,6 +124,7 @@ async def exact_match_dino(
         filter6: Optional[str] = None,
         created_at_start=None,
         created_at_end=None,
+        max_candidates: Optional[int] = None,
     ) -> list[dict]:
         """
         Exact (non-ANN) DINO similarity match, restricted to documents in
@@ -133,7 +142,20 @@ async def exact_match_dino(
         (`QueryGenerator.get_image_embedding_dino_exact_match`) never engages the
         HNSW index -- see that method's docstring -- so scores returned here
         are always exact, not approximate.
+
+        The candidate set is capped directly in SQL via `ORDER BY d.id LIMIT
+        max_candidates + 1` (clamped to `EXACT_MATCH_HARD_CEILING`), 
+        so an oversized candidate set never gets fully brute-forced. 
+        The extra "+1" row lets us detect an overflow  (raise a 422) 
+        using just the row count of the one query already run. 
+        The `threshold` filter is applied in Python (after the overflow check) 
+        rather than in SQL.
         """
+        effective_cap = min(
+            max_candidates or DEFAULT_EXACT_MATCH_MAX_CANDIDATES,
+            EXACT_MATCH_HARD_CEILING,
+        )
+
         data = {'image_data': image_data}
         internal_api_url = f'{inference_url}/inference/v1/query/embeddings'
         try:
@@ -176,7 +198,7 @@ async def exact_match_dino(
                     filter1,
                     document_date_start,
                     document_date_end,
-                    threshold,
+                    effective_cap,
                     filter2,
                     filter3,
                     filter4,
@@ -186,7 +208,7 @@ async def exact_match_dino(
                     created_at_end,
                 )
             )
-            return await self.knowledge_base_embeddings_repository.execute_query(
+            raw_rows = await self.knowledge_base_embeddings_repository.execute_query(
                 sql_query,
                 query_params,
             )
@@ -195,6 +217,18 @@ async def exact_match_dino(
                 f'Failed to execute the query for exact match retrieval: {e}'
             )
 
+        if len(raw_rows) > effective_cap:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f'More than {effective_cap} documents match the given filters, '
+                    'which exceeds the exact-match safety limit. Narrow your date '
+                    'range or filters and try again.'
+                ),
+            )
+
+        return [row for row in raw_rows if row['dino_score'] > threshold]
+
     async def image_retrieve_clip(
         self,
         clip_embedding,