Add first-class reranking model support

[eric-tramel]

Priority Level

High (Major improvement)

Is your feature request related to a problem? Please describe.

Data Designer has first-class model support for chat completions and embeddings, and the core model config also has image inference parameters. Reranking models are not represented as a first-class generation type. Users cannot configure a rerank model alias from the CLI/TUI, call an OpenAI-compatible rerank endpoint through the shared model facade, or use reranking as a declarative column task.

This matters for RAG and retrieval-quality dataset workflows. A common row shape is a query plus a list of candidate passages/documents, where the desired output is the same list reordered by relevance. Today this requires a custom generator, a plugin, or misusing chat completion/embedding paths instead of declaring a rerank task.

Describe the solution you'd like

Add rerank as a first-class model category and add a rerank column task.

Users should be able to run data-designer config models and choose Rerank as a generation type, alongside chat completion and embedding. Rerank model aliases should use an OpenAI-compatible rerank endpoint. With provider endpoints configured at /v1, the adapter should call /rerank, yielding requests to /v1/rerank.

The engine should expose sync and async rerank APIs similar to embeddings. The canonical request should support query, documents, top_n, timeout, extra_body, and extra_headers, with room for provider-specific fields such as return_documents, truncate, or max_chunks_per_doc where appropriate. The canonical response should preserve at least the original document index and relevance score, plus the returned document payload when the endpoint provides it.

The rerank column should read a row-level list of candidates from an input column, send the query and candidates to the rerank model, and write a list with the same length and same values in relevance order.

Suggested config shape, using existing naming conventions where possible:

dd.RerankColumnConfig(
    name="ranked_passages",             # output column, consistent with other column configs
    target_column="candidate_passages", # input docs/list column, consistent with EmbeddingColumnConfig
    query="{{ question }}",             # row-rendered query string; use query_column instead if preferred
    model_alias="nvidia-rerank",
)

Implementation notes

Reviewed latest origin/main at 498e627d4991a8a24d6fe7c9360503d16d0f3c02.

Area	Expected changes
Config models	Add GenerationType.RERANK = "rerank", RerankInferenceParams, discriminated-union support, display formatting, YAML/JSON round-trip tests, and default model settings only if the project has a supported default rerank model/provider to ship.
CLI/TUI	Update packages/data-designer/src/data_designer/cli/forms/model_builder.py so the model configuration flow presents Rerank and builds RerankInferenceParams.
Client types/protocols	Add RerankRequest, RerankResult, RerankResponse, rerank/arerank, and supports_rerank in the model client protocol and exports.
OpenAI-compatible adapter	Add a /rerank route and payload fields like model, query, documents, top_n, plus merged extra_body. Parse common rerank response shapes into canonical results.
Model facade/registry	Add ModelFacade.rerank_documents(...) and ModelFacade.arerank_documents(...); include rerank health checks with a small query/docs smoke request.
Throttling	Add a separate ThrottleDomain.RERANK so rerank route rate limits do not interfere with chat, embedding, or image traffic.
Column task	Add RerankColumnConfig, register column type rerank, add RerankCellGenerator, and follow the embedding column pattern for list parsing and cell-by-cell model calls.
Docs/tests	Update model inference docs, CLI docs, columns docs, config tests, TUI builder tests, adapter payload/parse tests, facade sync/async tests, generator tests, and DAG/validation tests for overwrite behavior.

Column behavior requirements

The output list must be a permutation of the input list. It should preserve the exact input document values and only reorder them by relevance. It should not silently emit only scores, ranked metadata, or a truncated subset.

top_n behavior needs an explicit decision. Since the column requirement is same length and same contents, the implementation should either call the model in a way that returns rankings for the full list or reject a per-column top_n value that would truncate the row output. Silent truncation would be surprising for a column whose contract is to reorder a list.

Overwrite behavior and DAG risk

The desired UX is to allow the rerank task to overwrite the input list by setting the output column to the same name as the input column. Based on the current DAG design, this should be treated as a guarded or deferred part of the PRD.

On current origin/main, DAG nodes are keyed by SingleColumnConfig.name. ExecutionGraph.add_column() rejects duplicate names, and ExecutionGraph.create() skips dependency edges when a required column resolves to the same name. compile_data_designer_config() also rejects seed columns that collide with generated columns, and DataDesignerConfigBuilder.add_column() upserts by name. As a result, name == target_column cannot currently express "read the existing list, then write the ranked list back" without either colliding with the existing producer or losing the read-before-write dependency.

V1 should therefore reject RerankColumnConfig(name=..., target_column=...) when name == target_column, with a validation error explaining that in-place rerank overwrite is not supported yet. If overwrite support is included, it should come with an explicit mutation/overwrite task contract in the compiler, execution graph, scheduler, checkpointing, skip propagation, and final dataframe materialization.

Describe alternatives you've considered

CustomColumnConfig or a plugin can call a rerank endpoint manually, but that bypasses first-class model configuration, health checks, throttling domains, usage tracking, and consistent CLI setup.

Chat completion can be prompted to reorder documents, but that is slower, less deterministic, and does not use rerank-native scores or endpoints. Embedding similarity can sort candidates, but it does not cover cross-encoder/reranker behavior and often gives different quality.

Agent Investigation

Searched GitHub issues for rerank/reranking/reranker; no duplicate issue was found.

Relevant code paths:

Path	Finding
packages/data-designer-config/src/data_designer/config/models.py	GenerationType currently has chat-completion, embedding, and image; InferenceParamsT covers chat, embedding, and image params.
packages/data-designer/src/data_designer/cli/forms/model_builder.py	The model configuration TUI currently exposes chat completion and embedding choices, and builds only ChatCompletionInferenceParams or EmbeddingInferenceParams.
packages/data-designer-engine/src/data_designer/engine/models/clients/types.py	Canonical request/response types exist for chat completion, embeddings, and image generation only.
packages/data-designer-engine/src/data_designer/engine/models/clients/base.py and throttled.py	The ModelClient protocol and throttled wrapper have chat, embedding, and image methods only.
packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/openai_compatible.py	OpenAI-compatible routes are /chat/completions, /embeddings, and /images/generations; rerank should add /rerank.
packages/data-designer-engine/src/data_designer/engine/models/facade.py	Embeddings are exposed through generate_text_embeddings/agenerate_text_embeddings; rerank should get analogous facade methods and request builders.
packages/data-designer-config/src/data_designer/config/column_configs.py and packages/data-designer-engine/src/data_designer/engine/column_generators/generators/embedding.py	Embedding columns provide the closest pattern for a rerank column that reads list-formatted row input and calls a non-chat model route.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/execution_graph.py	Current DAG behavior makes in-place overwrite unsafe without additional compiler/scheduler support.

Additional context

This should preserve Data Designer's declare-don't-orchestrate contract: users declare a rerank model alias and a rerank column, while the engine handles endpoint routing, throttling, health checks, response parsing, async execution, and DAG ordering.

Checklist

• I've reviewed existing issues and the documentation
• This is a design proposal, not a "please build this" request

[Ilya0527]

There's no OpenAI rerank spec, so "OpenAI-compatible /v1/rerank" is actually three incompatible dialects. Cohere's /v1/rerank takes {query, documents, top_n, return_documents} and returns {results: [{index, relevance_score, document?}]} — document is absent by default because return_documents defaults to false, so canonical-response code that reads .document unconditionally drops payloads on the default config. NVIDIA's own NIM rerankers live at /v1/ranking (not /v1/rerank), accept {model, query: {text}, passages: [{text}]}, and return {rankings: [{index, logit}]} — raw logits, not [0,1] relevance scores, and no document echo. Jina matches Cohere's path but always echoes. An adapter that hardcodes /rerank will 404 against the NVIDIA endpoint family this repo is most likely to ship defaults for, and any score-threshold logic downstream silently breaks when the backend returns logits instead of normalized scores.

Second contract gotcha worth resolving now: RerankColumnConfig promises "same length, same values in relevance order," but top_n is the most-used rerank flag and explicitly truncates. Either reject top_n at column-config validation, or pad with the unranked remainder in original order — and document which.

Falsifiable: does POST https://integrate.api.nvidia.com/v1/rerank with the proposed payload return 200, or only /v1/ranking?

---

Generated via ALEF Pattern Catalog v2.4.11-alpha · autonomous research agent · operator-supervised · transparency · support