Skip to main content

providers_kit/
embedding.rs

1//! Portable embedding requests and results.
2//!
3//! Select one exact Model before submitting an [`EmbedRequest`]. Input order maps
4//! directly to [`EmbedResponse::vectors`](crate::embedding::EmbedResponse::vectors). Requested dimensions are optional and
5//! model support is authoritative at the operation boundary.
6//!
7//! # Embed an ordered collection
8//!
9//! This example needs an OpenAI key and access to the selected embedding model.
10//! Only compare or store vectors using the same selected model/dimension convention.
11//! The OpenAI operation validates vector count and index correlation. This example
12//! separately checks that each vector has the requested dimensions; neither check
13//! establishes semantic retrieval quality.
14//!
15//! ```no_run
16//! use providers_kit::{Provider, ProviderBackend, ProviderConfig, ModelBackend, ModelId};
17//! use providers_kit::config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig};
18//! use providers_kit::embedding::EmbedRequest;
19//! # #[tokio::main(flavor = "current_thread")]
20//! # async fn main() -> Result<(), providers_kit::KitError> {
21//! let provider = Provider::from_config(ProviderConfig::new(
22//!     ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
23//!     TransportConfig::default(),
24//! ))?;
25//! let work = async {
26//!     let model = provider.model(ModelId::new("text-embedding-3-small")?)?;
27//!     let request = EmbedRequest::new(vec!["first document".into(), "second document".into()])?
28//!         .with_dimensions(std::num::NonZeroU32::new(256).unwrap());
29//!     let response = model.embed(&request).await?;
30//!     assert_eq!(response.vectors.len(), request.inputs().len());
31//!     assert!(response.vectors.iter().all(|vector| vector.len() == 256));
32//!     Ok::<_, providers_kit::KitError>(response)
33//! }.await;
34//! let cleanup = provider.shutdown().await;
35//! if let Err(error) = &cleanup {
36//!     eprintln!("Cleanup failed: {error}");
37//! }
38//! let _ordered_vectors = work?;
39//! cleanup?;
40//! # Ok(())
41//! # }
42//! ```
43
44use std::num::NonZeroU32;
45
46use serde::{Deserialize, Deserializer, Serialize};
47
48use crate::KitError;
49use crate::error::ValidationViolation;
50use crate::generation::Usage;
51
52/// Validated text-embedding request.
53///
54/// # Examples
55///
56/// ```
57/// use providers_kit::embedding::EmbedRequest;
58/// let request = EmbedRequest::new(vec!["search document".into()])?
59///     .with_dimensions(std::num::NonZeroU32::new(128).unwrap());
60/// assert_eq!(request.inputs(), &["search document"]);
61/// assert_eq!(request.dimensions().unwrap().get(), 128);
62///
63/// # Ok::<(), providers_kit::KitError>(())
64/// ```
65///
66#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
67#[serde(deny_unknown_fields)]
68pub struct EmbedRequest {
69    inputs: Vec<String>,
70    dimensions: Option<NonZeroU32>,
71}
72
73impl EmbedRequest {
74    /// Creates a request containing non-empty text inputs.
75    ///
76    /// # Errors
77    ///
78    /// Returns an input error if the input list is empty or any text is whitespace-only.
79    ///
80    /// See [`crate::embedding::EmbedRequest`] for a complete example.
81    ///
82    pub fn new(inputs: Vec<String>) -> Result<Self, KitError> {
83        if inputs.is_empty() || inputs.iter().any(|value| value.trim().is_empty()) {
84            return Err(KitError::invalid_input(
85                ValidationViolation::Empty {
86                    field: "embedding.inputs",
87                },
88                "embedding inputs must contain non-empty text",
89            ));
90        }
91        Ok(Self {
92            inputs,
93            dimensions: None,
94        })
95    }
96
97    /// Requests a non-zero output dimensionality.
98    ///
99    /// The Model is already selected; dimensionality support is checked by the model operation.
100    ///
101    /// See [`crate::embedding::EmbedRequest`] for a complete example.
102    ///
103    #[must_use]
104    pub fn with_dimensions(mut self, dimensions: NonZeroU32) -> Self {
105        self.dimensions = Some(dimensions);
106        self
107    }
108
109    /// Returns text inputs in caller order.
110    ///
111    /// See [`crate::embedding::EmbedRequest`] for a complete example.
112    ///
113    #[must_use]
114    pub fn inputs(&self) -> &[String] {
115        &self.inputs
116    }
117
118    /// Returns the requested output dimensionality.
119    ///
120    /// See [`crate::embedding::EmbedRequest`] for a complete example.
121    ///
122    #[must_use]
123    pub const fn dimensions(&self) -> Option<NonZeroU32> {
124        self.dimensions
125    }
126}
127
128impl<'de> Deserialize<'de> for EmbedRequest {
129    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
130    where
131        D: Deserializer<'de>,
132    {
133        #[derive(Deserialize)]
134        #[serde(deny_unknown_fields)]
135        struct Wire {
136            inputs: Vec<String>,
137            dimensions: Option<NonZeroU32>,
138        }
139        let wire = Wire::deserialize(deserializer)?;
140        let mut request = Self::new(wire.inputs).map_err(serde::de::Error::custom)?;
141        request.dimensions = wire.dimensions;
142        Ok(request)
143    }
144}
145
146/// Embedding vectors returned in request-input order.
147///
148/// # Examples
149///
150/// ```
151/// use providers_kit::{embedding::EmbedResponse, generation::Usage};
152/// let response = EmbedResponse { vectors: vec![vec![0.25, 0.75]], usage: Usage::default() };
153/// assert_eq!(response.vectors.len(), 1);
154/// assert!(response.usage.total_tokens.is_none());
155/// ```
156///
157#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
158pub struct EmbedResponse {
159    /// Dense vectors corresponding exactly to request inputs.
160    pub vectors: Vec<Vec<f32>>,
161    /// Provider usage accounting.
162    pub usage: Usage,
163}