Add query filters for service-backed retrieval

nemo_retriever/src/nemo_retriever/cli/query/app.py

@@ -22,9 +22,11 @@
     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,
+    QueryFilterOptions,
     QueryRerankOptions,
     QueryRequest,
     QueryRetrievalOptions,
@@ -155,6 +157,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 +187,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,
@@ -206,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(
@@ -258,6 +287,7 @@ def _request() -> QueryRequest:
                     content_types=content_types,
                     retrieval_mode=effective_retrieval_mode,
                 ),
+                filters=filters,
                 embed=QueryEmbedOptions(
                     embed_invoke_url=embed_invoke_url,
                     embed_model_name=embed_model_name,
@@ -295,6 +325,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 +345,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,

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/trusted-caller: raw LanceDB SQL predicate appended to structured query filters.",
+    ),
+]
 LanceDbUriOption = Annotated[
     str,
     typer.Option(

nemo_retriever/src/nemo_retriever/query/filters.py

@@ -0,0 +1,136 @@
+# 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 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] = {}
+
+    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

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)

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,

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),

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),

nemo_retriever/src/nemo_retriever/service/query_schema.py

@@ -12,6 +12,19 @@
 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 = 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):