Skip to main content

rig_core/
rerank.rs

1//! Provider-agnostic reranking abstractions.
2//!
3//! Reranking models reorder a list of documents by relevance to a query.
4//! The [`RerankModel`] trait defines the interface, and [`RerankResponse`]
5//! carries both the scored results and token usage.
6
7use crate::{
8    completion::Usage,
9    wasm_compat::{WasmCompatSend, WasmCompatSync},
10};
11use serde::{Deserialize, Serialize};
12
13crate::provider_response::provider_error_enum!(
14    RerankError, "reranking" {
15        /// URL construction or parsing failed while preparing a provider request.
16        #[error("UrlError: {0}")]
17        UrlError(#[from] url::ParseError),
18    }
19);
20
21/// Trait for reranking models that score documents by relevance to a query.
22pub trait RerankModel: WasmCompatSend + WasmCompatSync {
23    /// The maximum number of documents that can be reranked in a single request.
24    const MAX_DOCUMENTS: usize;
25
26    /// Provider client type used to construct this rerank model.
27    type Client;
28
29    /// Construct a model handle from a provider client and model identifier.
30    fn make(client: &Self::Client, model: impl Into<String>) -> Self;
31
32    /// Rerank a list of documents against a query.
33    fn rerank(
34        &self,
35        query: &str,
36        documents: Vec<String>,
37    ) -> impl std::future::Future<Output = Result<RerankResponse, RerankError>> + WasmCompatSend;
38}
39
40/// A single reranked document result.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct RerankResult {
43    /// Index of the document in the original input list.
44    pub index: usize,
45    /// The document text, if requested via `return_documents`.
46    pub document: Option<String>,
47    /// Relevance score between 0 and 1 (higher is more relevant).
48    pub relevance_score: f64,
49}
50
51/// Response from a reranking request.
52#[derive(Debug, Clone)]
53pub struct RerankResponse {
54    /// Reranked results sorted by relevance (highest first).
55    pub results: Vec<RerankResult>,
56    /// Model identifier used for this request.
57    pub model: String,
58    /// Token usage for this rerank request.
59    pub usage: Usage,
60}