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