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    http_client,
10    wasm_compat::{WasmCompatSend, WasmCompatSync},
11};
12use serde::{Deserialize, Serialize};
13
14#[derive(Debug, thiserror::Error)]
15pub enum RerankError {
16    #[error("HttpError: {0}")]
17    HttpError(#[from] http_client::Error),
18
19    #[error("JsonError: {0}")]
20    JsonError(#[from] serde_json::Error),
21
22    #[error("UrlError: {0}")]
23    UrlError(#[from] url::ParseError),
24
25    #[error("ResponseError: {0}")]
26    ResponseError(String),
27
28    #[error("ProviderError: {0}")]
29    ProviderError(String),
30}
31
32/// Trait for reranking models that score documents by relevance to a query.
33pub trait RerankModel: WasmCompatSend + WasmCompatSync {
34    /// The maximum number of documents that can be reranked in a single request.
35    const MAX_DOCUMENTS: usize;
36
37    /// Provider client type used to construct this rerank model.
38    type Client;
39
40    /// Construct a model handle from a provider client and model identifier.
41    fn make(client: &Self::Client, model: impl Into<String>) -> Self;
42
43    /// Rerank a list of documents against a query.
44    fn rerank(
45        &self,
46        query: &str,
47        documents: Vec<String>,
48    ) -> impl std::future::Future<Output = Result<RerankResponse, RerankError>> + WasmCompatSend;
49}
50
51/// A single reranked document result.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct RerankResult {
54    /// Index of the document in the original input list.
55    pub index: usize,
56    /// The document text, if requested via `return_documents`.
57    pub document: Option<String>,
58    /// Relevance score between 0 and 1 (higher is more relevant).
59    pub relevance_score: f64,
60}
61
62/// Response from a reranking request.
63#[derive(Debug, Clone)]
64pub struct RerankResponse {
65    /// Reranked results sorted by relevance (highest first).
66    pub results: Vec<RerankResult>,
67    /// Model identifier used for this request.
68    pub model: String,
69    /// Token usage for this rerank request.
70    pub usage: Usage,
71}
rig_core/rerank.rs

rig_core/
rerank.rs
