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}