diff --git a/wavefront/server/apps/floware/floware/config.ini b/wavefront/server/apps/floware/floware/config.ini index a4204048..2ecb6660 100644 --- a/wavefront/server/apps/floware/floware/config.ini +++ b/wavefront/server/apps/floware/floware/config.ini @@ -144,6 +144,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 586d0a07..9ba7fb40 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 @@ -210,6 +214,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( @@ -396,6 +413,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, @@ -411,6 +429,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 bd743c73..7e0b40d9 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 @@ -420,7 +420,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, @@ -437,20 +437,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, @@ -469,27 +470,26 @@ 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, - d.document_date::text AS document_date, - 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, + d.document_date::text AS document_date, + 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 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..7364172b 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,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,