diff --git a/wavefront/server/apps/floware/floware/config.ini b/wavefront/server/apps/floware/floware/config.ini index b60ef9a6..1ad67d10 100644 --- a/wavefront/server/apps/floware/floware/config.ini +++ b/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} diff --git a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/controllers/rag_retreival_controller.py b/wavefront/server/modules/knowledge_base_module/knowledge_base_module/controllers/rag_retreival_controller.py index 7885bec1..ac615d40 100644 --- a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/controllers/rag_retreival_controller.py +++ b/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) diff --git a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/queries/generate_query.py b/wavefront/server/modules/knowledge_base_module/knowledge_base_module/queries/generate_query.py index 8ea1321d..057ede37 100644 --- a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/queries/generate_query.py +++ b/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, diff --git a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/services/image_rag_retrieve.py b/wavefront/server/modules/knowledge_base_module/knowledge_base_module/services/image_rag_retrieve.py index 555d7cd2..a240c451 100644 --- a/wavefront/server/modules/knowledge_base_module/knowledge_base_module/services/image_rag_retrieve.py +++ b/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: