diff --git a/nemo_retriever/README.md b/nemo_retriever/README.md index 06af01044f..49b9476314 100644 --- a/nemo_retriever/README.md +++ b/nemo_retriever/README.md @@ -193,6 +193,12 @@ retriever ingest /your-example-dir \ --table-name nemo-retriever ``` +You do not need to choose a retrieval index mode for the normal workflow. The +default `index_mode=auto` creates a hybrid table (dense +vectors plus BM25/full-text search), and query mode `auto` uses that table +automatically. The explicit `dense`, `hybrid`, and `sparse` modes are advanced +overrides for experiments or specialized deployments. + Chunks land at `./lancedb/nemo-retriever`, which matches the storage settings used in [Run a recall query](#run-a-recall-query) below. With the `[local]` extra installed (see setup), defaults point at local-GPU extraction diff --git a/nemo_retriever/docs/cli/README.md b/nemo_retriever/docs/cli/README.md index 2f21575bbb..06eb22c85c 100644 --- a/nemo_retriever/docs/cli/README.md +++ b/nemo_retriever/docs/cli/README.md @@ -259,7 +259,7 @@ These options apply to `retriever ingest`, `retriever ingest local`, and | `--lancedb-uri` | `lancedb` | LanceDB database URI. | | `--table-name` | `nemo-retriever` | LanceDB table name. Must match query-time storage flags. | | `--overwrite/--append` | overwrite | Overwrite the table by default; use `--append` to add rows. | -| `--index-mode` | `dense` | Dense vector index by default; `hybrid` also builds BM25/FTS and `sparse` builds an FTS-only table. | +| `--index-mode` | `auto` | Recommended: leave this unset. `auto` creates a hybrid vector + BM25/FTS configuration for new tables and preserves an existing table on append. Use `dense`, `hybrid`, or `sparse` only for explicit experiments or specialized deployments. | | `--method` | planner default | PDF extraction method such as `pdfium` or `nemotron_parse`. | | `--extract-text`, `--extract-tables`, `--extract-charts` | planner default | Enable or disable extraction families. | | `--ocr-version` | planner default | OCR engine version for local extraction. | diff --git a/nemo_retriever/helm/README.md b/nemo_retriever/helm/README.md index 0cbf88db33..f529852fbb 100644 --- a/nemo_retriever/helm/README.md +++ b/nemo_retriever/helm/README.md @@ -339,11 +339,20 @@ The retriever service picks up the in-cluster ASR endpoint when `nimOperator.aud | `serviceConfig.agentic.requestTimeoutS` | `1800` | Gateway and MCP timeout for the multi-step agentic retrieval call. | | `serviceConfig.vectordb.enabled` | `true` | Deploy the LanceDB vectordb Pod. When `true` the chart **requires** a resolvable embed endpoint (refer to [VectorDB and the embed endpoint](#vectordb-and-the-embed-endpoint)); `helm install` / `helm upgrade` fails fast otherwise. | | `serviceConfig.vectordb.lancedbUri` | `/data/vectordb` | LanceDB on the vectordb Pod's PVC. | +| `serviceConfig.vectordb.indexMode` | `auto` | `auto`, `dense`, or `hybrid`. Fresh `auto` storage creates FTS and uses hybrid retrieval; persistent dense storage remains dense until `hybrid` is requested explicitly. | | `serviceConfig.vectordb.embedModel` | `nvidia/llama-nemotron-embed-vl-1b-v2` | Passed to vectordb + worker `embed_model_name`. | | `serviceConfig.vectordb.embedModelProviderPrefix` | `""` | Optional LiteLLM provider prefix prepended to the remote embed model name. | #### VectorDB and the embed endpoint { #vectordb-and-the-embed-endpoint } +The VectorDB storage default is `indexMode: auto`; most users should leave it +unchanged. A fresh table creates and waits for its FTS index after the first +write, while an existing dense table is left dense. Set +`serviceConfig.vectordb.indexMode=hybrid` only when you explicitly want to +upgrade an existing dense table. Incremental rows remain searchable through +LanceDB's unindexed-tail scan; the service performs incremental FTS maintenance +automatically and reports FTS and maintenance state from `/v1/health`. + The vectordb Pod's `/v1/query` handler embeds the incoming query text before searching LanceDB. It needs a NIM embedding endpoint to do that, and rendering the Deployment with an empty `--embed-endpoint` produces a diff --git a/nemo_retriever/helm/templates/deployment-vectordb.yaml b/nemo_retriever/helm/templates/deployment-vectordb.yaml index 1df0a86a3a..b55430745e 100644 --- a/nemo_retriever/helm/templates/deployment-vectordb.yaml +++ b/nemo_retriever/helm/templates/deployment-vectordb.yaml @@ -6,6 +6,10 @@ {{- $embedURL := include "nemo-retriever.nim.endpointURL" (dict "context" . "key" "vlm_embed" "serviceName" .Values.nimOperator.vlm_embed.nimServiceName "configKey" "embedInvokeUrl" "invokePath" "/v1/embeddings") -}} {{- $localEmbed := include "nemo-retriever.localEmbed.enabled" . | eq "true" -}} {{- $localModels := .Values.serviceConfig.localModels -}} +{{- $indexMode := .Values.serviceConfig.vectordb.indexMode | default "auto" -}} +{{- if not (has $indexMode (list "auto" "dense" "hybrid")) -}} +{{- fail "serviceConfig.vectordb.indexMode must be one of: auto, dense, hybrid" -}} +{{- end -}} {{- $agentic := .Values.serviceConfig.agentic -}} {{- $internalAuth := .Values.serviceConfig.vectordb.internalAuth -}} {{- /* @@ -77,6 +81,8 @@ spec: - {{ .Values.serviceConfig.vectordb.lancedbUri | quote }} - --table-name - {{ .Values.serviceConfig.vectordb.tableName | quote }} + - --index-mode + - {{ $indexMode | quote }} {{- if $embedURL }} - --embed-endpoint - {{ $embedURL | quote }} diff --git a/nemo_retriever/helm/values.yaml b/nemo_retriever/helm/values.yaml index 043433a3d7..c52a412416 100644 --- a/nemo_retriever/helm/values.yaml +++ b/nemo_retriever/helm/values.yaml @@ -691,6 +691,9 @@ serviceConfig: enabled: true lancedbUri: "/data/vectordb" tableName: "nemo_retriever" + # auto creates hybrid storage when fresh and preserves existing physical indexes. + # Use hybrid explicitly to upgrade an existing dense table by adding FTS. + indexMode: "auto" embedModel: "nvidia/llama-nemotron-embed-vl-1b-v2" embedModelProviderPrefix: "" # Optional dedicated gateway/worker-to-VectorDB authentication. When diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py index feefdb26b9..746495ba93 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/graph_commands.py @@ -222,7 +222,7 @@ def _graph_ingest_command( dedup_iou_threshold: opts.DedupIouThresholdOption = None, store_images_uri: opts.StoreImagesUriOption = None, overwrite: opts.OverwriteOption = True, - index_mode: opts.IndexModeOption = "dense", + index_mode: opts.IndexModeOption = "auto", ray_address: opts.RayAddressOption = None, ray_log_to_driver: opts.RayLogToDriverOption = None, page_elements_invoke_url: opts.PageElementsInvokeUrlOption = None, diff --git a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py index 8aae3b5ec8..e29ae7924b 100644 --- a/nemo_retriever/src/nemo_retriever/cli/ingest/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/ingest/options.py @@ -200,8 +200,9 @@ typer.Option( "--index-mode", help=( - "LanceDB index mode: dense, hybrid, or sparse. Dense is vector-only; hybrid also builds " - "BM25/FTS; sparse skips dense embedding and writes an FTS-only table." + "Recommended: leave unset. Auto creates a hybrid table for new indexes and preserves an existing " + "table on append. Dense, hybrid, and sparse are advanced overrides for experiments or specialized " + "deployments." ), ), ] diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 3fabce18d3..168dc7d892 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -136,8 +136,8 @@ typer.Option( "--retrieval-mode", help=( - "Expert LanceDB retrieval mode: auto, dense, hybrid, or sparse. Default auto inspects the table " - "and chooses the supported mode." + "Advanced override: auto, dense, hybrid, or sparse. Leave at auto to inspect the table and use " + "the supported default mode." ), ), ] diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index fbd543abf0..10b7eed6e7 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/nemo_retriever/src/nemo_retriever/common/vdb/README.md @@ -123,9 +123,9 @@ For `vdb_op="lancedb"`, **`LanceDB.retrieval`**: - Opens the table with `lancedb.connect(table_path).open_table(table_name)`. - For dense retrieval, each query vector uses **`table.search([vector], vector_column_name=..., **search_kwargs)`**, optional **`.where(where_clause)`** (Lance / DataFusion SQL; `metadata` / `source` are stored as JSON strings), then **`.limit(top_k).refine_factor(...).nprobes(...)`**. -- For hybrid retrieval, callers pass `hybrid=True` plus `query_texts` aligned with the vectors. LanceDB uses **`table.search(query_type="hybrid", vector_column_name=..., fts_columns="text").vector(vector).text(query_text)`** before applying the same `where`, limit, refine, probe, and select handling. +- For hybrid retrieval, callers pass `hybrid=True` plus `query_texts` aligned with the vectors. LanceDB uses **`table.search(query_type="hybrid", vector_column_name=..., fts_columns="text").vector(vector).text(query_text)`** before applying the same `where`, limit, refine, probe, and select handling. Product query paths also pass the shared weighted-RRF policy (`candidate_depth=50`, `dense_weight=0.8`, `rrf_k=10`), then truncate the fused ranking to `top_k`. Direct low-level callers opt into that behavior explicitly with `hybrid_fusion=HybridFusionPolicy(...)`. -Notable kwargs: `top_k`, `refine_factor`, `n_probe` / `nprobes`, `where` or `_filter`, `table_path`, `table_name`, `search_kwargs`, `hybrid`, and `query_texts`. `query_texts` is stripped from constructor kwargs and forwarded only for retrieval calls whose effective mode is hybrid. +Notable kwargs: `top_k`, `refine_factor`, `n_probe` / `nprobes`, `where` or `_filter`, `table_path`, `table_name`, `search_kwargs`, `hybrid`, `query_texts`, and `hybrid_fusion`. `query_texts` is stripped from constructor kwargs and forwarded only for retrieval calls whose effective mode is hybrid. Example of **direct** operator use (you supply vectors): diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/hybrid_fusion.py b/nemo_retriever/src/nemo_retriever/common/vdb/hybrid_fusion.py new file mode 100644 index 0000000000..d16c00b85b --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/vdb/hybrid_fusion.py @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed rank-fusion policy for LanceDB hybrid retrieval.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from typing import Any + +import pyarrow as pa +from lancedb.rerankers.base import Reranker + + +@dataclass(frozen=True) +class HybridFusionPolicy: + """Candidate depth and weighted-RRF parameters for one hybrid query.""" + + candidate_depth: int + dense_weight: float + rrf_k: int + + def __post_init__(self) -> None: + if self.candidate_depth <= 0: + raise ValueError("candidate_depth must be greater than zero") + if not 0.0 <= self.dense_weight <= 1.0: + raise ValueError("dense_weight must be between zero and one") + if self.rrf_k <= 0: + raise ValueError("rrf_k must be greater than zero") + + +DEFAULT_HYBRID_FUSION_POLICY = HybridFusionPolicy(candidate_depth=50, dense_weight=0.8, rrf_k=10) + + +class WeightedRRFReranker(Reranker): + """Fuse dense and FTS ranks while preferring dense order for score ties.""" + + def __init__(self, policy: HybridFusionPolicy) -> None: + super().__init__(return_score="relevance") + self.policy = policy + + def rerank_hybrid( + self, + query: str, + vector_results: pa.Table, + fts_results: pa.Table, + ) -> pa.Table: + del query + vector_ids = vector_results["_rowid"].to_pylist() if len(vector_results) else [] + fts_ids = fts_results["_rowid"].to_pylist() if len(fts_results) else [] + scores: defaultdict[Any, float] = defaultdict(float) + for weight, row_ids in ( + (self.policy.dense_weight, vector_ids), + (1.0 - self.policy.dense_weight, fts_ids), + ): + for rank, row_id in enumerate(row_ids, start=1): + scores[row_id] += weight / (self.policy.rrf_k + rank) + + combined = self.merge_results(vector_results, fts_results) + row_ids = combined["_rowid"].to_pylist() + combined = combined.append_column( + "_relevance_score", + pa.array([scores[row_id] for row_id in row_ids], type=pa.float64()), + ) + combined = combined.append_column("_fusion_order", pa.array(range(len(combined)), type=pa.int64())) + combined = combined.sort_by([("_relevance_score", "descending"), ("_fusion_order", "ascending")]) + combined = combined.drop_columns(["_fusion_order"]) + return self._keep_relevance_score(combined) diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py index ed538c0492..0cabeff407 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py @@ -9,7 +9,8 @@ import time from collections.abc import Iterable, Sequence -from datetime import timedelta +from contextlib import nullcontext +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import Any, Final, FrozenSet @@ -32,6 +33,10 @@ CollectionWriteResult, VDB, ) +from nemo_retriever.common.vdb.hybrid_fusion import ( + HybridFusionPolicy, + WeightedRRFReranker, +) from nemo_retriever.common.vdb.lancedb_capabilities import inspect_lancedb_table_object from nemo_retriever.common.vdb.lancedb_schema import ( build_lancedb_row, @@ -49,6 +54,10 @@ _NEMO_RETRIEVER_RETRIEVAL_MODE_METADATA_KEY: Final[bytes] = b"nemo_retriever.retrieval_mode" _EMBEDDING_MODEL_METADATA_KEY: Final[bytes] = b"nemo_retriever.embedding_model_name" _EMBEDDING_MODEL_REVISION_METADATA_KEY: Final[bytes] = b"nemo_retriever.embedding_model_revision" +# Appended rows remain searchable through LanceDB's unindexed-tail scan until +# optimize() folds them into FTS. These thresholds follow its recommended cadence. +_SERVICE_OPTIMIZE_WRITE_THRESHOLD: Final[int] = 20 +_SERVICE_OPTIMIZE_ROW_THRESHOLD: Final[int] = 100_000 def _normalize_on_bad_vectors(value: str) -> str: @@ -575,6 +584,7 @@ def __init__( ): create_index = kwargs.pop("create_index", None) service_table_schema = bool(kwargs.pop("_service_table_schema", False)) + service_index_mode = kwargs.pop("_service_index_mode", None) if build_index is None: build_index = True if create_index is None else bool(create_index) elif create_index is not None and bool(create_index) != bool(build_index): @@ -603,10 +613,83 @@ def __init__( self.validate_vector_length = bool(validate_vector_length) self.expiration_cleanup_enabled = bool(expiration_cleanup_enabled) self._service_table_schema = service_table_schema + self._service_index_mode = str(service_index_mode) if service_index_mode is not None else None + self._service_write_lock = threading.Lock() + self._writes_since_optimize = 0 + # Process-local status is diagnostic; persisted index statistics are authoritative. + self._last_optimization: dict[str, Any] = { + "status": "never", + "completed_at": None, + "error": None, + } self._collection_store: Any | None = None self._collection_store_init_failed = False self._collection_store_lock = threading.Lock() super().__init__(**kwargs) + if self._service_index_mode is not None and self.hybrid: + db = lancedb.connect(uri=self.uri) + if self.table_name in db.list_tables().tables: + self._ensure_fts_index(db.open_table(self.table_name)) + + @staticmethod + def _is_fts_index(index: Any) -> bool: + index_type = str(getattr(index, "index_type", "") or "").lower() + index_name = str(getattr(index, "name", "") or "").lower() + return index_type == "fts" or "text" in index_name or "fts" in index_name + + def _ensure_fts_index(self, table: Any) -> None: + if inspect_lancedb_table_object(table).has_fts: + return + started = time.perf_counter() + table.create_fts_index("text", language=self.fts_language, replace=True) + names = [index.name for index in table.list_indices() if self._is_fts_index(index)] + if names: + table.wait_for_index(names, timeout=timedelta(seconds=600)) + _record_timing("lancedb.fts_index_ready", time.perf_counter() - started) + + def _fts_unindexed_rows(self, table: Any) -> int | None: + values: list[int] = [] + for index in table.list_indices(): + if not self._is_fts_index(index): + continue + try: + stats = table.index_stats(index.name) + except Exception: + logger.debug("Unable to read LanceDB index stats for %s", index.name, exc_info=True) + continue + value = ( + stats.get("num_unindexed_rows") + if isinstance(stats, dict) + else getattr(stats, "num_unindexed_rows", None) + ) + if value is not None: + values.append(int(value)) + return sum(values) if values else None + + def _optimize_service_table_if_due(self, table: Any) -> None: + # Persisted FTS statistics keep the row threshold valid across restarts. + unindexed_rows = self._fts_unindexed_rows(table) + if self._writes_since_optimize < _SERVICE_OPTIMIZE_WRITE_THRESHOLD and ( + unindexed_rows is None or unindexed_rows < _SERVICE_OPTIMIZE_ROW_THRESHOLD + ): + return + try: + table.optimize() + except Exception as exc: + self._last_optimization = { + "status": "error", + "completed_at": datetime.now(timezone.utc).isoformat(), + "error": str(exc), + } + logger.exception("LanceDB optimization failed for table %r", self.table_name) + # Preserve the trigger state so a later write retries transient failures. + return + self._writes_since_optimize = 0 + self._last_optimization = { + "status": "ok", + "completed_at": datetime.now(timezone.utc).isoformat(), + "error": None, + } def _get_collection_store(self) -> Any: """Lazily initialize collection catalogs only when a collection API is used.""" @@ -792,16 +875,21 @@ def health(self) -> dict[str, Any]: total_rows = 0 effective_mode: str | None = None retrieval_strategies: list[str] = [] + table: Any | None = None + capabilities = None if table_exists: try: - total_rows = int(db.open_table(self.table_name).count_rows()) + table = db.open_table(self.table_name) + total_rows = int(table.count_rows()) except Exception: logger.warning( "Failed to count rows in the default LanceDB table", exc_info=True, ) try: - capabilities = inspect_lancedb_table_object(db.open_table(self.table_name)) + if table is None: + table = db.open_table(self.table_name) + capabilities = inspect_lancedb_table_object(table) mode = capabilities.retrieval_mode if mode in {"dense", "hybrid"}: effective_mode = str(mode) @@ -819,8 +907,22 @@ def health(self) -> dict[str, Any]: raise RuntimeError("Collection catalog initialization failed") store = self._collection_store collection_health = store.health() if store is not None else LanceDBCollectionStore.empty_health() + service_health: dict[str, Any] = {} + if self._service_index_mode is not None: + service_health = { + "configured_index_mode": self._service_index_mode, + "effective_index_mode": effective_mode, + "fts_present": bool(capabilities and capabilities.has_fts), + "fts_unindexed_rows": ( + self._fts_unindexed_rows(table) + if table is not None and capabilities is not None and capabilities.has_fts + else None + ), + "last_optimization": dict(self._last_optimization), + } return { **collection_health, + **service_health, "total_rows": total_rows, "table_exists": table_exists, "effective_retrieval_mode": effective_mode, @@ -1056,24 +1158,38 @@ def write_to_index( def run(self, records): """Orchestrate index creation and data ingestion.""" - table = self.create_index(records=records, table_name=self.table_name) - if self.build_index: - self.write_to_index( - records, - table=table, - index_type=self.index_type, - metric=self.metric, - num_partitions=self.num_partitions, - num_sub_vectors=self.num_sub_vectors, - hybrid=self.hybrid, - sparse=self.sparse, - fts_language=self.fts_language, - ) - else: - logger.info( - "Skipping LanceDB index creation for table %r because build_index=False.", - self.table_name, - ) + service_write = self._service_index_mode is not None + lock = self._service_write_lock if service_write else nullcontext() + with lock: + table_existed = False + if service_write: + db = lancedb.connect(uri=self.uri) + table_existed = self.table_name in db.list_tables().tables + + table = self.create_index(records=records, table_name=self.table_name) + if self.build_index: + self.write_to_index( + records, + table=table, + index_type=self.index_type, + metric=self.metric, + num_partitions=self.num_partitions, + num_sub_vectors=self.num_sub_vectors, + hybrid=self.hybrid, + sparse=self.sparse, + fts_language=self.fts_language, + ) + elif service_write: + if self.hybrid: + self._ensure_fts_index(table) + if self.hybrid and table_existed: + self._writes_since_optimize += 1 + self._optimize_service_table_if_due(table) + else: + logger.info( + "Skipping LanceDB index creation for table %r because build_index=False.", + self.table_name, + ) return records def put( @@ -1236,8 +1352,13 @@ def retrieval(self, vectors: Iterable[Sequence[float]], **kwargs: Any) -> list[l query_texts: Raw query strings aligned with ``vectors``. Required for ``hybrid=True`` and ignored for dense-only retrieval. + hybrid_fusion: + Optional :class:`HybridFusionPolicy`. When present, each hybrid leg + retrieves at least ``candidate_depth`` rows, applies weighted RRF, + and returns only the requested ``top_k`` rows. """ hybrid = kwargs.pop("hybrid", self.hybrid) + hybrid_fusion = kwargs.pop("hybrid_fusion", None) query_texts = kwargs.pop("query_texts", None) table_path = kwargs.pop("table_path", self.uri) table_name = kwargs.pop("table_name", self.table_name) @@ -1272,6 +1393,13 @@ def retrieval(self, vectors: Iterable[Sequence[float]], **kwargs: Any) -> list[l ) search_kwargs["query_type"] = "hybrid" search_kwargs.setdefault("fts_columns", "text") + elif hybrid_fusion is not None: + raise ValueError("hybrid_fusion requires hybrid=True") + + if hybrid_fusion is not None and not isinstance(hybrid_fusion, HybridFusionPolicy): + raise TypeError( + "hybrid_fusion must be a HybridFusionPolicy or None; " f"got {type(hybrid_fusion).__name__}" + ) where_clause = kwargs.pop("where", None) _filter_fallback = kwargs.pop("_filter", None) @@ -1306,10 +1434,15 @@ def retrieval(self, vectors: Iterable[Sequence[float]], **kwargs: Any) -> list[l query = table.search([vector], vector_column_name=vector_column_name, **search_kwargs) if where_clause is not None: query = query.where(where_clause) - query = query.limit(top_k).refine_factor(refine_factor).nprobes(n_probe) + query_limit = max(top_k, hybrid_fusion.candidate_depth) if hybrid_fusion is not None else top_k + query = query.limit(query_limit).refine_factor(refine_factor).nprobes(n_probe) + if hybrid_fusion is not None: + query = query.rerank(WeightedRRFReranker(hybrid_fusion)) if result_fields is not None: query = query.select(result_fields) results = query.to_list() + if hybrid_fusion is not None: + results = results[:top_k] search_results.append(results) return search_results diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_capabilities.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_capabilities.py index b675ff82b1..fbef9aa95f 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_capabilities.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_capabilities.py @@ -133,7 +133,8 @@ def inspect_lancedb_table_object(table: Any) -> LanceTableCapabilities: has_vector = vector_column is not None has_fts = bool(fts_columns) - retrieval_mode = _metadata_retrieval_mode(schema) or _mode_from_capabilities(has_vector, has_fts) + physical_mode = _mode_from_capabilities(has_vector, has_fts) + retrieval_mode = physical_mode if physical_mode != "unknown" else (_metadata_retrieval_mode(schema) or "unknown") return LanceTableCapabilities( has_vector=has_vector, diff --git a/nemo_retriever/src/nemo_retriever/graph/retriever.py b/nemo_retriever/src/nemo_retriever/graph/retriever.py index d342523a87..3e8a2d7407 100644 --- a/nemo_retriever/src/nemo_retriever/graph/retriever.py +++ b/nemo_retriever/src/nemo_retriever/graph/retriever.py @@ -16,6 +16,7 @@ filter_retrieval_kwargs, rerank_long_dataframe_to_hits, ) +from nemo_retriever.common.vdb.hybrid_fusion import DEFAULT_HYBRID_FUSION_POLICY from nemo_retriever.common.vdb.lancedb_capabilities import ( LanceRetrievalMode, LanceTableCapabilities, @@ -473,6 +474,7 @@ def queries( ] if mode == "hybrid": vdb_call_kwargs["hybrid"] = True + vdb_call_kwargs.setdefault("hybrid_fusion", DEFAULT_HYBRID_FUSION_POLICY) elif mode == "dense" and has_mode_override: vdb_call_kwargs["hybrid"] = False if caps.vector_column and caps.vector_column != "vector": diff --git a/nemo_retriever/src/nemo_retriever/ingest/index_mode.py b/nemo_retriever/src/nemo_retriever/ingest/index_mode.py new file mode 100644 index 0000000000..1b03b42330 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/ingest/index_mode.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Policy for resolving requested ingest modes against an existing LanceDB table.""" + +from __future__ import annotations + +from typing import Literal, cast + +from nemo_retriever.common.vdb.lancedb_capabilities import inspect_lancedb_table_object + +RequestedIngestIndexMode = Literal["auto", "dense", "hybrid", "sparse"] +ResolvedIngestIndexMode = Literal["dense", "hybrid", "sparse"] + +SUPPORTED_INGEST_INDEX_MODES: tuple[RequestedIngestIndexMode, ...] = ( + "auto", + "dense", + "hybrid", + "sparse", +) + + +def validate_requested_index_mode(index_mode: str) -> RequestedIngestIndexMode: + """Normalize and validate the public ingest index-mode vocabulary.""" + normalized = index_mode.strip().lower() + if normalized not in SUPPORTED_INGEST_INDEX_MODES: + raise ValueError(f"index_mode must be one of {', '.join(SUPPORTED_INGEST_INDEX_MODES)}, got {index_mode!r}.") + return cast(RequestedIngestIndexMode, normalized) + + +def resolve_ingest_index_mode( + requested_mode: RequestedIngestIndexMode, + *, + overwrite: bool, + existing_mode: ResolvedIngestIndexMode | None, +) -> ResolvedIngestIndexMode: + """Resolve one ingest request without mutating storage. + + ``auto`` makes fresh and overwritten tables hybrid, but preserves the + physical mode of a table during append. The sole compatible mode-changing + append is an explicit dense-to-hybrid upgrade. + """ + if overwrite or existing_mode is None: + return "hybrid" if requested_mode == "auto" else cast(ResolvedIngestIndexMode, requested_mode) + + if requested_mode == "auto" or requested_mode == existing_mode: + return existing_mode + + if existing_mode == "dense" and requested_mode == "hybrid": + return "hybrid" + + raise ValueError( + f"Cannot append with index_mode={requested_mode!r} to an existing {existing_mode!r} table. " + "Use index_mode='auto' to preserve the table mode, request 'hybrid' to upgrade a dense table, " + "or overwrite the table to replace it." + ) + + +def inspect_existing_lancedb_mode(uri: str, table_name: str) -> ResolvedIngestIndexMode | None: + """Return the physical mode of an existing table, or ``None`` when absent.""" + import lancedb # type: ignore + + db = lancedb.connect(uri) + if table_name not in db.list_tables().tables: + return None + + capabilities = inspect_lancedb_table_object(db.open_table(table_name)) + if capabilities.retrieval_mode == "unknown": + raise ValueError( + f"Cannot determine physical retrieval capabilities for LanceDB table {table_name!r} at {uri!r}." + ) + return cast(ResolvedIngestIndexMode, capabilities.retrieval_mode) diff --git a/nemo_retriever/src/nemo_retriever/ingest/plan.py b/nemo_retriever/src/nemo_retriever/ingest/plan.py index 93456d8a7c..1b03a1b823 100644 --- a/nemo_retriever/src/nemo_retriever/ingest/plan.py +++ b/nemo_retriever/src/nemo_retriever/ingest/plan.py @@ -39,13 +39,19 @@ expand_input_file_patterns, resolve_input_files, ) +from nemo_retriever.ingest.index_mode import ( + RequestedIngestIndexMode, + inspect_existing_lancedb_mode, + resolve_ingest_index_mode, + validate_requested_index_mode, +) from nemo_retriever.models import resolve_embed_model from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision IngestRunModeValue = Literal["inprocess", "batch"] IngestInputTypeValue = Literal["auto", "pdf", "doc", "txt", "html", "image", "audio", "video"] IngestProfileValue = Literal["auto", "fast-text"] -IngestIndexModeValue = Literal["dense", "hybrid", "sparse"] +IngestIndexModeValue = RequestedIngestIndexMode AudioSplitTypeValue = Literal["size", "time", "frame"] LocalIngestEmbedBackendValue = Literal["vllm", "hf"] OcrLangValue = OCRLang @@ -53,7 +59,6 @@ TableOutputFormatValue = Literal["pseudo_markdown", "markdown"] _SUPPORTED_RUN_MODES: tuple[IngestRunModeValue, ...] = ("inprocess", "batch") _SUPPORTED_PROFILES: tuple[IngestProfileValue, ...] = ("auto", "fast-text") -_SUPPORTED_INDEX_MODES: tuple[IngestIndexModeValue, ...] = ("dense", "hybrid", "sparse") _SUPPORTED_AUDIO_SPLIT_TYPES: tuple[AudioSplitTypeValue, ...] = ("size", "time", "frame") _SUPPORTED_INPUT_TYPES: tuple[IngestInputTypeValue, ...] = ( "auto", @@ -206,7 +211,7 @@ class IngestStorageOptions: lancedb_uri: str = "lancedb" table_name: str = "nemo-retriever" overwrite: bool = True - index_mode: IngestIndexModeValue = "dense" + index_mode: IngestIndexModeValue = "auto" @dataclass(frozen=True) @@ -242,10 +247,7 @@ def validate_ingest_profile(profile: str) -> IngestProfileValue: def validate_ingest_index_mode(index_mode: str) -> IngestIndexModeValue: - normalized = index_mode.strip().lower() - if normalized not in _SUPPORTED_INDEX_MODES: - raise ValueError(f"index_mode must be one of {', '.join(_SUPPORTED_INDEX_MODES)}, got {index_mode!r}.") - return cast(IngestIndexModeValue, normalized) + return validate_requested_index_mode(index_mode) def _validate_audio_split_type(split_type: str) -> AudioSplitTypeValue: @@ -597,7 +599,15 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: validated_run_mode = _validate_run_mode(runtime.run_mode) validated_profile = validate_ingest_profile(source.profile) validated_input_type = validate_ingest_input_type(source.input_type) - validated_index_mode = validate_ingest_index_mode(storage.index_mode) + requested_index_mode = validate_ingest_index_mode(storage.index_mode) + existing_index_mode = ( + None if storage.overwrite else inspect_existing_lancedb_mode(storage.lancedb_uri, storage.table_name) + ) + resolved_index_mode = resolve_ingest_index_mode( + requested_index_mode, + overwrite=storage.overwrite, + existing_mode=existing_index_mode, + ) validated_audio_split_type = _validate_audio_split_type(media.audio_split_type) document_list = expand_ingest_documents(source.documents, input_type=validated_input_type) branches = plan_extraction_branches(build_input_manifest(document_list)) @@ -636,7 +646,7 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: if extract_tuning is not None: extract_kwargs["batch_tuning"] = extract_tuning - embedding_model_name = None if validated_index_mode == "sparse" else resolve_embed_model(embed.embed_model_name) + embedding_model_name = None if resolved_index_mode == "sparse" else resolve_embed_model(embed.embed_model_name) embedding_model_revision = None if embedding_model_name is not None and not str(embed.embed_invoke_url or "").strip(): embedding_model_revision = resolve_embed_model_revision(embedding_model_name, None) @@ -662,16 +672,16 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: embed_gpus_per_actor=embed.batch.embed_gpus_per_actor, ) extract_params = ExtractParams(**extract_kwargs) - embed_params = None if validated_index_mode == "sparse" else EmbedParams(**embed_kwargs) if embed_kwargs else None + embed_params = None if resolved_index_mode == "sparse" else EmbedParams(**embed_kwargs) if embed_kwargs else None vdb_upload_kwargs = { "uri": storage.lancedb_uri, "table_name": storage.table_name, "overwrite": bool(storage.overwrite), } # Keep dense ingest kwargs unchanged unless the index mode needs additional LanceDB behavior. - if validated_index_mode == "sparse": + if resolved_index_mode == "sparse": vdb_upload_kwargs["sparse"] = True - elif validated_index_mode == "hybrid": + elif resolved_index_mode == "hybrid": vdb_upload_kwargs["hybrid"] = True if embedding_model_name is not None: vdb_upload_kwargs["embedding_model_name"] = embedding_model_name @@ -752,5 +762,5 @@ def resolve_ingest_plan(request: IngestPlanRequest) -> ResolvedIngestPlan: vdb_params=vdb_params, lancedb_uri=storage.lancedb_uri, table_name=storage.table_name, - sparse=validated_index_mode == "sparse", + sparse=resolved_index_mode == "sparse", ) diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index 3f6adaaa36..9be81aa8e3 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -20,7 +20,7 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path -from typing import Any, AsyncIterator, Union +from typing import Any, AsyncIterator, Literal, Union, cast import uvicorn from fastapi import FastAPI, Header, HTTPException, Query, Request, Response @@ -50,7 +50,13 @@ VDBResourceNotFound, ) from nemo_retriever.common.vdb.factory import get_vdb_op_cls +from nemo_retriever.common.vdb.hybrid_fusion import DEFAULT_HYBRID_FUSION_POLICY from nemo_retriever.common.vdb.records import RetrievalContractError +from nemo_retriever.ingest.index_mode import ( + inspect_existing_lancedb_mode, + resolve_ingest_index_mode, + validate_requested_index_mode, +) from nemo_retriever.operators.vdb import IngestVdbOperator, RetrieveVdbOperator from nemo_retriever.query.evidence import build_evidence_result from nemo_retriever.service.agentic_query import run_agentic_query @@ -65,6 +71,16 @@ logger = logging.getLogger(__name__) +ServiceIndexMode = Literal["auto", "dense", "hybrid"] + + +def _validate_service_index_mode(index_mode: str) -> ServiceIndexMode: + normalized = validate_requested_index_mode(index_mode) + if normalized == "sparse": + raise ValueError(f"index_mode must be one of auto, dense, hybrid; got {index_mode!r}.") + return cast(ServiceIndexMode, normalized) + + MAX_CONCURRENT_QUERIES = 4 MAX_CONCURRENT_AGENTIC_QUERIES = 2 @@ -215,8 +231,17 @@ def _production_vdb( lancedb_uri: str, table_name: str, expiration_cleanup_enabled: bool, + index_mode: ServiceIndexMode, ) -> VDB: """Construct the sole production VDB implementation for this service.""" + existing_mode = inspect_existing_lancedb_mode(lancedb_uri, table_name) + effective_mode = resolve_ingest_index_mode( + index_mode, + overwrite=False, + existing_mode=existing_mode, + ) + if effective_mode == "sparse": + raise ValueError("The VectorDB service requires a dense vector column; sparse-only tables are unsupported.") vdb_cls = get_vdb_op_cls("lancedb") return vdb_cls( uri=lancedb_uri, @@ -224,7 +249,9 @@ def _production_vdb( vector_dim=None, overwrite=False, build_index=False, + hybrid=effective_mode == "hybrid", _service_table_schema=True, + _service_index_mode=index_mode, expiration_cleanup_enabled=expiration_cleanup_enabled, ) @@ -261,6 +288,7 @@ def create_vectordb_app( hf_cache_dir: str | None = None, device: str | None = None, gpu_memory_utilization: float = 0.45, + index_mode: ServiceIndexMode = "auto", internal_api_token: str | None = None, reconciliation_interval_seconds: int = 60, expiration_cleanup_enabled: bool = True, @@ -270,6 +298,7 @@ def create_vectordb_app( """Build the VectorDB FastAPI application around an injected VDB contract.""" if reconciliation_interval_seconds < 0: raise ValueError("reconciliation_interval_seconds must be non-negative") + index_mode = _validate_service_index_mode(index_mode) agentic_config = agentic_config or AgenticConfig() state: VectorDBState | None = None @@ -283,6 +312,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: lancedb_uri=lancedb_uri, table_name=table_name, expiration_cleanup_enabled=expiration_cleanup_enabled, + index_mode=index_mode, ) state = VectorDBState( vdb=backend, @@ -699,12 +729,18 @@ async def query( raise RetrievalContractError("Collection retrieval did not return strategies") hits_per_query, strategies = result else: + hybrid = backend_health.get("effective_retrieval_mode") == "hybrid" + retrieval_kwargs: dict[str, Any] = { + "query_texts": queries, + "top_k": req.top_k, + "hybrid": hybrid, + } + if hybrid: + retrieval_kwargs["hybrid_fusion"] = DEFAULT_HYBRID_FUSION_POLICY hits_per_query = await asyncio.to_thread( current.retrieve_operator.run, vectors, - query_texts=queries, - top_k=req.top_k, - hybrid=backend_health.get("effective_retrieval_mode") == "hybrid", + **retrieval_kwargs, ) if not isinstance(hits_per_query, list): raise RetrievalContractError("Legacy retrieval returned an invalid shape") @@ -796,6 +832,12 @@ def main() -> None: parser = argparse.ArgumentParser(description="NeMo Retriever VectorDB service") parser.add_argument("--lancedb-uri", default="/data/vectordb", help="LanceDB directory") parser.add_argument("--table-name", default="nemo_retriever", help="Vector table name") + parser.add_argument( + "--index-mode", + default="auto", + choices=("auto", "dense", "hybrid"), + help="Fresh-table index mode; auto creates hybrid and preserves existing table capabilities.", + ) parser.add_argument("--embed-endpoint", default="", help="Remote NIM/OpenAI-compatible embed URL") parser.add_argument("--embed-model", default="nvidia/llama-nemotron-embed-vl-1b-v2") parser.add_argument( @@ -889,6 +931,7 @@ def main() -> None: hf_cache_dir=args.hf_cache_dir or None, device=args.device or None, gpu_memory_utilization=args.gpu_memory_utilization, + index_mode=args.index_mode, internal_api_token=args.internal_api_token or None, reconciliation_interval_seconds=args.reconciliation_interval_seconds, expiration_cleanup_enabled=not args.disable_expiration_cleanup, diff --git a/nemo_retriever/tests/helm/test_vectordb_index_mode.py b/nemo_retriever/tests/helm/test_vectordb_index_mode.py new file mode 100644 index 0000000000..2f19f1c0c2 --- /dev/null +++ b/nemo_retriever/tests/helm/test_vectordb_index_mode.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +_ROOT = Path(__file__).resolve().parents[3] +_CHART = _ROOT / "nemo_retriever" / "helm" + + +def _helm_template(*args: str) -> subprocess.CompletedProcess[str]: + if shutil.which("helm") is None: + pytest.skip("helm executable is not installed") + return subprocess.run( + [ + "helm", + "template", + "retriever", + str(_CHART), + "--set", + "serviceConfig.nimEndpoints.embedInvokeUrl=http://embed.example/v1/embeddings", + *args, + ], + capture_output=True, + check=False, + text=True, + ) + + +@pytest.mark.parametrize("mode", ["auto", "dense", "hybrid"]) +def test_vectordb_index_mode_renders_into_deployment(mode: str) -> None: + result = _helm_template("--set", f"serviceConfig.vectordb.indexMode={mode}") + assert result.returncode == 0, result.stderr + documents = [doc for doc in yaml.safe_load_all(result.stdout) if doc] + deployment = next( + doc for doc in documents if doc.get("kind") == "Deployment" and doc["metadata"]["name"].endswith("-vectordb") + ) + args = deployment["spec"]["template"]["spec"]["containers"][0]["args"] + assert args[args.index("--index-mode") + 1] == mode + + +def test_vectordb_index_mode_rejects_invalid_value() -> None: + result = _helm_template("--set", "serviceConfig.vectordb.indexMode=sparse") + assert result.returncode != 0 + assert "indexMode must be one of" in result.stderr diff --git a/nemo_retriever/tests/ingest/test_index_mode.py b/nemo_retriever/tests/ingest/test_index_mode.py new file mode 100644 index 0000000000..f32662646c --- /dev/null +++ b/nemo_retriever/tests/ingest/test_index_mode.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from nemo_retriever.ingest.index_mode import resolve_ingest_index_mode + + +@pytest.mark.parametrize( + ("requested", "existing", "expected"), + [ + ("auto", None, "hybrid"), + ("dense", None, "dense"), + ("hybrid", None, "hybrid"), + ("sparse", None, "sparse"), + ("auto", "dense", "dense"), + ("dense", "dense", "dense"), + ("hybrid", "dense", "hybrid"), + ("auto", "hybrid", "hybrid"), + ("hybrid", "hybrid", "hybrid"), + ("auto", "sparse", "sparse"), + ("sparse", "sparse", "sparse"), + ], +) +def test_resolve_ingest_index_mode_compatible_transitions(requested, existing, expected) -> None: + assert resolve_ingest_index_mode(requested, overwrite=False, existing_mode=existing) == expected + + +@pytest.mark.parametrize("requested", ["auto", "dense", "hybrid", "sparse"]) +def test_resolve_ingest_index_mode_overwrite_ignores_existing_mode(requested) -> None: + expected = "hybrid" if requested == "auto" else requested + assert resolve_ingest_index_mode(requested, overwrite=True, existing_mode="sparse") == expected + + +@pytest.mark.parametrize( + ("requested", "existing"), + [ + ("dense", "hybrid"), + ("sparse", "hybrid"), + ("dense", "sparse"), + ("hybrid", "sparse"), + ("sparse", "dense"), + ], +) +def test_resolve_ingest_index_mode_rejects_incompatible_append(requested, existing) -> None: + with pytest.raises(ValueError, match="Cannot append"): + resolve_ingest_index_mode(requested, overwrite=False, existing_mode=existing) diff --git a/nemo_retriever/tests/test_lancedb_capabilities.py b/nemo_retriever/tests/test_lancedb_capabilities.py index 778b582720..ade84ac213 100644 --- a/nemo_retriever/tests/test_lancedb_capabilities.py +++ b/nemo_retriever/tests/test_lancedb_capabilities.py @@ -13,8 +13,15 @@ pa = pytest.importorskip("pyarrow") import nemo_retriever.graph.retriever as retriever_module # noqa: E402 -from nemo_retriever.common.vdb.lancedb_capabilities import LanceTableCapabilities, inspect_lancedb_table # noqa: E402 +from nemo_retriever.common.vdb.hybrid_fusion import ( + DEFAULT_HYBRID_FUSION_POLICY, # noqa: E402 +) from nemo_retriever.common.vdb.lancedb import LanceDB # noqa: E402 +from nemo_retriever.common.vdb.lancedb_capabilities import ( # noqa: E402 + LanceTableCapabilities, + inspect_lancedb_table, + inspect_lancedb_table_object, +) from nemo_retriever.graph.retriever import Retriever # noqa: E402 from nemo_retriever.operators.vdb import RetrieveVdbOperator # noqa: E402 @@ -159,6 +166,36 @@ def test_detector_returns_hybrid_for_vector_plus_fts_table(tmp_path) -> None: assert caps.retrieval_mode == "hybrid" +def test_physical_fts_overrides_stale_dense_schema_metadata(tmp_path) -> None: + uri = str(tmp_path / "db") + _create_vector_table(uri, "upgraded", fts=True) + table = lancedb.connect(uri).open_table("upgraded") + stale_table = type( + "StaleSchemaTable", + (), + {"schema": table.schema.with_metadata({b"retrieval_mode": b"dense"}), "list_indices": table.list_indices}, + )() + + caps = inspect_lancedb_table_object(stale_table) + + assert caps.retrieval_mode == "hybrid" + + +def test_physical_vector_only_table_overrides_stale_hybrid_metadata(tmp_path) -> None: + uri = str(tmp_path / "db") + _create_vector_table(uri, "legacy") + table = lancedb.connect(uri).open_table("legacy") + stale_table = type( + "StaleSchemaTable", + (), + {"schema": table.schema.with_metadata({b"retrieval_mode": b"hybrid"}), "list_indices": table.list_indices}, + )() + + caps = inspect_lancedb_table_object(stale_table) + + assert caps.retrieval_mode == "dense" + + def test_detector_returns_sparse_for_fts_only_table(tmp_path) -> None: uri = str(tmp_path / "db") _create_sparse_table(uri, "sparse") @@ -191,7 +228,7 @@ def test_hybrid_table_query_automatically_enables_hybrid(monkeypatch, tmp_path) Retriever(vdb_kwargs={"uri": uri, "table_name": "hybrid"}).query("alpha", top_k=1) - assert calls == [{"hybrid": True}] + assert calls == [{"hybrid": True, "hybrid_fusion": DEFAULT_HYBRID_FUSION_POLICY}] def test_existing_dense_query_behavior_is_unchanged(monkeypatch, tmp_path) -> None: @@ -221,7 +258,7 @@ def test_explicit_hybrid_override_on_hybrid_table(monkeypatch, tmp_path) -> None Retriever(vdb_kwargs={"uri": uri, "table_name": "hybrid", "retrieval_mode": "hybrid"}).query("alpha", top_k=1) - assert calls == [{"hybrid": True}] + assert calls == [{"hybrid": True, "hybrid_fusion": DEFAULT_HYBRID_FUSION_POLICY}] def test_explicit_sparse_override_on_hybrid_table_uses_sparse_retrieval(monkeypatch, tmp_path) -> None: diff --git a/nemo_retriever/tests/test_lancedb_retrieval_where.py b/nemo_retriever/tests/test_lancedb_retrieval_where.py index c22774a300..40d79cba03 100644 --- a/nemo_retriever/tests/test_lancedb_retrieval_where.py +++ b/nemo_retriever/tests/test_lancedb_retrieval_where.py @@ -14,6 +14,7 @@ lancedb = pytest.importorskip("lancedb") +from nemo_retriever.common.vdb.hybrid_fusion import HybridFusionPolicy from nemo_retriever.common.vdb.lancedb import LanceDB from nemo_retriever.common.vdb.records import to_client_vdb_records @@ -47,6 +48,34 @@ def _tiny_table(uri: str, *, create_fts_index: bool = False) -> None: table.create_fts_index("text", replace=True) +def _hybrid_weighting_table(uri: str) -> None: + schema = pa.schema( + [ + pa.field("vector", pa.list_(pa.float32(), 2)), + pa.field("text", pa.string()), + pa.field("metadata", pa.string()), + pa.field("source", pa.string()), + ] + ) + rows = [ + { + "vector": vector, + "text": text, + "metadata": json.dumps({"doc_id": doc_id}), + "source": "{}", + } + for doc_id, vector, text in ( + ("dense", [1.0, 0.0], "semantic dense anchor"), + ("middle_1", [0.9, 0.1], "middle one"), + ("middle_2", [0.8, 0.2], "middle two"), + ("middle_3", [0.7, 0.3], "middle three"), + ("sparse", [0.0, 1.0], "rare unicorn identifier"), + ) + ] + table = lancedb.connect(uri).create_table("t", rows, schema=schema, mode="overwrite") + table.create_fts_index("text", replace=True) + + def test_retrieval_where_filters_rows() -> None: d = tempfile.mkdtemp() _tiny_table(d) @@ -124,6 +153,34 @@ def test_hybrid_retrieval_uses_query_texts() -> None: assert results[0][0]["text"] == "alpha" +def test_hybrid_retrieval_applies_explicit_weighted_rrf_policy() -> None: + d = tempfile.mkdtemp() + _hybrid_weighting_table(d) + op = LanceDB(uri=d, table_name="t", overwrite=False, vector_dim=2, validate_vector_length=False) + + equal_rrf = op.retrieval( + [[1.0, 0.0]], + top_k=5, + table_path=d, + table_name="t", + hybrid=True, + query_texts=["unicorn"], + ) + tuned = op.retrieval( + [[1.0, 0.0]], + top_k=1, + table_path=d, + table_name="t", + hybrid=True, + query_texts=["unicorn"], + hybrid_fusion=HybridFusionPolicy(candidate_depth=5, dense_weight=0.8, rrf_k=10), + ) + + assert json.loads(equal_rrf[0][0]["metadata"])["doc_id"] == "sparse" + assert json.loads(tuned[0][0]["metadata"])["doc_id"] == "dense" + assert len(tuned[0]) == 1 + + def test_hybrid_ingestion_builds_searchable_fts_index_from_record_text() -> None: """`LanceDB.run(..., hybrid=True)` builds the BM25/FTS side of hybrid search.""" d = tempfile.mkdtemp() diff --git a/nemo_retriever/tests/test_root_cli_workflow.py b/nemo_retriever/tests/test_root_cli_workflow.py index fe7bc7b50a..177c8ba8e2 100644 --- a/nemo_retriever/tests/test_root_cli_workflow.py +++ b/nemo_retriever/tests/test_root_cli_workflow.py @@ -161,6 +161,7 @@ def fake_create_ingestor(**kwargs: Any) -> Any: "uri": "lancedb", "table_name": "nemo-retriever", "overwrite": True, + "hybrid": True, "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -203,6 +204,7 @@ def test_root_ingest_without_mode_accepts_local_options_before_documents(monkeyp "uri": "/tmp/default-lancedb", "table_name": "nemo-retriever", "overwrite": False, + "hybrid": True, "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -400,6 +402,7 @@ def fake_create_ingestor(**kwargs: Any) -> Any: "uri": "/tmp/lancedb", "table_name": "docs", "overwrite": True, + "hybrid": True, "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -420,6 +423,7 @@ def test_root_ingest_append_forwards_overwrite_false(monkeypatch, tmp_path) -> N "uri": "lancedb", "table_name": "nemo-retriever", "overwrite": False, + "hybrid": True, "embedding_model_name": "nvidia/llama-nemotron-embed-vl-1b-v2", "embedding_model_revision": "4ef1bfa6da3a909de6bd00611950b7ed99203117", } @@ -1735,6 +1739,66 @@ def test_root_ingest_rejects_redundant_no_dedup_flag(tmp_path) -> None: assert "No such option" in result.output +def test_root_ingest_default_builds_vector_and_fts_table(monkeypatch, tmp_path) -> None: + lancedb = pytest.importorskip("lancedb") + from nemo_retriever.common.vdb.lancedb import LanceDB + + fake_ingestor = _make_fake_ingestor() + doc = tmp_path / "a.pdf" + doc.write_bytes(b"%PDF-1.4\n") + db_path = tmp_path / "db" + records = [ + [ + { + "document_type": "text", + "metadata": { + "embedding": [1.0, 0.0], + "content": "alpha hybrid manual", + "content_metadata": {"id": "alpha", "page_number": 1}, + "source_metadata": {"source_id": str(doc)}, + }, + }, + { + "document_type": "text", + "metadata": { + "embedding": [0.0, 1.0], + "content": "beta hybrid manual", + "content_metadata": {"id": "beta", "page_number": 2}, + "source_metadata": {"source_id": str(doc)}, + }, + }, + ] + ] + + def write_with_real_lancedb(params) -> Any: + LanceDB( + **params.vdb_kwargs, + vector_dim=2, + num_partitions=1, + num_sub_vectors=1, + ).run(records) + return fake_ingestor + + fake_ingestor.vdb_upload.side_effect = write_with_real_lancedb + monkeypatch.setattr(ingest_execution, "create_ingestor", lambda **_: fake_ingestor) + monkeypatch.setattr( + ingest_execution, + "_count_lancedb_rows", + lambda uri, table_name: lancedb.connect(uri).open_table(table_name).count_rows(), + ) + + result = RUNNER.invoke( + cli_main.app, + ["ingest", str(doc), "--lancedb-uri", str(db_path), "--table-name", "hybrid_docs"], + ) + + assert result.exit_code == 0, result.output + table = lancedb.connect(str(db_path)).open_table("hybrid_docs") + assert "vector" in table.schema.names + index_names = {index.name.lower() for index in table.list_indices()} + assert any("text" in name or "fts" in name for name in index_names) + + def test_root_ingest_index_mode_sparse_skips_embedding_and_writes_fts_table(monkeypatch, tmp_path) -> None: lancedb = pytest.importorskip("lancedb") fake_ingestor = _make_fake_ingestor() diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index 2d41f04a88..752ec56b90 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -31,6 +31,7 @@ VDBResourceConflict, VDBResourceNotFound, ) +from nemo_retriever.common.vdb.hybrid_fusion import DEFAULT_HYBRID_FUSION_POLICY from nemo_retriever.common.vdb.lancedb import LanceDB from nemo_retriever.common.vdb.records import RetrievalContractError from nemo_retriever.service.vectordb_app import ( @@ -53,6 +54,7 @@ def __init__(self, *, table_exists: bool = False) -> None: self.legacy_rows = 1 if table_exists else 0 self.last_write_context: CollectionWriteContext | None = None self.last_retrieval: dict[str, Any] | None = None + self.last_legacy_retrieval: dict[str, Any] | None = None self.health_calls = 0 def create_index(self, **kwargs): @@ -63,6 +65,7 @@ def write_to_index(self, records: list, **kwargs): self.table_exists = self.legacy_rows > 0 def retrieval(self, queries: list, **kwargs): + self.last_legacy_retrieval = dict(kwargs) return [ [ { @@ -405,32 +408,38 @@ def retrieve_collection(self, *args: Any, **kwargs: Any): assert backend.legacy_retrieval_calls == 0 -def test_production_vdb_preserves_legacy_service_write_without_index_rebuild( - tmp_path, - monkeypatch, -) -> None: +def test_production_vdb_defaults_to_hybrid_without_vector_index_build(tmp_path) -> None: backend = vectordb_module._production_vdb( lancedb_uri=str(tmp_path), table_name="legacy", expiration_cleanup_enabled=True, + index_mode="auto", ) assert isinstance(backend, LanceDB) assert backend.build_index is False + assert backend.hybrid is True - index_writes = [] - monkeypatch.setattr( - backend, - "create_index", - lambda records, table_name: object(), - ) - monkeypatch.setattr( - backend, - "write_to_index", - lambda *args, **kwargs: index_writes.append((args, kwargs)), + +def test_legacy_hybrid_query_uses_default_weighted_fusion() -> None: + backend = FakeVDB(table_exists=True) + hybrid_health = backend.health() + hybrid_health.update( + effective_retrieval_mode="hybrid", + retrieval_strategies=["hybrid"], ) - assert backend.run([]) == [] - assert index_writes == [] + with patch.object(backend, "health", return_value=hybrid_health), patch.object( + VectorDBState, + "embed_queries", + return_value=[[1.0, 0.0]], + ): + with TestClient(_app(backend)) as client: + response = client.post("/v1/query", json={"query": "legacy"}) + + assert response.status_code == 200 + assert backend.last_legacy_retrieval is not None + assert backend.last_legacy_retrieval["hybrid"] is True + assert backend.last_legacy_retrieval["hybrid_fusion"] == DEFAULT_HYBRID_FUSION_POLICY def test_legacy_write_and_query_keep_existing_vdb_path() -> None: diff --git a/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py b/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py index a70df982bc..ca17fdde59 100644 --- a/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py +++ b/nemo_retriever/tests/test_service_vectordb_hybrid_integration.py @@ -4,18 +4,20 @@ from __future__ import annotations -from datetime import timedelta from unittest.mock import patch import lancedb import pytest from fastapi.testclient import TestClient +import nemo_retriever.common.vdb.lancedb as lancedb_module from nemo_retriever.common.vdb.lancedb import LanceDB from nemo_retriever.common.vdb.lancedb_capabilities import inspect_lancedb_table_object -from nemo_retriever.service.vectordb_app import VectorDBState, create_vectordb_app - -_DIM = 4 +from nemo_retriever.service.vectordb_app import ( + VectorDBState, + _production_vdb, + create_vectordb_app, +) def _record( @@ -40,15 +42,15 @@ def _record( _RECORD = _record() -def _backend(tmp_path, *, hybrid: bool = False) -> LanceDB: - return LanceDB( - uri=str(tmp_path), +def _backend(tmp_path, *, index_mode: str = "auto") -> LanceDB: + backend = _production_vdb( + lancedb_uri=str(tmp_path), table_name="nemo_retriever", - vector_dim=_DIM, - overwrite=False, - build_index=False, - hybrid=hybrid, + expiration_cleanup_enabled=True, + index_mode=index_mode, ) + assert isinstance(backend, LanceDB) + return backend def _write(backend: LanceDB, *records: dict) -> int: @@ -62,122 +64,112 @@ def _capabilities(backend: LanceDB): return inspect_lancedb_table_object(table) -def _prebuild_fts_index(uri: str, table_name: str) -> None: - """Simulate an ingestion pipeline that wrote a table with a BM25/FTS index. - - The VectorDB service itself never builds FTS; the query path only detects - an index that was created at ingestion time. - """ - table = lancedb.connect(uri).open_table(table_name) - table.create_fts_index("text", replace=True) - for stub in table.list_indices(): - if "text" in stub.name.lower() or "fts" in stub.name.lower(): - table.wait_for_index([stub.name], timeout=timedelta(seconds=600)) +def _app(backend: LanceDB): + return create_vectordb_app( + vdb=backend, + embed_endpoint="http://embed.example/v1/embeddings", + embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", + reconciliation_interval_seconds=0, + ) @pytest.mark.integration -def test_write_rows_persists_rows_without_building_fts(tmp_path) -> None: +def test_fresh_auto_write_builds_fts_and_query_uses_hybrid(tmp_path) -> None: backend = _backend(tmp_path) - assert _write(backend, _RECORD) == 1 + app = _app(backend) - caps = _capabilities(backend) - assert caps is not None - assert caps.has_vector - # The service must not build an FTS index on write; the table stays dense. - assert not caps.has_fts - assert backend.health()["effective_retrieval_mode"] == "dense" + with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): + with TestClient(app) as client: + write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) + response = client.post( + "/v1/query", + json={"query": "revenue", "top_k": 5, "format": "evidence"}, + ) + health = client.get("/v1/health") + + assert write.status_code == 200, write.text + assert response.status_code == 200, response.text + assert _capabilities(backend).has_fts + assert response.json()["results"][0]["coverage"]["strategies_used"] == ["hybrid"] + assert health.json()["configured_index_mode"] == "auto" + assert health.json()["effective_index_mode"] == "hybrid" + assert health.json()["fts_present"] is True @pytest.mark.integration -def test_append_does_not_build_or_mutate_fts(tmp_path) -> None: - backend = _backend(tmp_path) - assert _write(backend, _RECORD) == 1 +def test_auto_preserves_existing_dense_table(tmp_path) -> None: + dense = _backend(tmp_path, index_mode="dense") + assert _write(dense, _RECORD) == 1 - appended = _record( - vector=[0.0, 1.0, 0.0, 0.0], - text="Zephyr quarterly guidance mentions unicorn synergy.", - ) - assert _write(backend, appended) == 2 + reopened = _backend(tmp_path) + caps = _capabilities(reopened) - table = lancedb.connect(str(tmp_path)).open_table("nemo_retriever") - assert table.count_rows() == 2 - # Still no FTS index — appends only persist rows. - caps = _capabilities(backend) + assert reopened.hybrid is False + assert caps.has_vector assert not caps.has_fts + assert reopened.health()["effective_retrieval_mode"] == "dense" @pytest.mark.integration -def test_auto_resolves_hybrid_when_fts_prebuilt(tmp_path) -> None: - # Ingestion built the table with both a vector column and an FTS index. - seed = _backend(tmp_path, hybrid=True) - _write(seed, _RECORD) - _prebuild_fts_index(str(tmp_path), "nemo_retriever") - - backend = _backend(tmp_path, hybrid=True) - caps = _capabilities(backend) - assert caps.has_vector - assert caps.has_fts - assert backend.health()["effective_retrieval_mode"] == "hybrid" +def test_explicit_hybrid_upgrades_existing_dense_table(tmp_path) -> None: + dense = _backend(tmp_path, index_mode="dense") + _write(dense, _RECORD) + upgraded = _backend(tmp_path, index_mode="hybrid") -@pytest.mark.integration -def test_auto_resolves_dense_when_no_fts(tmp_path) -> None: - seed = _backend(tmp_path) - _write(seed, _RECORD) - - backend = _backend(tmp_path) - caps = _capabilities(backend) - assert caps.has_vector - assert not caps.has_fts - assert backend.health()["effective_retrieval_mode"] == "dense" + assert upgraded.hybrid is True + assert _capabilities(upgraded).has_fts + assert upgraded.health()["effective_retrieval_mode"] == "hybrid" @pytest.mark.integration -def test_query_auto_selects_hybrid_when_fts_prebuilt(tmp_path) -> None: - backend = _backend(tmp_path, hybrid=True) - app = create_vectordb_app( - vdb=backend, - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - ) +def test_explicit_dense_rejects_existing_hybrid_table(tmp_path) -> None: + hybrid = _backend(tmp_path) + _write(hybrid, _RECORD) - with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): - with TestClient(app) as client: - write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) - assert write.status_code == 200, write.text + with pytest.raises(ValueError, match="Cannot append"): + _backend(tmp_path, index_mode="dense") - # Ingestion builds the FTS index; the query path detects it. - _prebuild_fts_index(str(tmp_path), "nemo_retriever") - resp = client.post( - "/v1/query", - json={"query": "revenue", "top_k": 5, "format": "evidence"}, - ) +@pytest.mark.integration +@pytest.mark.parametrize( + ("write_threshold", "row_threshold"), + [(1, 100_000), (20, 1)], + ids=("write-count", "unindexed-rows"), +) +def test_incremental_maintenance_thresholds_optimize_and_update_health( + tmp_path, + monkeypatch, + write_threshold: int, + row_threshold: int, +) -> None: + backend = _backend(tmp_path) + _write(backend, _RECORD) + table = lancedb.connect(backend.uri).open_table(backend.table_name) + monkeypatch.setattr(lancedb_module, "_SERVICE_OPTIMIZE_WRITE_THRESHOLD", write_threshold) + monkeypatch.setattr(lancedb_module, "_SERVICE_OPTIMIZE_ROW_THRESHOLD", row_threshold) + + with patch.object(type(table), "optimize", autospec=True) as optimize: + _write(backend, _record(text="second row")) - assert resp.status_code == 200, resp.text - coverage = resp.json()["results"][0]["coverage"] - assert coverage["strategies_used"] == ["hybrid"] + optimize.assert_called_once() + assert backend.health()["last_optimization"]["status"] == "ok" @pytest.mark.integration -def test_query_auto_selects_dense_when_no_fts(tmp_path) -> None: - backend = _backend(tmp_path) - app = create_vectordb_app( - vdb=backend, - embed_endpoint="http://embed.example/v1/embeddings", - embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", - ) +def test_explicit_dense_query_remains_dense(tmp_path) -> None: + backend = _backend(tmp_path, index_mode="dense") + app = _app(backend) with patch.object(VectorDBState, "embed_queries", return_value=[[1.0, 0.0, 0.0, 0.0]]): with TestClient(app) as client: write = client.post("/internal/vectordb/write", json={"records": [[_RECORD]]}) - assert write.status_code == 200, write.text - - resp = client.post( + response = client.post( "/v1/query", json={"query": "revenue", "top_k": 5, "format": "evidence"}, ) - assert resp.status_code == 200, resp.text - coverage = resp.json()["results"][0]["coverage"] - assert coverage["strategies_used"] == ["dense"] + assert write.status_code == 200, write.text + assert response.status_code == 200, response.text + assert not _capabilities(backend).has_fts + assert response.json()["results"][0]["coverage"]["strategies_used"] == ["dense"]