aimux_core/embedding_model.rs
1//! The `EmbeddingModel` trait — the provider-facing interface for text embeddings.
2//!
3//! Aligned with Vercel AI SDK `EmbeddingModelV4`
4//! (`reference/ai/packages/provider/src/embedding-model/v4/`).
5//!
6//! Providers implement [`EmbeddingModel::do_embed`]; users never call it
7//! directly.
8
9use async_trait::async_trait;
10use serde::{Deserialize, Serialize};
11use ts_rs::TS;
12
13use crate::error::AiMuxError;
14use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning};
15use crate::{AbortSignal, retry, timeout};
16
17/// A single embedding vector.
18///
19/// The TS spec types this as `Array<number>`; we use `f32` because that is the
20/// de-facto precision returned by every major embedding provider (OpenAI,
21/// Voyage, Cohere, …) and halves the memory footprint of `f64`.
22pub type Embedding = Vec<f32>;
23
24/// Options passed to [`EmbeddingModel::do_embed`].
25///
26/// Aligned with V4 `EmbeddingModelV4CallOptions`.
27#[derive(Debug, Clone, Serialize, Deserialize, TS)]
28#[ts(export)]
29pub struct EmbeddingCallOptions {
30 /// List of text values to generate embeddings for.
31 pub values: Vec<String>,
32
33 /// Abort signal for cancelling the operation.
34 #[serde(skip)]
35 #[ts(skip)]
36 pub abort_signal: Option<AbortSignal>,
37
38 /// Per-call retry override. `None` uses the model default.
39 pub max_retries: Option<u32>,
40
41 /// Per-call operation timeout.
42 pub timeout: Option<crate::options::TimeoutConfiguration>,
43
44 /// Additional provider-specific options, keyed by provider name.
45 pub provider_options: Option<SharedProviderOptions>,
46
47 /// Additional HTTP headers to send with the request.
48 pub headers: Option<SharedHeaders>,
49}
50
51impl EmbeddingCallOptions {
52 /// Create options for a single value with everything else unset.
53 pub fn new(value: impl Into<String>) -> Self {
54 Self {
55 values: vec![value.into()],
56 abort_signal: None,
57 max_retries: None,
58 timeout: None,
59 provider_options: None,
60 headers: None,
61 }
62 }
63}
64
65/// Token usage for an embedding call. Embeddings only report input tokens.
66#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
67#[ts(export)]
68pub struct EmbeddingUsage {
69 /// Number of input tokens consumed.
70 pub tokens: u32,
71}
72
73/// The result of [`EmbeddingModel::do_embed`].
74///
75/// Aligned with V4 `EmbeddingModelV4Result`.
76#[derive(Debug, Clone, Serialize, Deserialize, TS)]
77#[ts(export)]
78pub struct EmbeddingResult {
79 /// Generated embeddings, in the same order as the input `values`.
80 pub embeddings: Vec<Embedding>,
81
82 /// Token usage (input tokens only).
83 pub usage: Option<EmbeddingUsage>,
84
85 /// Additional provider-specific metadata.
86 pub provider_metadata: Option<SharedProviderMetadata>,
87
88 /// Optional response information for debugging.
89 pub response: Option<EmbeddingResponse>,
90
91 /// Warnings for the call, e.g. unsupported settings.
92 pub warnings: Vec<Warning>,
93}
94
95/// Debugging response info specific to embedding calls.
96#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
97#[ts(export)]
98pub struct EmbeddingResponse {
99 /// Response headers.
100 pub headers: Option<SharedHeaders>,
101 /// The response body (opaque JSON).
102 pub body: Option<serde_json::Value>,
103}
104
105/// The unified embedding model trait (provider-facing).
106///
107/// Aligned with V4 `EmbeddingModelV4`. Specific to text embeddings.
108///
109/// # Implementation notes
110///
111/// - `do_embed` is **not** user-facing API.
112/// - `embeddings` in the result must be in the same order as the input
113/// `values`.
114/// - Providers that cannot serve a request in a single call because
115/// `values.len() > max_embeddings_per_call()` should batch internally only
116/// when [`supports_parallel_calls`](Self::supports_parallel_calls) returns
117/// `true`; otherwise they should return an error.
118#[async_trait]
119pub trait EmbeddingModel: Send + Sync {
120 /// Specification version (always `"v4"`).
121 fn specification_version(&self) -> &'static str {
122 "v4"
123 }
124
125 /// Provider name, e.g. `"openai"`.
126 fn provider(&self) -> &str;
127
128 /// Provider-specific model ID, e.g. `"text-embedding-3-small"`.
129 fn model_id(&self) -> &str;
130
131 fn retry_config(&self) -> crate::retry::RetryConfig {
132 crate::retry::RetryConfig::default()
133 }
134
135 /// Limit of how many embeddings can be generated in a single API call.
136 ///
137 /// `None` means the model has no fixed limit. The TS spec allows this to
138 /// be a promise/function; in Rust we resolve it to a plain `Option<u32>`.
139 fn max_embeddings_per_call(&self) -> Option<u32>;
140
141 /// `true` if the model can handle multiple embedding calls in parallel.
142 fn supports_parallel_calls(&self) -> bool;
143
144 /// Generate a list of embeddings for the given input text.
145 ///
146 /// Naming: the `do_` prefix prevents accidental direct usage by users.
147 async fn do_embed(&self, options: &EmbeddingCallOptions)
148 -> Result<EmbeddingResult, AiMuxError>;
149}
150
151/// User-facing embedding operation with Core-owned retry and timeout.
152///
153/// # Errors
154///
155/// Returns the provider failure, retry exhaustion, timeout, or caller abort.
156pub async fn embed(
157 model: &dyn EmbeddingModel,
158 options: EmbeddingCallOptions,
159) -> Result<EmbeddingResult, 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_embed(&options)),
169 abort_signal.as_ref(),
170 timeout,
171 )
172 .await
173}