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}