From e2d3e9457a2150bf3fd360eba827a6b5e3377e24 Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Mon, 22 Jun 2026 16:59:52 +0000 Subject: [PATCH 1/3] Add root query source and page filters --- .../src/nemo_retriever/cli/query/app.py | 36 +++++ .../src/nemo_retriever/cli/query/options.py | 29 ++++ .../src/nemo_retriever/query/filters.py | 131 ++++++++++++++++++ .../src/nemo_retriever/query/options.py | 10 ++ .../src/nemo_retriever/query/service.py | 2 +- .../src/nemo_retriever/query/workflow.py | 8 +- .../src/nemo_retriever/service/client.py | 12 +- .../nemo_retriever/service/query_schema.py | 4 + .../nemo_retriever/service/vectordb_app.py | 25 +++- nemo_retriever/tests/test_query_filters.py | 64 +++++++++ .../tests/test_query_workflow_options.py | 45 +++++- nemo_retriever/tests/test_root_query_cli.py | 73 ++++++++++ .../tests/test_service_query_client.py | 34 +++++ .../tests/test_service_vectordb_app.py | 54 ++++++++ 14 files changed, 518 insertions(+), 9 deletions(-) create mode 100644 nemo_retriever/src/nemo_retriever/query/filters.py create mode 100644 nemo_retriever/tests/test_query_filters.py diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 1c6d59d96..529bf9243 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -25,6 +25,7 @@ from nemo_retriever.query.options import ( QueryAgenticOptions, QueryEmbedOptions, + QueryFilterOptions, QueryRerankOptions, QueryRequest, QueryRetrievalOptions, @@ -155,6 +156,21 @@ def _retrieval_options( ) +def _filter_options( + *, + source_id: str | None, + source: str | None, + page_number: int | None, + where: str | None, +) -> QueryFilterOptions: + return QueryFilterOptions( + source_id=source_id, + source=source, + page_number=page_number, + where=where, + ) + + @app.command( "_local", hidden=True, @@ -170,6 +186,10 @@ def _local_command( candidate_k: opts.CandidateKOption = None, page_dedup: opts.PageDedupOption = False, content_types: opts.ContentTypesOption = None, + source_id: opts.SourceIdOption = None, + source: opts.SourceOption = None, + page_number: opts.PageNumberOption = None, + where: opts.WhereOption = None, lancedb_uri: opts.LanceDbUriOption = "lancedb", table_name: opts.TableNameOption = "nemo-retriever", embed_invoke_url: opts.EmbedInvokeUrlOption = None, @@ -258,6 +278,12 @@ def _request() -> QueryRequest: content_types=content_types, retrieval_mode=effective_retrieval_mode, ), + filters=_filter_options( + source_id=source_id, + source=source, + page_number=page_number, + where=where, + ), embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, @@ -295,6 +321,10 @@ def _service_command( candidate_k: opts.CandidateKOption = None, page_dedup: opts.PageDedupOption = False, content_types: opts.ContentTypesOption = None, + source_id: opts.SourceIdOption = None, + source: opts.SourceOption = None, + page_number: opts.PageNumberOption = None, + where: opts.WhereOption = None, output_format: opts.OutputFormatOption = "hits", max_text_chars: opts.MaxTextCharsOption = None, ) -> None: @@ -311,6 +341,12 @@ def _service_command( page_dedup=page_dedup, content_types=content_types, ), + filters=_filter_options( + source_id=source_id, + source=source, + page_number=page_number, + where=where, + ), service=QueryServiceOptions( service_url=service_url, service_api_token=service_api_token, diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 0c60b78d2..730f02843 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -47,6 +47,35 @@ ), ), ] +SourceIdOption = Annotated[ + str | None, + typer.Option( + "--source-id", + help="Only search chunks from this source identifier.", + ), +] +SourceOption = Annotated[ + str | None, + typer.Option( + "--source", + help="Only search chunks from this source name or identifier.", + ), +] +PageNumberOption = Annotated[ + int | None, + typer.Option( + "--page-number", + min=0, + help="Only search chunks from this page number.", + ), +] +WhereOption = Annotated[ + str | None, + typer.Option( + "--where", + help="Advanced: raw LanceDB SQL predicate appended to structured query filters.", + ), +] LanceDbUriOption = Annotated[ str, typer.Option( diff --git a/nemo_retriever/src/nemo_retriever/query/filters.py b/nemo_retriever/src/nemo_retriever/query/filters.py new file mode 100644 index 000000000..a15c38c7a --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/query/filters.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Translate root query filter options into LanceDB's native ``where`` predicate. + +LanceDB owns predicate execution. This module only maps stable root query +options such as source and page filters onto the table columns used by both +older local indexes and the VectorDB service, where ``metadata`` and ``source`` +are JSON strings at rest. +""" + +from __future__ import annotations + +import json +from typing import Any + +from nemo_retriever.query.options import QueryFilterOptions + + +def _clean_str(value: str | None) -> str | None: + cleaned = (value or "").strip() + return cleaned or None + + +def _sql_string(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _like_clause(column: str, pattern: str, *, escape: bool = False) -> str: + clause = f"{column} LIKE {_sql_string(pattern)}" + if escape: + clause += " ESCAPE '\\'" + return clause + + +def _like_pattern_literal(value: str) -> str: + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") + + +def _json_member_patterns(key: str, value: str) -> list[str]: + encoded_value = _like_pattern_literal(json.dumps(value, ensure_ascii=False)) + return [ + f'%"{key}":{encoded_value}%', + f'%"{key}": {encoded_value}%', + ] + + +def _source_json_member_clause(key: str, value: str) -> str: + return " OR ".join(_like_clause("source", pattern, escape=True) for pattern in _json_member_patterns(key, value)) + + +def _source_id_clause(source_id: str) -> str: + clauses = [ + _source_json_member_clause("source_id", source_id), + f"source = {_sql_string(source_id)}", + ] + return "(" + " OR ".join(clauses) + ")" + + +def _source_clause(source: str) -> str: + clauses = [ + _source_json_member_clause("source_id", source), + _source_json_member_clause("source_name", source), + f"source = {_sql_string(source)}", + ] + return "(" + " OR ".join(clauses) + ")" + + +def _page_number_clause(page_number: int) -> str: + page = int(page_number) + if page < 0: + raise ValueError("page_number must be greater than or equal to 0.") + patterns = [ + f'%"page_number":{page},%', + f'%"page_number":{page}}}%', + f'%"page_number": {page},%', + f'%"page_number": {page}}}%', + f'%"page_number":"{page}"%', + f'%"page_number": "{page}"%', + ] + return "(" + " OR ".join(_like_clause("metadata", pattern) for pattern in patterns) + ")" + + +def build_query_where_clause(options: QueryFilterOptions) -> str | None: + """Build the LanceDB predicate for root query filters.""" + clauses: list[str] = [] + + source_id = _clean_str(options.source_id) + if source_id is not None: + clauses.append(_source_id_clause(source_id)) + + source = _clean_str(options.source) + if source is not None: + clauses.append(_source_clause(source)) + + if options.page_number is not None: + clauses.append(_page_number_clause(options.page_number)) + + where = _clean_str(options.where) + if where is not None: + clauses.append(where if not clauses else f"({where})") + + if not clauses: + return None + return " AND ".join(clauses) + + +def query_filter_payload(options: QueryFilterOptions) -> dict[str, Any]: + """Return the public service query filter payload for non-empty fields.""" + payload: dict[str, Any] = {} + + source_id = _clean_str(options.source_id) + if source_id is not None: + payload["source_id"] = source_id + + source = _clean_str(options.source) + if source is not None: + payload["source"] = source + + if options.page_number is not None: + page_number = int(options.page_number) + if page_number < 0: + raise ValueError("page_number must be greater than or equal to 0.") + payload["page_number"] = page_number + + where = _clean_str(options.where) + if where is not None: + payload["where"] = where + + return payload diff --git a/nemo_retriever/src/nemo_retriever/query/options.py b/nemo_retriever/src/nemo_retriever/query/options.py index 237001a6e..1a2e15269 100644 --- a/nemo_retriever/src/nemo_retriever/query/options.py +++ b/nemo_retriever/src/nemo_retriever/query/options.py @@ -22,6 +22,14 @@ class QueryRetrievalOptions: retrieval_mode: QueryRetrievalMode = "auto" +@dataclass(frozen=True) +class QueryFilterOptions: + source_id: str | None = None + source: str | None = None + page_number: int | None = None + where: str | None = None + + @dataclass(frozen=True) class QueryEmbedOptions: embed_invoke_url: str | None = None @@ -67,6 +75,7 @@ class QueryAgenticOptions: class QueryRequest: query: str retrieval: QueryRetrievalOptions = field(default_factory=QueryRetrievalOptions) + filters: QueryFilterOptions = field(default_factory=QueryFilterOptions) embed: QueryEmbedOptions = field(default_factory=QueryEmbedOptions) rerank: QueryRerankOptions = field(default_factory=QueryRerankOptions) storage: QueryStorageOptions = field(default_factory=QueryStorageOptions) @@ -77,4 +86,5 @@ class QueryRequest: class ServiceQueryRequest: query: str retrieval: QueryRetrievalOptions = field(default_factory=QueryRetrievalOptions) + filters: QueryFilterOptions = field(default_factory=QueryFilterOptions) service: QueryServiceOptions = field(default_factory=QueryServiceOptions) diff --git a/nemo_retriever/src/nemo_retriever/query/service.py b/nemo_retriever/src/nemo_retriever/query/service.py index ef36f135e..db9287550 100644 --- a/nemo_retriever/src/nemo_retriever/query/service.py +++ b/nemo_retriever/src/nemo_retriever/query/service.py @@ -27,7 +27,7 @@ def query_documents(request: ServiceQueryRequest) -> list[RetrievalHit]: base_url=request.service.service_url, api_token=request.service.service_api_token, ) - raw_result_sets = client.query(request.query, top_k=retrieval_top_k) + raw_result_sets = client.query(request.query, top_k=retrieval_top_k, filters=request.filters) raw_hits = raw_result_sets[0] return shape_query_hits( raw_hits, diff --git a/nemo_retriever/src/nemo_retriever/query/workflow.py b/nemo_retriever/src/nemo_retriever/query/workflow.py index 5c777d537..5ceb6dbe5 100644 --- a/nemo_retriever/src/nemo_retriever/query/workflow.py +++ b/nemo_retriever/src/nemo_retriever/query/workflow.py @@ -13,6 +13,7 @@ from nemo_retriever.common.vdb.records import RetrievalHit from nemo_retriever.graph.retriever import Retriever from nemo_retriever.models import VL_RERANK_MODEL +from nemo_retriever.query.filters import build_query_where_clause from nemo_retriever.query.options import QueryRequest, QueryRerankOptions _LOCAL_VL_RERANK_MODEL = VL_RERANK_MODEL @@ -62,6 +63,7 @@ class ResolvedQueryPlan: candidate_k: int | None page_dedup: bool content_types: str | None + where: str | None lancedb_uri: str table_name: str retrieval_mode: str @@ -93,11 +95,14 @@ def create_retriever(self) -> Retriever: return Retriever(**self.retriever_kwargs()) def query_kwargs(self) -> dict[str, Any]: - return { + kwargs: dict[str, Any] = { "candidate_k": self.candidate_k, "page_dedup": self.page_dedup, "content_types": self.content_types, } + if self.where is not None: + kwargs["vdb_kwargs"] = {"where": self.where} + return kwargs def resolve_query_plan(request: QueryRequest) -> ResolvedQueryPlan: @@ -112,6 +117,7 @@ def resolve_query_plan(request: QueryRequest) -> ResolvedQueryPlan: candidate_k=request.retrieval.candidate_k, page_dedup=bool(request.retrieval.page_dedup), content_types=content_types, + where=build_query_where_clause(request.filters), lancedb_uri=str(request.storage.lancedb_uri), table_name=str(request.storage.table_name), retrieval_mode=str(request.retrieval.retrieval_mode), diff --git a/nemo_retriever/src/nemo_retriever/service/client.py b/nemo_retriever/src/nemo_retriever/service/client.py index 52bb76314..0509c4ab1 100644 --- a/nemo_retriever/src/nemo_retriever/service/client.py +++ b/nemo_retriever/src/nemo_retriever/service/client.py @@ -50,6 +50,8 @@ TimeRemainingColumn, ) +from nemo_retriever.query.filters import query_filter_payload +from nemo_retriever.query.options import QueryFilterOptions from nemo_retriever.service.query_schema import QueryResponse logger = logging.getLogger(__name__) @@ -216,11 +218,19 @@ def _auth_headers(self) -> dict[str, str]: # Query # ------------------------------------------------------------------ - def query(self, query: str | list[str], *, top_k: int) -> list[list[dict[str, Any]]]: + def query( + self, + query: str | list[str], + *, + top_k: int, + filters: QueryFilterOptions | None = None, + ) -> list[list[dict[str, Any]]]: """Search ingested documents through ``POST /v1/query``.""" url = f"{self._base_url}/v1/query" expected_results = len(query) if isinstance(query, list) else 1 payload: dict[str, Any] = {"query": query, "top_k": int(top_k)} + if filters is not None: + payload.update(query_filter_payload(filters)) try: with httpx.Client( timeout=httpx.Timeout(300.0, connect=30.0), diff --git a/nemo_retriever/src/nemo_retriever/service/query_schema.py b/nemo_retriever/src/nemo_retriever/service/query_schema.py index b38e31d47..f7dfbcc48 100644 --- a/nemo_retriever/src/nemo_retriever/service/query_schema.py +++ b/nemo_retriever/src/nemo_retriever/service/query_schema.py @@ -12,6 +12,10 @@ class QueryRequest(BaseModel): query: str | list[str] top_k: int = Field(default=10, ge=1, le=1000) + source_id: str | None = None + source: str | None = None + page_number: int | None = Field(default=None, ge=0) + where: str | None = None class QueryResult(BaseModel): diff --git a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py index 35560b963..fa5d4ced9 100644 --- a/nemo_retriever/src/nemo_retriever/service/vectordb_app.py +++ b/nemo_retriever/src/nemo_retriever/service/vectordb_app.py @@ -41,6 +41,8 @@ from fastapi import FastAPI, HTTPException from pydantic import BaseModel +from nemo_retriever.query.filters import build_query_where_clause +from nemo_retriever.query.options import QueryFilterOptions from nemo_retriever.service.query_schema import QueryRequest, QueryResponse, QueryResult logger = logging.getLogger(__name__) @@ -195,7 +197,12 @@ def total_rows(self) -> int: except Exception: return 0 - def search(self, vectors: list[list[float]], top_k: int) -> list[list[dict[str, Any]]]: + def search( + self, + vectors: list[list[float]], + top_k: int, + where: str | None = None, + ) -> list[list[dict[str, Any]]]: """Search the LanceDB table with precomputed query vectors.""" if not self._table_exists: return [[] for _ in vectors] @@ -203,9 +210,13 @@ def search(self, vectors: list[list[float]], top_k: int) -> list[list[dict[str, from nemo_retriever.common.vdb.records import normalize_retrieval_results table = self._db.open_table(self.table_name) + where_clause = str(where).strip() if where is not None else None raw_results = [] for vector in vectors: - results = table.search(vector).limit(top_k).to_list() + query = table.search(vector) + if where_clause: + query = query.where(where_clause) + results = query.limit(top_k).to_list() raw_results.append(results) return normalize_retrieval_results(raw_results) @@ -351,7 +362,15 @@ async def query(req: QueryRequest) -> QueryResponse: async with _query_semaphore: vectors = await asyncio.to_thread(_state.embed_queries, queries) - hits_per_query = await asyncio.to_thread(_state.search, vectors, req.top_k) + where_clause = build_query_where_clause( + QueryFilterOptions( + source_id=req.source_id, + source=req.source, + page_number=req.page_number, + where=req.where, + ) + ) + hits_per_query = await asyncio.to_thread(_state.search, vectors, req.top_k, where_clause) results = [QueryResult(hits=hits) for hits in hits_per_query] return QueryResponse(results=results) diff --git a/nemo_retriever/tests/test_query_filters.py b/nemo_retriever/tests/test_query_filters.py new file mode 100644 index 000000000..7db55a1a0 --- /dev/null +++ b/nemo_retriever/tests/test_query_filters.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from nemo_retriever.query.filters import build_query_where_clause, query_filter_payload +from nemo_retriever.query.options import QueryFilterOptions + + +def test_structured_source_id_and_page_filters_build_lancedb_predicate() -> None: + assert build_query_where_clause(QueryFilterOptions(source_id="docs/a.pdf", page_number=3)) == ( + "(source LIKE '%\"source_id\":\"docs/a.pdf\"%' ESCAPE '\\' " + "OR source LIKE '%\"source_id\": \"docs/a.pdf\"%' ESCAPE '\\' " + "OR source = 'docs/a.pdf') " + "AND (metadata LIKE '%\"page_number\":3,%' " + "OR metadata LIKE '%\"page_number\":3}%' " + "OR metadata LIKE '%\"page_number\": 3,%' " + "OR metadata LIKE '%\"page_number\": 3}%' " + 'OR metadata LIKE \'%"page_number":"3"%\' ' + 'OR metadata LIKE \'%"page_number": "3"%\')' + ) + + +def test_source_filter_matches_source_identifier_name_or_raw_source() -> None: + assert build_query_where_clause(QueryFilterOptions(source="a.pdf")) == ( + "(source LIKE '%\"source_id\":\"a.pdf\"%' ESCAPE '\\' " + "OR source LIKE '%\"source_id\": \"a.pdf\"%' ESCAPE '\\' " + "OR source LIKE '%\"source_name\":\"a.pdf\"%' ESCAPE '\\' " + "OR source LIKE '%\"source_name\": \"a.pdf\"%' ESCAPE '\\' " + "OR source = 'a.pdf')" + ) + + +def test_source_filters_escape_lancedb_like_wildcards() -> None: + where = build_query_where_clause(QueryFilterOptions(source_id="my_report%2026.pdf")) + + assert where is not None + assert "my\\_report\\%2026.pdf" in where + assert "ESCAPE '\\'" in where + assert "source = 'my_report%2026.pdf'" in where + + +def test_advanced_where_is_pass_through_and_combines_with_structured_filters() -> None: + assert build_query_where_clause(QueryFilterOptions(where="text = 'alpha'")) == "text = 'alpha'" + assert build_query_where_clause(QueryFilterOptions(page_number=2, where="text = 'alpha'")) == ( + "(metadata LIKE '%\"page_number\":2,%' " + "OR metadata LIKE '%\"page_number\":2}%' " + "OR metadata LIKE '%\"page_number\": 2,%' " + "OR metadata LIKE '%\"page_number\": 2}%' " + 'OR metadata LIKE \'%"page_number":"2"%\' ' + 'OR metadata LIKE \'%"page_number": "2"%\') ' + "AND (text = 'alpha')" + ) + + +def test_negative_page_filter_is_rejected() -> None: + with pytest.raises(ValueError, match="page_number must be greater than or equal to 0"): + build_query_where_clause(QueryFilterOptions(page_number=-1)) + + with pytest.raises(ValueError, match="page_number must be greater than or equal to 0"): + query_filter_payload(QueryFilterOptions(page_number=-1)) diff --git a/nemo_retriever/tests/test_query_workflow_options.py b/nemo_retriever/tests/test_query_workflow_options.py index f17d274bd..29552267b 100644 --- a/nemo_retriever/tests/test_query_workflow_options.py +++ b/nemo_retriever/tests/test_query_workflow_options.py @@ -10,8 +10,10 @@ import nemo_retriever.query.service as query_service import nemo_retriever.query.workflow as query_workflow +from nemo_retriever.query.filters import build_query_where_clause from nemo_retriever.query.options import ( QueryEmbedOptions, + QueryFilterOptions, QueryRerankOptions, QueryRequest, QueryRetrievalOptions, @@ -182,6 +184,36 @@ def query(self, query: str, **_kwargs: Any) -> list[dict[str, Any]]: assert result.strategies == strategies +def test_query_documents_threads_filter_options_into_vdb_kwargs(monkeypatch) -> None: + query_calls: list[tuple[str, dict[str, Any]]] = [] + + class FakeRetriever: + def __init__(self, **_kwargs: Any) -> None: + pass + + def query(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: + query_calls.append((query, kwargs)) + return [] + + monkeypatch.setattr(query_workflow, "Retriever", FakeRetriever) + + filters = QueryFilterOptions(source_id="docs/a.pdf", page_number=3) + request = QueryRequest(query="deployment?", filters=filters) + + assert query_workflow.query_documents(request) == [] + assert query_calls == [ + ( + "deployment?", + { + "candidate_k": None, + "page_dedup": False, + "content_types": None, + "vdb_kwargs": {"where": build_query_where_clause(filters)}, + }, + ) + ] + + def test_service_query_uses_candidate_pool_and_preserves_local_shaping(monkeypatch) -> None: client_calls: list[dict[str, Any]] = [] @@ -189,8 +221,14 @@ class FakeServiceClient: def __init__(self, *, base_url: str, api_token: str | None = None, **_kwargs: Any) -> None: client_calls.append({"base_url": base_url, "api_token": api_token}) - def query(self, query: str, *, top_k: int) -> list[list[dict[str, Any]]]: - client_calls.append({"query": query, "top_k": top_k}) + def query( + self, + query: str, + *, + top_k: int, + filters: QueryFilterOptions | None = None, + ) -> list[list[dict[str, Any]]]: + client_calls.append({"query": query, "top_k": top_k, "filters": filters}) return [ [ {"text": "keep", "source": "doc.pdf", "page_number": 1, "metadata": {"type": "text"}}, @@ -209,6 +247,7 @@ def query(self, query: str, *, top_k: int) -> list[list[dict[str, Any]]]: page_dedup=True, content_types="text", ), + filters=QueryFilterOptions(source="doc.pdf"), service=QueryServiceOptions(service_url="http://svc:7670", service_api_token="secret"), ) @@ -217,7 +256,7 @@ def query(self, query: str, *, top_k: int) -> list[list[dict[str, Any]]]: ] assert client_calls == [ {"base_url": "http://svc:7670", "api_token": "secret"}, - {"query": "deployment?", "top_k": 3}, + {"query": "deployment?", "top_k": 3, "filters": QueryFilterOptions(source="doc.pdf")}, ] diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 28bfeec9b..97992ca4d 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -12,6 +12,8 @@ import nemo_retriever.query.workflow as query_core from nemo_retriever.models import VL_EMBED_MODEL, VL_RERANK_MODEL +from nemo_retriever.query.filters import build_query_where_clause +from nemo_retriever.query.options import QueryFilterOptions RUNNER = CliRunner() cli_main = importlib.import_module("nemo_retriever.cli.main") @@ -113,6 +115,66 @@ def query(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: ] +def test_root_query_passes_source_page_filters_without_changing_output_shape(monkeypatch) -> None: + query_kwargs: list[dict[str, Any]] = [] + + class FakeRetriever: + def __init__(self, **_kwargs: Any) -> None: + pass + + def query(self, query: str, **kwargs: Any) -> list[dict[str, Any]]: + query_kwargs.append(kwargs) + return [ + { + "text": "text row", + "metadata": {"type": "text"}, + "page_number": 3, + "source": "docs/a.pdf", + "_distance": 0.2, + }, + ] + + monkeypatch.setattr(query_core, "Retriever", FakeRetriever) + + result = RUNNER.invoke( + cli_main.app, + [ + "query", + "deployment?", + "--source-id", + "docs/a.pdf", + "--page-number", + "3", + "--where", + "text != ''", + ], + ) + + assert result.exit_code == 0 + assert json.loads(result.output) == [ + { + "modality": "text", + "page_number": 3, + "score": 0.2, + "source": "docs/a.pdf", + "text": "text row", + }, + ] + assert set(json.loads(result.output)[0]) == {"modality", "page_number", "score", "source", "text"} + assert query_kwargs == [ + { + "candidate_k": None, + "page_dedup": False, + "content_types": None, + "vdb_kwargs": { + "where": build_query_where_clause( + QueryFilterOptions(source_id="docs/a.pdf", page_number=3, where="text != ''") + ) + }, + } + ] + + def test_root_query_passes_embed_options(monkeypatch) -> None: retriever_calls: list[dict[str, Any]] = [] query_calls: list[str] = [] @@ -587,6 +649,10 @@ def test_root_query_service_help_hides_local_only_options() -> None: assert "--top-k" in result.output assert "--candidate-k" in result.output assert "--content-types" in result.output + assert "--source-id" in result.output + assert "--source" in result.output + assert "--page-number" in result.output + assert "--where" in result.output assert "--format" in result.output assert "--max-text-chars" in result.output assert "--run-mode" not in result.output @@ -630,6 +696,12 @@ def fake_query_documents(request: Any) -> list[dict[str, Any]]: "--page-dedup", "--content-types", "text", + "--source-id", + "docs/a.pdf", + "--page-number", + "3", + "--where", + "text != ''", ], ) @@ -642,6 +714,7 @@ def fake_query_documents(request: Any) -> list[dict[str, Any]]: assert request.retrieval.candidate_k == 5 assert request.retrieval.page_dedup is True assert request.retrieval.content_types == "text" + assert request.filters == QueryFilterOptions(source_id="docs/a.pdf", page_number=3, where="text != ''") assert json.loads(result.output) == [ {"modality": "text", "page_number": 3, "score": 0.2, "source": "doc.pdf", "text": "service passage"}, ] diff --git a/nemo_retriever/tests/test_service_query_client.py b/nemo_retriever/tests/test_service_query_client.py index 755b2db4d..541d5471c 100644 --- a/nemo_retriever/tests/test_service_query_client.py +++ b/nemo_retriever/tests/test_service_query_client.py @@ -9,6 +9,7 @@ import pytest import nemo_retriever.service.client as service_client_module +from nemo_retriever.query.options import QueryFilterOptions from nemo_retriever.service.client import RetrieverServiceClient @@ -61,6 +62,39 @@ def test_service_client_query_posts_to_v1_query_with_auth(monkeypatch) -> None: } +def test_service_client_query_posts_filters_to_v1_query(monkeypatch) -> None: + calls: list[dict[str, Any]] = [] + _install_query_response( + monkeypatch, + {"results": [{"hits": [{"text": "passage", "source": "doc.pdf"}]}]}, + calls, + ) + + client = RetrieverServiceClient(base_url="http://svc:7670", api_token="secret") + + assert client.query( + "deployment?", + top_k=2, + filters=QueryFilterOptions( + source_id="docs/a.pdf", + source="a.pdf", + page_number=3, + where="text != ''", + ), + ) == [[{"text": "passage", "source": "doc.pdf"}]] + assert calls[1] == { + "url": "http://svc:7670/v1/query", + "json": { + "query": "deployment?", + "top_k": 2, + "source_id": "docs/a.pdf", + "source": "a.pdf", + "page_number": 3, + "where": "text != ''", + }, + } + + def test_service_client_query_accepts_empty_hits(monkeypatch) -> None: _install_query_response(monkeypatch, {"results": [{"hits": []}]}) diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index 910e42370..52b4640ce 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -82,3 +82,57 @@ def test_vector_db_state_local_embed_queries() -> None: assert vectors == [[1.0, 2.0]] mock_embedder.embed_queries.assert_called_once_with(["hello"]) + + +def test_query_rejects_negative_page_filter(tmp_path) -> None: + app = create_vectordb_app( + lancedb_uri=str(tmp_path), + table_name="test_table", + embed_endpoint="http://embed.example/v1/embeddings", + embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", + ) + with TestClient(app) as client: + resp = client.post("/v1/query", json={"query": "hello", "top_k": 3, "page_number": -1}) + + assert resp.status_code == 422 + + +def test_vectordb_search_applies_where_before_limit() -> None: + calls: list[tuple[str, object]] = [] + + class FakeQuery: + def where(self, where: str) -> "FakeQuery": + calls.append(("where", where)) + return self + + def limit(self, top_k: int) -> "FakeQuery": + calls.append(("limit", top_k)) + return self + + def to_list(self) -> list[dict[str, object]]: + return [{"text": "match", "metadata": "{}", "source": "{}"}] + + class FakeTable: + def search(self, vector: list[float]) -> FakeQuery: + calls.append(("search", vector)) + return FakeQuery() + + class FakeDB: + def open_table(self, table_name: str) -> FakeTable: + calls.append(("open_table", table_name)) + return FakeTable() + + state = object.__new__(VectorDBState) + state._table_exists = True + state._db = FakeDB() + state.table_name = "docs" + + result = state.search([[1.0, 0.0]], 2, "metadata LIKE '%\"page_number\":3%'") + + assert result[0][0]["text"] == "match" + assert calls == [ + ("open_table", "docs"), + ("search", [1.0, 0.0]), + ("where", "metadata LIKE '%\"page_number\":3%'"), + ("limit", 2), + ] From e4240640128420652610d3c5482667eba0a46968 Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Wed, 24 Jun 2026 18:50:15 +0000 Subject: [PATCH 2/3] Clarify service query filter contract --- .../src/nemo_retriever/cli/query/options.py | 2 +- .../nemo_retriever/service/query_schema.py | 11 ++- .../tests/test_lancedb_retrieval_where.py | 39 +++++++++++ .../tests/test_service_vectordb_app.py | 68 +++++++++++++++++++ 4 files changed, 118 insertions(+), 2 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/cli/query/options.py b/nemo_retriever/src/nemo_retriever/cli/query/options.py index 730f02843..71520f575 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/options.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/options.py @@ -73,7 +73,7 @@ str | None, typer.Option( "--where", - help="Advanced: raw LanceDB SQL predicate appended to structured query filters.", + help="Advanced/trusted-caller: raw LanceDB SQL predicate appended to structured query filters.", ), ] LanceDbUriOption = Annotated[ diff --git a/nemo_retriever/src/nemo_retriever/service/query_schema.py b/nemo_retriever/src/nemo_retriever/service/query_schema.py index f7dfbcc48..7a79211b3 100644 --- a/nemo_retriever/src/nemo_retriever/service/query_schema.py +++ b/nemo_retriever/src/nemo_retriever/service/query_schema.py @@ -15,7 +15,16 @@ class QueryRequest(BaseModel): source_id: str | None = None source: str | None = None page_number: int | None = Field(default=None, ge=0) - where: str | None = None + where: str | None = Field( + default=None, + description=( + "Advanced trusted-caller LanceDB/DataFusion predicate. The VectorDB service " + "passes this to table.search(...).where(...) as a read-only pre-filter before " + "limit(). Prefer source_id, source, and page_number for common filtering. " + "Deployers should expose this only inside the same auth/trust boundary as " + "unfiltered /v1/query access." + ), + ) class QueryResult(BaseModel): diff --git a/nemo_retriever/tests/test_lancedb_retrieval_where.py b/nemo_retriever/tests/test_lancedb_retrieval_where.py index 0d614aef2..a637d99fd 100644 --- a/nemo_retriever/tests/test_lancedb_retrieval_where.py +++ b/nemo_retriever/tests/test_lancedb_retrieval_where.py @@ -97,6 +97,45 @@ def test_retrieval_metadata_like_predicate() -> None: assert cm["doc_id"] == "x" +def test_retrieval_metadata_like_matches_compact_sidecar_metadata() -> None: + d = tempfile.mkdtemp() + 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": [1.0, 0.0], + "text": "alpha chunk", + "metadata": json.dumps({"meta_a": "alpha", "page_number": 1}, separators=(",", ":")), + "source": "{}", + }, + { + "vector": [0.0, 1.0], + "text": "bravo chunk", + "metadata": json.dumps({"meta_a": "bravo", "page_number": 2}, separators=(",", ":")), + "source": "{}", + }, + ] + lancedb.connect(d).create_table("t", rows, schema=schema, mode="overwrite") + op = LanceDB(uri=d, table_name="t", overwrite=False, vector_dim=2, validate_vector_length=False, hybrid=False) + + filtered = op.retrieval( + [[1.0, 0.0]], + top_k=10, + table_path=d, + table_name="t", + where='metadata LIKE \'%"meta_a":"alpha"%\'', + ) + + assert [hit["text"] for hit in filtered[0]] == ["alpha chunk"] + assert json.loads(filtered[0][0]["metadata"])["meta_a"] == "alpha" + + def test_retrieval_search_kwargs_must_be_dict() -> None: d = tempfile.mkdtemp() _tiny_table(d) diff --git a/nemo_retriever/tests/test_service_vectordb_app.py b/nemo_retriever/tests/test_service_vectordb_app.py index 52b4640ce..b8b92e226 100644 --- a/nemo_retriever/tests/test_service_vectordb_app.py +++ b/nemo_retriever/tests/test_service_vectordb_app.py @@ -8,6 +8,9 @@ from fastapi.testclient import TestClient +import nemo_retriever.service.vectordb_app as vectordb_app_module +from nemo_retriever.query.filters import build_query_where_clause +from nemo_retriever.query.options import QueryFilterOptions from nemo_retriever.service.vectordb_app import ( VectorDBState, _tensor_to_embedding_rows, @@ -136,3 +139,68 @@ def open_table(self, table_name: str) -> FakeTable: ("where", "metadata LIKE '%\"page_number\":3%'"), ("limit", 2), ] + + +def test_query_route_accepts_filters_and_applies_where(monkeypatch, tmp_path) -> None: + calls: list[tuple[str, object]] = [] + + class FakeState: + table_exists = True + embed_endpoint = "http://embed.example/v1/embeddings" + embed_mode = "remote" + + def total_rows(self) -> int: + return 1 + + def embed_queries(self, texts: list[str]) -> list[list[float]]: + calls.append(("embed", texts)) + return [[1.0, 0.0]] + + def search( + self, + vectors: list[list[float]], + top_k: int, + where: str | None = None, + ) -> list[list[dict[str, object]]]: + calls.append(("search", {"vectors": vectors, "top_k": top_k, "where": where})) + return [[{"text": "match", "metadata": '{"meta_a":"alpha"}', "source": "{}"}]] + + monkeypatch.setattr(vectordb_app_module, "VectorDBState", lambda **_kwargs: FakeState()) + + app = vectordb_app_module.create_vectordb_app( + lancedb_uri=str(tmp_path), + table_name="test_table", + embed_endpoint="http://embed.example/v1/embeddings", + embed_model="nvidia/llama-nemotron-embed-vl-1b-v2", + ) + with TestClient(app) as client: + resp = client.post( + "/v1/query", + json={ + "query": "hello", + "top_k": 3, + "source_id": "docs/a.pdf", + "page_number": 1, + "where": 'metadata LIKE \'%"meta_a":"alpha"%\'', + }, + ) + + assert resp.status_code == 200 + assert resp.json()["results"][0]["hits"][0]["text"] == "match" + assert calls == [ + ("embed", ["hello"]), + ( + "search", + { + "vectors": [[1.0, 0.0]], + "top_k": 3, + "where": build_query_where_clause( + QueryFilterOptions( + source_id="docs/a.pdf", + page_number=1, + where='metadata LIKE \'%"meta_a":"alpha"%\'', + ) + ), + }, + ), + ] From 9f4df77e660a46147634f4a0aaf2d882cf06e3a0 Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Wed, 24 Jun 2026 19:14:15 +0000 Subject: [PATCH 3/3] Reject unsupported agentic query filters --- .../src/nemo_retriever/cli/query/app.py | 16 ++++++++++------ .../src/nemo_retriever/query/filters.py | 5 +++++ nemo_retriever/tests/test_query_filters.py | 8 +++++++- nemo_retriever/tests/test_root_query_cli.py | 18 ++++++++++++++++++ 4 files changed, 40 insertions(+), 7 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/cli/query/app.py b/nemo_retriever/src/nemo_retriever/cli/query/app.py index 529bf9243..4f4eb3d87 100644 --- a/nemo_retriever/src/nemo_retriever/cli/query/app.py +++ b/nemo_retriever/src/nemo_retriever/cli/query/app.py @@ -22,6 +22,7 @@ silence_noisy_libraries, ) from nemo_retriever.common.vdb.records import RetrievalHit +from nemo_retriever.query.filters import has_query_filters from nemo_retriever.query.options import ( QueryAgenticOptions, QueryEmbedOptions, @@ -226,8 +227,16 @@ def _local_command( try: reranker_api_key = _api_key_from_env_option(reranker_api_key_env) if reranker_invoke_url else None effective_retrieval_mode = _query_retrieval_mode(ctx, retrieval_mode, hybrid) + filters = _filter_options( + source_id=source_id, + source=source, + page_number=page_number, + where=where, + ) if agentic: + if has_query_filters(filters): + raise ValueError("--source-id, --source, --page-number, and --where are not supported with --agentic.") request = QueryRequest( query=query, retrieval=_retrieval_options( @@ -278,12 +287,7 @@ def _request() -> QueryRequest: content_types=content_types, retrieval_mode=effective_retrieval_mode, ), - filters=_filter_options( - source_id=source_id, - source=source, - page_number=page_number, - where=where, - ), + filters=filters, embed=QueryEmbedOptions( embed_invoke_url=embed_invoke_url, embed_model_name=embed_model_name, diff --git a/nemo_retriever/src/nemo_retriever/query/filters.py b/nemo_retriever/src/nemo_retriever/query/filters.py index a15c38c7a..42c0654b3 100644 --- a/nemo_retriever/src/nemo_retriever/query/filters.py +++ b/nemo_retriever/src/nemo_retriever/query/filters.py @@ -106,6 +106,11 @@ def build_query_where_clause(options: QueryFilterOptions) -> str | None: return " AND ".join(clauses) +def has_query_filters(options: QueryFilterOptions) -> bool: + """Return whether any query filter option is set.""" + return build_query_where_clause(options) is not None + + def query_filter_payload(options: QueryFilterOptions) -> dict[str, Any]: """Return the public service query filter payload for non-empty fields.""" payload: dict[str, Any] = {} diff --git a/nemo_retriever/tests/test_query_filters.py b/nemo_retriever/tests/test_query_filters.py index 7db55a1a0..3096e7e73 100644 --- a/nemo_retriever/tests/test_query_filters.py +++ b/nemo_retriever/tests/test_query_filters.py @@ -6,7 +6,7 @@ import pytest -from nemo_retriever.query.filters import build_query_where_clause, query_filter_payload +from nemo_retriever.query.filters import build_query_where_clause, has_query_filters, query_filter_payload from nemo_retriever.query.options import QueryFilterOptions @@ -56,6 +56,12 @@ def test_advanced_where_is_pass_through_and_combines_with_structured_filters() - ) +def test_has_query_filters_reports_non_empty_options() -> None: + assert not has_query_filters(QueryFilterOptions()) + assert has_query_filters(QueryFilterOptions(source_id="docs/a.pdf")) + assert has_query_filters(QueryFilterOptions(where="text = 'alpha'")) + + def test_negative_page_filter_is_rejected() -> None: with pytest.raises(ValueError, match="page_number must be greater than or equal to 0"): build_query_where_clause(QueryFilterOptions(page_number=-1)) diff --git a/nemo_retriever/tests/test_root_query_cli.py b/nemo_retriever/tests/test_root_query_cli.py index 97992ca4d..0edd17f77 100644 --- a/nemo_retriever/tests/test_root_query_cli.py +++ b/nemo_retriever/tests/test_root_query_cli.py @@ -470,6 +470,24 @@ def test_root_query_agentic_requires_llm_model() -> None: assert "requires --agentic-llm-model" in result.output +def test_root_query_agentic_rejects_filter_flags() -> None: + result = RUNNER.invoke( + cli_main.app, + [ + "query", + "hello", + "--agentic", + "--agentic-llm-model", + "nvidia/llama-3.3-nemotron-super-49b-v1.5", + "--source-id", + "docs/a.pdf", + ], + ) + + assert result.exit_code == 1 + assert "not supported with --agentic" in result.output + + def test_root_query_agentic_plumbs_rerank_into_config(monkeypatch) -> None: """`--rerank` with `--agentic` wires the reranker config into AgenticRetrievalConfig (reranker model + endpoint + backend), so the agent's retrieval backend reranks."""