Add typed root query workflow options#2222
Conversation
Greptile SummaryThis PR extracts the root query execution logic from the CLI adapter layer into a new
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/query/workflow.py | New module housing Retriever construction and execution logic extracted from the CLI adapter; private helpers are correctly marked; local GPU rerank branch lacks a dedicated unit test in the companion test file. |
| nemo_retriever/src/nemo_retriever/query/options.py | Introduces five frozen dataclass option types as user-facing API; SPDX header present; pydantic/validation concern raised in a previous thread. |
| nemo_retriever/src/nemo_retriever/query/init.py | Minimal package init with SPDX header and docstring; option types are not re-exported here but are importable via the submodule path. |
| nemo_retriever/src/nemo_retriever/adapters/cli/query_workflow.py | Slimmed down to a thin passthrough facade delegating to nemo_retriever.query.workflow.query_documents; correct and consistent with adapter-layer design. |
| nemo_retriever/src/nemo_retriever/adapters/cli/main.py | Updated to construct typed QueryRequest and sub-option objects before calling query_documents; flag parsing and default wiring look correct. |
| nemo_retriever/tests/test_query_workflow_options.py | New unit tests covering the no-rerank and remote-rerank paths; local GPU rerank branch (empty reranker_invoke_url) is not exercised here, only through the CLI integration tests. |
| nemo_retriever/tests/test_root_query_cli.py | Updated to patch Retriever at nemo_retriever.query.workflow (query_core) instead of the old adapter module; patch site is now correct for the refactored call path. |
Sequence Diagram
sequenceDiagram
participant CLI as adapters/cli/main.py
participant CW as adapters/cli/query_workflow.py
participant QW as query/workflow.py
participant R as Retriever
CLI->>CLI: parse flags → build QueryRequest
CLI->>CW: query_documents(QueryRequest)
CW->>QW: run_query_documents(QueryRequest)
QW->>QW: _build_retriever_kwargs(request)
QW->>QW: build_embed_option_kwargs(...)
alt rerank.enabled
QW->>QW: _build_rerank_kwargs(request.rerank)
end
QW->>R: "Retriever(**retriever_kwargs)"
QW->>R: retriever.query(query, candidate_k, page_dedup, content_types)
R-->>QW: list[RetrievalHit]
QW-->>CW: list[RetrievalHit]
CW-->>CLI: list[RetrievalHit]
CLI->>CLI: print JSON output
Prompt To Fix All With AI
Fix the following 2 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 2
nemo_retriever/tests/test_query_workflow_options.py:89-124
**Local GPU rerank branch uncovered in unit tests**
`workflow._build_rerank_kwargs` has two branches: remote (non-empty `reranker_invoke_url`) and local GPU (empty URL, picks `_LOCAL_VL_RERANK_MODEL` and optional `reranker_backend`). The new `test_query_workflow_options.py` exercises only the remote branch. The local path is tested indirectly via `test_root_query_rerank_flag_enables_local_rerank` in the CLI test file, but a direct unit test here (calling `query_documents` with `QueryRerankOptions(enabled=True)` and no URL) would verify the fallback model name and the `local_reranker_backend` forwarding without requiring a full CLI stack.
### Issue 2 of 2
nemo_retriever/src/nemo_retriever/query/workflow.py:47-51
**Redundant guard on always-truthy `rerank_kwargs`**
`_build_rerank_kwargs` always returns a non-empty dict (either the remote path with at least `rerank_invoke_url`, or the local path with at least `model_name`), so `if rerank_kwargs:` is always `True`. The guard is harmless today, but if `_build_rerank_kwargs` is ever changed to return an empty dict on a new code path, the inconsistency between `rerank=True` (already set) and the missing `rerank_kwargs` key would cause a silent misconfiguration at runtime.
```suggestion
if request.rerank.enabled:
rerank_kwargs = _build_rerank_kwargs(request.rerank)
retriever_kwargs["rerank"] = True
retriever_kwargs["rerank_kwargs"] = rerank_kwargs
```
Reviews (2): Last reviewed commit: "Add typed query workflow options" | Re-trigger Greptile
94adfa5 to
ab82441
Compare
|
Greptile follow-up:
|
Summary
Closes #2218.
Validation
Note: pytest is not installed in the local venv available in this worktree.