Skip to content

Add typed root query workflow options#2222

Open
jioffe502 wants to merge 1 commit into
NVIDIA:mainfrom
jioffe502:codex/query-core-2218
Open

Add typed root query workflow options#2222
jioffe502 wants to merge 1 commit into
NVIDIA:mainfrom
jioffe502:codex/query-core-2218

Conversation

@jioffe502

@jioffe502 jioffe502 commented Jun 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add query-owned typed request/options models for root query workflow
  • move root query Retriever construction into nemo_retriever.query while keeping Retriever.query() as the public facade
  • keep the CLI focused on flag/env parsing and preserve existing query JSON output behavior

Closes #2218.

Validation

  • py_compile on touched Python files
  • git diff --check
  • commit hooks passed on amend (trailing whitespace, EOF, AST, black, flake8)
  • live root query smoke against existing LanceDB index:
    • lancedb URI: /tmp/nemo_retriever_jp20_root_full
    • table: jp20_root_full
    • query returned the expected page 5 hit for the Issues Paper objective

Note: pytest is not installed in the local venv available in this worktree.

@jioffe502 jioffe502 requested review from a team as code owners June 9, 2026 16:59
@jioffe502 jioffe502 requested a review from nkmcalli June 9, 2026 16:59
@greptile-apps

greptile-apps Bot commented Jun 9, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR extracts the root query execution logic from the CLI adapter layer into a new nemo_retriever.query package, introducing typed frozen-dataclass option models (QueryRequest, QueryRetrievalOptions, QueryEmbedOptions, QueryRerankOptions, QueryStorageOptions) and keeping the CLI as a thin facade over the new query_documents entry point.

  • New nemo_retriever/query/ package: options.py defines the typed request/options dataclasses; workflow.py hosts _build_retriever_kwargs, _build_rerank_kwargs, and the public query_documents function previously living in adapters/cli/query_workflow.py.
  • adapters/cli/main.py: Updated to construct a QueryRequest from raw CLI flags before calling query_documents, keeping flag-parsing concerns out of the core query logic.
  • Tests: New test_query_workflow_options.py covers the no-rerank and remote-rerank paths at the unit level; test_root_query_cli.py patch sites updated to the new module location.

Confidence Score: 5/5

The refactor is a clean mechanical extraction with no behavior change; existing CLI tests and new unit tests pass through the full call chain correctly.

The diff is a well-scoped structural reorganization: the Retriever construction logic is moved verbatim from the CLI adapter into a dedicated query/ package, and the CLI is updated to construct typed option objects before delegating. No logic is altered, defaults are preserved, and the test patch sites are correctly updated to the new module. The two observations noted are a dead guard inherited from prior code and a gap in unit test coverage for the local GPU rerank branch — neither affects runtime correctness of the changed path.

No files require special attention; nemo_retriever/src/nemo_retriever/query/workflow.py has a minor inherited redundancy in the rerank guard, and nemo_retriever/tests/test_query_workflow_options.py is missing a unit test for the local GPU rerank branch.

Important Files Changed

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
Loading
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

Comment thread nemo_retriever/src/nemo_retriever/query/options.py
Comment thread nemo_retriever/src/nemo_retriever/query/workflow.py Outdated
Comment thread nemo_retriever/src/nemo_retriever/query/workflow.py Outdated
@jioffe502 jioffe502 force-pushed the codex/query-core-2218 branch from 94adfa5 to ab82441 Compare June 9, 2026 17:12
@jioffe502

Copy link
Copy Markdown
Collaborator Author

Greptile follow-up:

  • Addressed the public-helper/API-surface concern in ab82441: query package __init__ no longer re-exports the option models, resolver helpers are private (_build_*), and tests now exercise the public query_documents() boundary rather than importing helper internals.
  • Keeping the option/request types as frozen dataclasses is intentional for this PR. The merged Split root CLI ingest and query workflows #2209 ingest boundary uses the same pattern for IngestPlanRequest and the Ingest*Options types in nemo_retriever.ingest.plan; this Introduce typed query option models for root query workflow #2218 follow-up is aligning query with that design rather than introducing a different Pydantic contract.
  • Validation rerun after the cleanup: py_compile, git diff --check, commit hooks, and the live root-query smoke against the existing LanceDB index.

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.

Introduce typed query option models for root query workflow

1 participant