aimux_core/
reranking_model.rs

1//! The `RerankingModel` trait — the provider-facing interface for reranking.
2//!
3//! Aligned with Vercel AI SDK `RerankingModelV4`
4//! (`reference/ai/packages/provider/src/reranking-model/v4/`).
5
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use ts_rs::TS;
9
10use crate::error::AiMuxError;
11use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning};
12use crate::{AbortSignal, retry, timeout};
13
14/// Documents to rerank: either a list of texts or a list of JSON objects.
15///
16/// Aligned with V4 `RerankingModelV4CallOptions.documents`.
17#[derive(Debug, Clone, Serialize, Deserialize, TS)]
18#[ts(export)]
19pub enum RerankingDocuments {
20    /// A list of plain-text documents.
21    Text { values: Vec<String> },
22    /// A list of JSON-object documents.
23    Object { values: Vec<serde_json::Value> },
24}
25
26/// Options passed to [`RerankingModel::do_rerank`].
27///
28/// Aligned with V4 `RerankingModelV4CallOptions`.
29#[derive(Debug, Clone, Serialize, Deserialize, TS)]
30#[ts(export)]
31pub struct RerankingCallOptions {
32    /// Documents to rerank.
33    pub documents: RerankingDocuments,
34
35    /// The query to rerank the documents against.
36    pub query: String,
37
38    /// Optional limit: return only the top `n` documents.
39    pub top_n: Option<u32>,
40
41    /// Abort signal for cancelling the operation.
42    #[serde(skip)]
43    #[ts(skip)]
44    pub abort_signal: Option<AbortSignal>,
45
46    /// Per-call retry override. `None` uses the model default.
47    pub max_retries: Option<u32>,
48
49    /// Per-call operation timeout.
50    pub timeout: Option<crate::options::TimeoutConfiguration>,
51
52    /// Additional provider-specific options, keyed by provider name.
53    pub provider_options: Option<SharedProviderOptions>,
54
55    /// Additional HTTP headers to send with the request.
56    pub headers: Option<SharedHeaders>,
57}
58
59impl RerankingCallOptions {
60    /// Create options with text documents and a query.
61    pub fn new(query: impl Into<String>, documents: RerankingDocuments) -> Self {
62        Self {
63            documents,
64            query: query.into(),
65            top_n: None,
66            abort_signal: None,
67            max_retries: None,
68            timeout: None,
69            provider_options: None,
70            headers: None,
71        }
72    }
73}
74
75/// A single reranked entry: the original index and its relevance score.
76#[derive(Debug, Clone, Serialize, Deserialize, TS)]
77#[ts(export)]
78pub struct RerankingRank {
79    /// The index of the document in the original list (before reranking).
80    pub index: u32,
81    /// The relevance score of the document after reranking.
82    pub relevance_score: f64,
83}
84
85/// The result of [`RerankingModel::do_rerank`].
86///
87/// Aligned with V4 `RerankingModelV4Result`.
88#[derive(Debug, Clone, Serialize, Deserialize, TS)]
89#[ts(export)]
90pub struct RerankingResult {
91    /// Ordered list of reranked documents, sorted by descending relevance
92    /// score. Each entry's `index` refers to the position in the original
93    /// `documents` list.
94    pub ranking: Vec<RerankingRank>,
95
96    /// Additional provider-specific metadata, keyed by provider name.
97    pub provider_metadata: Option<SharedProviderMetadata>,
98
99    /// Warnings for the call, e.g. unsupported settings.
100    pub warnings: Option<Vec<Warning>>,
101
102    /// Optional response information for debugging.
103    pub response: Option<RerankingResponse>,
104}
105
106/// Optional response information for a reranking call.
107#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
108#[ts(export)]
109pub struct RerankingResponse {
110    /// ID for the generated response, if the provider sends one.
111    pub id: Option<String>,
112    /// Timestamp for the start of the generated response (ISO 8601 string).
113    pub timestamp: Option<String>,
114    /// The ID of the model that was used to generate the response.
115    pub model_id: Option<String>,
116    /// Response headers.
117    pub headers: Option<SharedHeaders>,
118    /// Response body (opaque JSON).
119    pub body: Option<serde_json::Value>,
120}
121
122/// The unified reranking model trait (provider-facing).
123///
124/// Aligned with V4 `RerankingModelV4`.
125#[async_trait]
126pub trait RerankingModel: Send + Sync {
127    /// Specification version (always `"v4"`).
128    fn specification_version(&self) -> &'static str {
129        "v4"
130    }
131
132    /// Provider name, e.g. `"cohere"`.
133    fn provider(&self) -> &str;
134
135    /// Provider-specific model ID, e.g. `"rerank-english-v3.0"`.
136    fn model_id(&self) -> &str;
137
138    fn retry_config(&self) -> crate::retry::RetryConfig {
139        crate::retry::RetryConfig::default()
140    }
141
142    /// Rerank a list of documents using the query.
143    ///
144    /// Naming: the `do_` prefix prevents accidental direct usage by users.
145    async fn do_rerank(
146        &self,
147        options: &RerankingCallOptions,
148    ) -> Result<RerankingResult, AiMuxError>;
149}
150
151/// User-facing reranking with Core-owned retry and timeout.
152///
153/// # Errors
154///
155/// Returns the provider failure, retry exhaustion, timeout, or caller abort.
156pub async fn rerank(
157    model: &dyn RerankingModel,
158    options: RerankingCallOptions,
159) -> Result<RerankingResult, AiMuxError> {
160    let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?;
161    let abort_signal = options.abort_signal.clone();
162    let retries = retry::prepare_retries(
163        options.max_retries,
164        model.retry_config(),
165        abort_signal.clone(),
166    );
167    timeout::run(
168        retries.retry(|| model.do_rerank(&options)),
169        abort_signal.as_ref(),
170        timeout,
171    )
172    .await
173}
aimux_core/reranking_model.rs

aimux_core/
reranking_model.rs
