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

rig_core/
rerank.rs
