ferrin_spec/
embedding_model.rs

1//! Embedding model interface.
2
3use std::future::Future;
4
5use serde::Deserialize;
6use serde::Serialize;
7use tokio_util::sync::CancellationToken;
8
9use crate::error::ProviderError;
10use crate::language_model::ResponseMetadata;
11use crate::shared::Headers;
12use crate::shared::ModelId;
13use crate::shared::ProviderId;
14use crate::shared::ProviderMetadata;
15use crate::shared::ProviderOptions;
16use crate::shared::Warning;
17
18/// An embedding vector.
19pub type Embedding = Vec<f64>;
20
21/// A model that turns text into embedding vectors.
22///
23/// Implementations report their batching limits through
24/// [`max_embeddings_per_call`](Self::max_embeddings_per_call) and
25/// [`supports_parallel_calls`](Self::supports_parallel_calls); the core splits
26/// inputs and schedules calls accordingly.
27pub trait EmbeddingModel: Send + Sync + 'static {
28    /// Provider identifier.
29    fn provider(&self) -> &ProviderId;
30
31    /// Model identifier.
32    fn model_id(&self) -> &ModelId;
33
34    /// Maximum number of values per call, or `None` when unlimited.
35    fn max_embeddings_per_call(&self) -> Option<usize>;
36
37    /// Maximum total UTF-8 size of the values of one call in bytes, or
38    /// `None` when unlimited.
39    fn max_input_bytes_per_call(&self) -> Option<usize> {
40        None
41    }
42
43    /// Whether several calls may run concurrently against this model.
44    fn supports_parallel_calls(&self) -> bool;
45
46    /// Embeds `options.values` in one provider call.
47    fn do_embed(
48        &self,
49        options: EmbedOptions,
50    ) -> impl Future<Output = Result<EmbedResult, ProviderError>> + Send;
51}
52
53/// Options for a single embedding call.
54#[derive(Debug, Clone, Default)]
55pub struct EmbedOptions {
56    /// Texts to embed.
57    pub values: Vec<String>,
58    /// Additional request headers.
59    pub headers: Headers,
60    /// Provider-specific options keyed by provider name.
61    pub provider_options: ProviderOptions,
62    /// Cancellation token.
63    pub cancellation: CancellationToken,
64}
65
66impl EmbedOptions {
67    /// Creates options for `values`.
68    #[must_use]
69    pub fn new(values: Vec<String>) -> Self {
70        Self {
71            values,
72            ..Self::default()
73        }
74    }
75}
76
77/// Token usage of an embedding call.
78#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
79pub struct EmbeddingUsage {
80    /// Input tokens consumed.
81    pub tokens: u64,
82}
83
84/// Result of an embedding call.
85#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
86pub struct EmbedResult {
87    /// One embedding per input value, in input order.
88    pub embeddings: Vec<Embedding>,
89    /// Token usage, if reported.
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub usage: Option<EmbeddingUsage>,
92    /// Provider-specific metadata.
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub provider_metadata: Option<ProviderMetadata>,
95    /// Response metadata (headers and body).
96    #[serde(default)]
97    pub response: ResponseMetadata,
98    /// Warnings.
99    #[serde(default)]
100    pub warnings: Vec<Warning>,
101}
ferrin_spec/embedding_model.rs

ferrin_spec/
embedding_model.rs
