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

@@ -488,6 +488,50 @@ def get_image_embedding_dino_exact_match(
 
         return sql_query, params
 
+    def get_filtered_document_count_query(
+        self,
+        kb_id: str,
+        filter1: Optional[str] = None,
+        document_date_start: Optional[Any] = None,
+        document_date_end: Optional[Any] = None,
+        filter2: Optional[str] = None,
+        filter3: Optional[str] = None,
+        filter4: Optional[str] = None,
+        filter5: Optional[str] = None,
+        filter6: Optional[str] = None,
+        created_at_start: Optional[Any] = None,
+        created_at_end: Optional[Any] = None,
+    ) -> Tuple[str, Dict[str, Any]]:
+        """
+        Cheap, index-only count of documents in `kb_id` matching the same
+        filters as `get_image_embedding_dino_exact_match`. Used as a
+        pre-flight check against a candidate-count cap before running that
+        much more expensive brute-force query.
+        """
+        filter_columns_clause, params = self.build_filter_columns_clause(
+            filter1,
+            filter2,
+            filter3,
+            filter4,
+            filter5,
+            filter6,
+            document_date_start,
+            document_date_end,
+            table_alias='d',
+            created_at_start=created_at_start,
+            created_at_end=created_at_end,
+        )
+        params['kb_id'] = str(kb_id)
+
+        sql_query = f"""
+        SELECT COUNT(*) AS candidate_count
+        FROM {KnowledgeBaseDocuments.__tablename__} d
+        WHERE d.knowledge_base_id = :kb_id
+            {filter_columns_clause}
+        """
+
+        return sql_query, params
+
     def get_documents_list_query(
         self,
         kb_id: str,

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

@@ -9,6 +9,16 @@
 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 (e.g. older
+# callers, or config missing the `knowledge_base.exact_match_max_candidates`
+# key). 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 +126,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 +144,52 @@ 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.
+
+        Before any of that, runs a cheap count of matching documents and
+        rejects with a 422 if it exceeds `max_candidates` (clamped to
+        `EXACT_MATCH_HARD_CEILING`), so an oversized candidate set fails
+        fast instead of brute-forcing distances over it.
         """
+        effective_cap = min(
+            max_candidates or DEFAULT_EXACT_MATCH_MAX_CANDIDATES,
+            EXACT_MATCH_HARD_CEILING,
+        )
+        try:
+            count_query, count_params = (
+                self.query_generator.get_filtered_document_count_query(
+                    kb_id,
+                    filter1,
+                    document_date_start,
+                    document_date_end,
+                    filter2,
+                    filter3,
+                    filter4,
+                    filter5,
+                    filter6,
+                    created_at_start,
+                    created_at_end,
+                )
+            )
+            count_rows = await self.knowledge_base_embeddings_repository.execute_query(
+                count_query,
+                count_params,
+            )
+        except SQLAlchemyError as e:
+            raise RuntimeError(
+                f'Failed to execute the candidate-count query for exact match retrieval: {e}'
+            )
+
+        candidate_count = int(count_rows[0]['candidate_count']) if count_rows else 0
+        if candidate_count > effective_cap:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=(
+                    f'{candidate_count} documents match the given filters, which '
+                    f'exceeds the exact-match safety limit of {effective_cap}. '
+                    'Narrow your date range or filters and try again.'
+                ),
+            )
+
         data = {'image_data': image_data}
         internal_api_url = f'{inference_url}/inference/v1/query/embeddings'
         try: