Skip to content

Add root service query mode#2228

Merged
jioffe502 merged 5 commits into
NVIDIA:mainfrom
jioffe502:codex/root-query-service-mode
Jun 24, 2026
Merged

Add root service query mode#2228
jioffe502 merged 5 commits into
NVIDIA:mainfrom
jioffe502:codex/root-query-service-mode

Conversation

@jioffe502

@jioffe502 jioffe502 commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add root service query as retriever query service, matching the mode-subcommand structure that landed in Move service ingest construction into ingest core #2221 for root ingest
  • keep retriever query ... as the default local LanceDB query path while moving query CLI parsing into nemo_retriever.cli.query
  • add a query-owned service workflow plus RetrieverServiceClient.query() for POST /v1/query
  • preserve the current root query output contracts: default hits JSON includes source/page/text/modality/score, --format evidence still emits {evidence, coverage}, and --max-text-chars still applies to hits output
  • apply candidate widening, page dedup, and content-type shaping client-side after service results

Review Notes

  • This is now rebased directly onto upstream/main; Move service ingest construction into ingest core #2221 is merged and no longer part of the diff.
  • The CLI direction follows Move service ingest construction into ingest core #2221: root main.py is thin, mode-specific parsing lives under cli/query/, and shared result shaping lives under query/shaping.py so local and service query use the same rules.
  • Service query intentionally exposes only service-appropriate options: service URL/token, top/candidate/page/content controls, and output-format controls. Local LanceDB/embed/rerank/hybrid options stay on the default local query command.
  • This PR does not move eval, recall, BEIR, harness, or reporting code, and does not expand pipeline service mode.

Validation

  • PYTHONPATH=/localhome/local-jioffe/NeMo-Retriever.worktrees/service-query-integration-base/nemo_retriever/src /localhome/local-jioffe/retriever-skills/nemo_retriever/.venv/bin/python -m pytest nemo_retriever/tests/test_query_workflow_options.py nemo_retriever/tests/test_root_query_cli.py nemo_retriever/tests/test_service_query_client.py nemo_retriever/tests/test_retriever_queries.py nemo_retriever/tests/test_root_cli_workflow.py — 116 passed
  • pre-commit run --files ... on all touched files — passed
  • git diff --check upstream/main...HEAD — passed
  • Live Helm service jp20 smoke against http://localhost:30670:
    • retriever ingest service /localhome/local-jioffe/datasets/nv-ingest/jp20 --service-url http://localhost:30670 --service-concurrency 4 — 20 files ingested, 3352 rows written
    • GET /v1/ingest/job/91e945956f7741dcbcbfacae15c07523 — completed, 20/20 documents, 0 failed, elapsed 65.1079s
    • retriever query service "In FY19, what were the earnings per share of the Saudi Research and Marketing Group?" --service-url http://localhost:30670 --top-k 3 --max-text-chars 260 — returned service hits JSON
    • retriever query service ... --candidate-k 5 --page-dedup --format evidence — returned {evidence, coverage} JSON with semantic hits from the ingested jp20 corpus

@jioffe502
jioffe502 force-pushed the codex/root-query-service-mode branch from daa852b to b0bf625 Compare June 22, 2026 16:28
@jioffe502
jioffe502 marked this pull request as ready for review June 22, 2026 16:37
@jioffe502
jioffe502 requested review from a team as code owners June 22, 2026 16:37
@jioffe502
jioffe502 requested a review from drobison00 June 22, 2026 16:37
@greptile-apps

greptile-apps Bot commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds retriever query service as a new subcommand that routes query traffic to a running Retriever service deployment via POST /v1/query, mirroring the retriever ingest service pattern from #2221. The query CLI parsing is moved from main.py into nemo_retriever/cli/query/, and result-shaping logic is extracted into a shared query/shaping.py so local and service query paths apply identical post-retrieval filtering rules.

  • New cli/query/ package implements DefaultLocalQueryGroup to preserve retriever query <text> as shorthand for the local LanceDB path while introducing retriever query service as an explicit mode.
  • New query/service.py orchestrates a service query by computing the candidate pool size, calling RetrieverServiceClient.query(), and applying shared shaping (page dedup, content-type filtering, top-k truncation) client-side.
  • service/query_schema.py extracts QueryRequest/QueryResult/QueryResponse into a shared schema used by both vectordb_app.py (server side) and the new RetrieverServiceClient.query() method; the JSON wire format is unchanged.

Confidence Score: 5/5

Safe to merge; the refactoring is clean, the wire format for POST /v1/query is unchanged, and all 116 tests pass.

The new service query path is well-isolated: candidate-pool validation happens before any network call, HTTP and schema errors are wrapped in RuntimeError that the CLI boundary catches, and the shared shaping logic is an exact extraction from the pre-existing graph/retriever.py with no behavioural changes. The only concerns are naming ambiguity and missing Raises documentation on newly-public functions — neither affects runtime correctness.

No files require special attention. The QueryRequest name collision in service/query_schema.py is worth addressing before this schema is used more widely, but it has no runtime impact today.

Important Files Changed

Filename Overview
nemo_retriever/src/nemo_retriever/cli/query/app.py New file implementing DefaultLocalQueryGroup (transparent fallback to _local), _local_command (preserves all existing local query options), and _service_command (new service mode); shares _emit_query_output and _query_cli_hit helpers. Prior review flagged _GROUP_OPTIONS and bearer token as P2; logic is sound and well-tested.
nemo_retriever/src/nemo_retriever/query/service.py New module orchestrating service-path queries: validates candidate_k >= top_k, constructs RetrieverServiceClient, calls POST /v1/query, and applies shared shaping client-side. Logic is clean; candidate-pool/top-k split correctly handled.
nemo_retriever/src/nemo_retriever/query/shaping.py Extracted from graph/retriever.py with no logic changes; now shared between local and service query paths. normalize_query_content_type_allowlist promoted to public; all functions carry over correct behaviour verified by existing tests.
nemo_retriever/src/nemo_retriever/service/client.py Adds synchronous query() method using httpx.Client; handles HTTP errors, JSON parse failures, and schema validation errors with specific RuntimeError messages. Timeout hardcoded at 300s (read) / 30s (connect).
nemo_retriever/src/nemo_retriever/service/query_schema.py New shared schema file; QueryRequest/QueryResult/QueryResponse extracted from vectordb_app. Wire format is unchanged. Class name QueryRequest duplicates the one in query/options.py — ambiguous for developers browsing the codebase.
nemo_retriever/src/nemo_retriever/service/vectordb_app.py Replaced inline QueryRequest/QueryResponse with imports from query_schema; result construction now builds typed QueryResult objects. JSON wire format is identical; tighter server-side validation is a net improvement.
nemo_retriever/src/nemo_retriever/cli/query/options.py New file with typed Annotated aliases for all CLI options; cleanly separates option definitions from command logic. ServiceApiTokenOption correctly includes envvar fallback.
nemo_retriever/src/nemo_retriever/query/options.py Adds QueryServiceOptions and ServiceQueryRequest dataclasses following the existing frozen-dataclass pattern. No validation on service_url format, consistent with existing QueryStorageOptions.
nemo_retriever/src/nemo_retriever/cli/main.py Correctly delegates to query_app; removes the large inline query_command and associated helpers, all now living in cli/query/app.py.
nemo_retriever/tests/test_service_query_client.py New test file covering RetrieverServiceClient.query(): auth headers, URL construction, empty hits, malformed responses (five parametrized cases), multi-query result-count mismatch. Good boundary coverage.
nemo_retriever/tests/test_root_query_cli.py Extended with six new tests covering service subcommand help output, option routing (rejects --lancedb-uri on service path), full request-object propagation, evidence format, and exit-code on unknown option.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User
    participant CLI as retriever query service
    participant App as cli/query/app.py _service_command
    participant SvcMod as query/service.py query_documents
    participant Client as RetrieverServiceClient.query()
    participant HTTP as POST /v1/query
    participant Shape as query/shaping.py shape_query_hits

    User->>CLI: retriever query service "question" --top-k 3 --candidate-k 5 --page-dedup
    CLI->>App: parse args → ServiceQueryRequest
    App->>App: _validate_output_options()
    App->>SvcMod: query_documents(ServiceQueryRequest)
    SvcMod->>SvcMod: "_service_retrieval_top_k() → candidate_k=5"
    SvcMod->>Client: "client.query("question", top_k=5)"
    Client->>HTTP: "POST /v1/query {query, top_k:5}"
    HTTP-->>Client: "200 {results: [{hits: [...]}]}"
    Client->>Client: QueryResponse.model_validate(body)
    Client-->>SvcMod: [[hit1, hit2, hit3, hit4, hit5]]
    SvcMod->>Shape: "shape_query_hits(raw_hits, top_k=3, page_dedup=True)"
    Shape-->>SvcMod: [hit1, hit2, hit3] (deduped)
    SvcMod-->>App: list[RetrievalHit]
    App->>App: "_emit_query_output(hits, strategies=["semantic"])"
    App-->>User: JSON output (hits or evidence)
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User
    participant CLI as retriever query service
    participant App as cli/query/app.py _service_command
    participant SvcMod as query/service.py query_documents
    participant Client as RetrieverServiceClient.query()
    participant HTTP as POST /v1/query
    participant Shape as query/shaping.py shape_query_hits

    User->>CLI: retriever query service "question" --top-k 3 --candidate-k 5 --page-dedup
    CLI->>App: parse args → ServiceQueryRequest
    App->>App: _validate_output_options()
    App->>SvcMod: query_documents(ServiceQueryRequest)
    SvcMod->>SvcMod: "_service_retrieval_top_k() → candidate_k=5"
    SvcMod->>Client: "client.query("question", top_k=5)"
    Client->>HTTP: "POST /v1/query {query, top_k:5}"
    HTTP-->>Client: "200 {results: [{hits: [...]}]}"
    Client->>Client: QueryResponse.model_validate(body)
    Client-->>SvcMod: [[hit1, hit2, hit3, hit4, hit5]]
    SvcMod->>Shape: "shape_query_hits(raw_hits, top_k=3, page_dedup=True)"
    Shape-->>SvcMod: [hit1, hit2, hit3] (deduped)
    SvcMod-->>App: list[RetrievalHit]
    App->>App: "_emit_query_output(hits, strategies=["semantic"])"
    App-->>User: JSON output (hits or evidence)
Loading

Reviews (7): Last reviewed commit: "Restore runtime ffmpeg deployment docs" | Re-trigger Greptile

@jioffe502
jioffe502 force-pushed the codex/root-query-service-mode branch 2 times, most recently from 926bc81 to 1dc4039 Compare June 22, 2026 18:33
# Conflicts:
#	nemo_retriever/src/nemo_retriever/cli/main.py
#	nemo_retriever/src/nemo_retriever/query/options.py
@jioffe502
jioffe502 merged commit 09071b6 into NVIDIA:main Jun 24, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants