Add root service query mode#2228
Conversation
daa852b to
b0bf625
Compare
Greptile SummaryThis PR adds
|
| 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)
%%{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)
Reviews (7): Last reviewed commit: "Restore runtime ffmpeg deployment docs" | Re-trigger Greptile
926bc81 to
1dc4039
Compare
# Conflicts: # nemo_retriever/src/nemo_retriever/cli/main.py # nemo_retriever/src/nemo_retriever/query/options.py
Summary
retriever query service, matching the mode-subcommand structure that landed in Move service ingest construction into ingest core #2221 for root ingestretriever query ...as the default local LanceDB query path while moving query CLI parsing intonemo_retriever.cli.queryRetrieverServiceClient.query()forPOST /v1/queryhitsJSON includessource/page/text/modality/score,--format evidencestill emits{evidence, coverage}, and--max-text-charsstill applies to hits outputReview Notes
upstream/main; Move service ingest construction into ingest core #2221 is merged and no longer part of the diff.main.pyis thin, mode-specific parsing lives undercli/query/, and shared result shaping lives underquery/shaping.pyso local and service query use the same rules.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 passedpre-commit run --files ...on all touched files — passedgit diff --check upstream/main...HEAD— passedhttp://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 writtenGET /v1/ingest/job/91e945956f7741dcbcbfacae15c07523— completed, 20/20 documents, 0 failed, elapsed 65.1079sretriever 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 JSONretriever query service ... --candidate-k 5 --page-dedup --format evidence— returned{evidence, coverage}JSON with semantic hits from the ingested jp20 corpus