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

aimux_core/
embedding_model.rs
