openai_tools/embedding/response.rs
1//! OpenAI Embeddings API Response Types
2//!
3//! This module provides data structures for deserializing responses from the
4//! OpenAI Embeddings API. It includes all necessary types to handle complete
5//! API responses, including embedding vectors, metadata, and usage statistics.
6//!
7//! # Response Structure
8//!
9//! The main response structure contains:
10//! - `object`: Always "list" for embedding responses
11//! - `data`: Array of embedding objects with vectors and indices
12//! - `model`: The model used for generating embeddings
13//! - `usage`: Token usage information
14//!
15//! # Examples
16//!
17//! ## Parsing a Simple Embedding Response
18//!
19//! ```rust,no_run
20//! use openai_tools::embedding::response::Response;
21//!
22//! let json = r#"{
23//! "object": "list",
24//! "data": [{
25//! "object": "embedding",
26//! "embedding": [0.1, 0.2, 0.3, 0.4, 0.5],
27//! "index": 0
28//! }],
29//! "model": "text-embedding-3-small",
30//! "usage": {
31//! "prompt_tokens": 5,
32//! "total_tokens": 5
33//! }
34//! }"#;
35//!
36//! let response: Response = serde_json::from_str(json).unwrap();
37//! assert_eq!(response.object, "list");
38//! assert_eq!(response.data.len(), 1);
39//!
40//! let embedding = response.data[0].embedding.as_1d().unwrap();
41//! assert_eq!(embedding.len(), 5);
42//! ```
43//!
44//! ## Working with Batch Embeddings
45//!
46//! ```rust,no_run
47//! use openai_tools::embedding::response::Response;
48//!
49//! let json = r#"{
50//! "object": "list",
51//! "data": [
52//! {"object": "embedding", "embedding": [0.1, 0.2], "index": 0},
53//! {"object": "embedding", "embedding": [0.3, 0.4], "index": 1},
54//! {"object": "embedding", "embedding": [0.5, 0.6], "index": 2}
55//! ],
56//! "model": "text-embedding-3-small",
57//! "usage": {"prompt_tokens": 15, "total_tokens": 15}
58//! }"#;
59//!
60//! let response: Response = serde_json::from_str(json).unwrap();
61//!
62//! for data in &response.data {
63//! let vector = data.embedding.as_1d().unwrap();
64//! println!("Index {}: {:?}", data.index, vector);
65//! }
66//! ```
67//!
68//! # Embedding Dimensions
69//!
70//! The [`Embedding`] enum supports multiple dimensionalities:
71//! - **1D**: Standard embedding vector (most common)
72//! - **2D**: Matrix of embeddings
73//! - **3D**: Tensor of embeddings
74//!
75//! Use the `is_*d()` and `as_*d()` methods to check and access the appropriate dimension.
76
77use serde::Deserialize;
78
79/// Embedding vector that supports multiple dimensionalities.
80///
81/// The OpenAI API typically returns 1D vectors, but this enum provides
82/// flexibility for future API changes or custom use cases. The `#[serde(untagged)]`
83/// attribute allows automatic deserialization based on the JSON structure.
84///
85/// # Variants
86///
87/// * `OneDim` - Standard 1D embedding vector (e.g., `[0.1, 0.2, 0.3]`)
88/// * `TwoDim` - 2D embedding matrix (e.g., `[[0.1, 0.2], [0.3, 0.4]]`)
89/// * `ThreeDim` - 3D embedding tensor
90///
91/// # Examples
92///
93/// ```rust
94/// use openai_tools::embedding::response::Embedding;
95///
96/// // Parse a 1D embedding
97/// let json = r#"[0.1, 0.2, 0.3]"#;
98/// let embedding: Embedding = serde_json::from_str(json).unwrap();
99///
100/// assert!(embedding.is_1d());
101/// let vector = embedding.as_1d().unwrap();
102/// assert_eq!(vector.len(), 3);
103/// ```
104#[derive(Debug, Clone, Deserialize)]
105#[serde(untagged)]
106pub enum Embedding {
107 /// 1D embedding: Vec<f32>
108 OneDim(Vec<f32>),
109 /// 2D embedding: Vec<Vec<f32>>
110 TwoDim(Vec<Vec<f32>>),
111 /// 3D embedding: Vec<Vec<Vec<f32>>>
112 ThreeDim(Vec<Vec<Vec<f32>>>),
113}
114
115impl Embedding {
116 /// Returns the embedding as a 1D vector if it is 1D, otherwise returns None.
117 ///
118 /// # Returns
119 ///
120 /// * `Some(&Vec<f32>)` - Reference to the 1D vector if the embedding is 1D
121 /// * `None` - If the embedding is 2D or 3D
122 ///
123 /// # Example
124 ///
125 /// ```rust
126 /// use openai_tools::embedding::response::Embedding;
127 ///
128 /// let embedding: Embedding = serde_json::from_str("[0.1, 0.2, 0.3]").unwrap();
129 /// if let Some(vec) = embedding.as_1d() {
130 /// println!("Dimension: {}", vec.len());
131 /// }
132 /// ```
133 pub fn as_1d(&self) -> Option<&Vec<f32>> {
134 match self {
135 Embedding::OneDim(v) => Some(v),
136 _ => None,
137 }
138 }
139
140 /// Returns the embedding as a 2D vector if it is 2D, otherwise returns None.
141 ///
142 /// # Returns
143 ///
144 /// * `Some(&Vec<Vec<f32>>)` - Reference to the 2D vector if the embedding is 2D
145 /// * `None` - If the embedding is 1D or 3D
146 pub fn as_2d(&self) -> Option<&Vec<Vec<f32>>> {
147 match self {
148 Embedding::TwoDim(v) => Some(v),
149 _ => None,
150 }
151 }
152
153 /// Returns the embedding as a 3D vector if it is 3D, otherwise returns None.
154 ///
155 /// # Returns
156 ///
157 /// * `Some(&Vec<Vec<Vec<f32>>>)` - Reference to the 3D vector if the embedding is 3D
158 /// * `None` - If the embedding is 1D or 2D
159 pub fn as_3d(&self) -> Option<&Vec<Vec<Vec<f32>>>> {
160 match self {
161 Embedding::ThreeDim(v) => Some(v),
162 _ => None,
163 }
164 }
165
166 /// Returns true if the embedding is 1D.
167 ///
168 /// # Example
169 ///
170 /// ```rust
171 /// use openai_tools::embedding::response::Embedding;
172 ///
173 /// let embedding: Embedding = serde_json::from_str("[0.1, 0.2]").unwrap();
174 /// assert!(embedding.is_1d());
175 /// ```
176 pub fn is_1d(&self) -> bool {
177 matches!(self, Embedding::OneDim(_))
178 }
179
180 /// Returns true if the embedding is 2D.
181 pub fn is_2d(&self) -> bool {
182 matches!(self, Embedding::TwoDim(_))
183 }
184
185 /// Returns true if the embedding is 3D.
186 pub fn is_3d(&self) -> bool {
187 matches!(self, Embedding::ThreeDim(_))
188 }
189}
190
191/// Single embedding data item from the API response.
192///
193/// Each input text produces one `EmbeddingData` object containing
194/// the embedding vector and its index in the input array.
195///
196/// # Fields
197///
198/// * `object` - Type identifier, always "embedding"
199/// * `embedding` - The numerical vector representation of the input text
200/// * `index` - The position of this embedding in the input array (0-indexed)
201#[derive(Debug, Clone, Deserialize)]
202pub struct EmbeddingData {
203 /// Type identifier for this object (always "embedding")
204 pub object: String,
205 /// The embedding vector generated from the input text
206 pub embedding: Embedding,
207 /// Index of this embedding in the input array
208 pub index: usize,
209}
210
211/// Token usage statistics for the embedding request.
212///
213/// Provides information about the number of tokens processed,
214/// which is useful for cost estimation and rate limiting.
215///
216/// # Note
217///
218/// Unlike chat completions, embeddings only count prompt tokens
219/// (the input text) since there are no completion tokens generated.
220#[derive(Debug, Clone, Deserialize)]
221pub struct EmbeddingUsage {
222 /// Number of tokens in the input text(s)
223 pub prompt_tokens: usize,
224 /// Total tokens processed (equal to prompt_tokens for embeddings)
225 pub total_tokens: usize,
226}
227
228/// Complete response from the OpenAI Embeddings API.
229///
230/// This is the top-level structure returned by the API, containing
231/// all embedding data and metadata about the request.
232///
233/// # Example
234///
235/// ```rust,no_run
236/// use openai_tools::embedding::request::Embedding;
237///
238/// #[tokio::main]
239/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
240/// let mut embedding = Embedding::new()?;
241/// let response = embedding
242/// .model("text-embedding-3-small")
243/// .input_text("Hello, world!")
244/// .embed()
245/// .await?;
246///
247/// println!("Model: {}", response.model);
248/// println!("Embeddings: {}", response.data.len());
249/// println!("Tokens used: {}", response.usage.total_tokens);
250/// Ok(())
251/// }
252/// ```
253#[derive(Debug, Clone, Deserialize)]
254pub struct Response {
255 /// Type identifier for this response (always "list")
256 pub object: String,
257 /// Array of embedding data, one per input text
258 pub data: Vec<EmbeddingData>,
259 /// The model used to generate the embeddings
260 pub model: String,
261 /// Token usage statistics for this request
262 pub usage: EmbeddingUsage,
263}
264
265#[cfg(test)]
266mod tests {
267 use super::*;
268
269 #[test]
270 fn test_embedding_variants() {
271 // Test 1D embedding
272 let json_1d = r#"[0.1, 0.2, 0.3]"#;
273 let embedding_1d: Embedding = serde_json::from_str(json_1d).unwrap();
274 assert!(embedding_1d.is_1d());
275 assert!(!embedding_1d.is_2d());
276 assert!(!embedding_1d.is_3d());
277 assert_eq!(embedding_1d.as_1d(), Some(&vec![0.1, 0.2, 0.3]));
278 assert_eq!(embedding_1d.as_2d(), None);
279 assert_eq!(embedding_1d.as_3d(), None);
280
281 // Test 2D embedding
282 let json_2d = r#"[[0.1, 0.2], [0.3, 0.4]]"#;
283 let embedding_2d: Embedding = serde_json::from_str(json_2d).unwrap();
284 assert!(!embedding_2d.is_1d());
285 assert!(embedding_2d.is_2d());
286 assert!(!embedding_2d.is_3d());
287 assert_eq!(embedding_2d.as_1d(), None);
288 assert_eq!(embedding_2d.as_2d(), Some(&vec![vec![0.1, 0.2], vec![0.3, 0.4]]));
289 assert_eq!(embedding_2d.as_3d(), None);
290
291 // Test 3D embedding
292 let json_3d = r#"[[[0.1, 0.2], [0.3, 0.4]], [[0.5, 0.6], [0.7, 0.8]]]"#;
293 let embedding_3d: Embedding = serde_json::from_str(json_3d).unwrap();
294 assert!(!embedding_3d.is_1d());
295 assert!(!embedding_3d.is_2d());
296 assert!(embedding_3d.is_3d());
297 assert_eq!(embedding_3d.as_1d(), None);
298 assert_eq!(embedding_3d.as_2d(), None);
299 assert_eq!(embedding_3d.as_3d(), Some(&vec![vec![vec![0.1, 0.2], vec![0.3, 0.4]], vec![vec![0.5, 0.6], vec![0.7, 0.8]]]));
300 }
301}