From 2de151232c3ff5c14f83b966f2fb4253eee8c78c Mon Sep 17 00:00:00 2001 From: Thomas Tomy <69713148+thomastomy5@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:10:26 +0530 Subject: [PATCH 1/3] added limit of 5k for exact search --- .../server/apps/floware/floware/config.ini | 3 + .../controllers/rag_retreival_controller.py | 27 +++++++- .../queries/generate_query.py | 49 ++++++++++++++ .../services/image_rag_retrieve.py | 65 +++++++++++++++++++ 4 files changed, 143 insertions(+), 1 deletion(-) 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..d192bbfb 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,25 @@ async def _resolve_image_data( return (image_data_b64, None) +def _resolve_exact_match_candidate_cap(config: dict) -> int: + """ + Resolve the exact-match candidate-count safety cap: a soft, + env-configurable limit (`knowledge_base.exact_match_max_candidates` / + `KB_EXACT_MATCH_MAX_CANDIDATES`, see config.ini) clamped to a hardcoded, + non-configurable ceiling (`EXACT_MATCH_HARD_CEILING`) so a misconfigured + or accidentally-widened env var can never fully disable the guard. + """ + 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 +418,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 +434,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..cf2b6db1 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,55 @@ 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 + `filterN`/`document_date`/`created_at` window used by + `get_image_embedding_dino_exact_match`. Intended as a pre-flight + safety check -- callers should run this first and compare the + result against a candidate-count cap before running the (much more + expensive) exact-match brute-force distance computation, since this + query only touches the real, indexed columns on + `knowledge_base_documents` and never reads `knowledge_base_embeddings` + or computes any vector distance. + """ + 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..d19e89d5 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,19 @@ from db_repo_module.repositories.sql_alchemy_repository import SQLAlchemyRepository from sqlalchemy.exc import SQLAlchemyError +# Absolute, non-configurable ceiling on how many candidate documents +# `exact_match_dino` will brute-force score in a single request, regardless +# of what `max_candidates`/config says. This exists so a misconfigured or +# accidentally-widened env var (`KB_EXACT_MATCH_MAX_CANDIDATES`) can never +# fully disable the safety guard -- see `exact_match_dino` below. +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 +129,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 +147,58 @@ 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 doing any of that, this first runs a cheap, index-only count + of how many documents match the given filters and rejects the + request with a 422 if that count exceeds `max_candidates` (itself + clamped to `EXACT_MATCH_HARD_CEILING`). This protects the DB from an + accidentally wide date range / filter turning into a brute-force + distance computation over a huge fraction of a large KB (see + `get_filtered_document_count_query`) -- the count check runs before + the outbound call to the inference service too, so an + already-doomed request fails fast instead of paying for an + embedding computation it won't use. """ + 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: From df65a24891a373b1a3b1a870914a4de8db28971d Mon Sep 17 00:00:00 2001 From: Thomas Tomy <69713148+thomastomy5@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:14:31 +0530 Subject: [PATCH 2/3] saving --- .../controllers/rag_retreival_controller.py | 8 +----- .../queries/generate_query.py | 11 +++----- .../services/image_rag_retrieve.py | 26 +++++-------------- 3 files changed, 11 insertions(+), 34 deletions(-) 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 d192bbfb..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 @@ -214,13 +214,7 @@ async def _resolve_image_data( def _resolve_exact_match_candidate_cap(config: dict) -> int: - """ - Resolve the exact-match candidate-count safety cap: a soft, - env-configurable limit (`knowledge_base.exact_match_max_candidates` / - `KB_EXACT_MATCH_MAX_CANDIDATES`, see config.ini) clamped to a hardcoded, - non-configurable ceiling (`EXACT_MATCH_HARD_CEILING`) so a misconfigured - or accidentally-widened env var can never fully disable the guard. - """ + """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( 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 cf2b6db1..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 @@ -504,14 +504,9 @@ def get_filtered_document_count_query( ) -> Tuple[str, Dict[str, Any]]: """ Cheap, index-only count of documents in `kb_id` matching the same - `filterN`/`document_date`/`created_at` window used by - `get_image_embedding_dino_exact_match`. Intended as a pre-flight - safety check -- callers should run this first and compare the - result against a candidate-count cap before running the (much more - expensive) exact-match brute-force distance computation, since this - query only touches the real, indexed columns on - `knowledge_base_documents` and never reads `knowledge_base_embeddings` - or computes any vector distance. + 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, 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 d19e89d5..acb87548 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,17 +9,11 @@ from db_repo_module.repositories.sql_alchemy_repository import SQLAlchemyRepository from sqlalchemy.exc import SQLAlchemyError -# Absolute, non-configurable ceiling on how many candidate documents -# `exact_match_dino` will brute-force score in a single request, regardless -# of what `max_candidates`/config says. This exists so a misconfigured or -# accidentally-widened env var (`KB_EXACT_MATCH_MAX_CANDIDATES`) can never -# fully disable the safety guard -- see `exact_match_dino` below. +# 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 candidate cap when the caller doesn't supply one. DEFAULT_EXACT_MATCH_MAX_CANDIDATES = 1_000 @@ -148,16 +142,10 @@ async def exact_match_dino( HNSW index -- see that method's docstring -- so scores returned here are always exact, not approximate. - Before doing any of that, this first runs a cheap, index-only count - of how many documents match the given filters and rejects the - request with a 422 if that count exceeds `max_candidates` (itself - clamped to `EXACT_MATCH_HARD_CEILING`). This protects the DB from an - accidentally wide date range / filter turning into a brute-force - distance computation over a huge fraction of a large KB (see - `get_filtered_document_count_query`) -- the count check runs before - the outbound call to the inference service too, so an - already-doomed request fails fast instead of paying for an - embedding computation it won't use. + 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, From b21cb86390f8b98615d9633bacf14c6ce9d42de2 Mon Sep 17 00:00:00 2001 From: Thomas Tomy <69713148+thomastomy5@users.noreply.github.com> Date: Sat, 5 Sep 2026 12:15:24 +0530 Subject: [PATCH 3/3] saving v2 --- .../knowledge_base_module/services/image_rag_retrieve.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 acb87548..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 @@ -13,7 +13,10 @@ # exceed this, so a misconfiguration can't fully disable the safety guard. EXACT_MATCH_HARD_CEILING = 5_000 -# Default candidate cap when the caller doesn't supply one. +# 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