openai_compat/types/
embeddings.rs

1//! Embeddings types, mirroring `openai-python/src/openai/types/embedding*.py`.
2
3use serde::{Deserialize, Serialize};
4
5use super::common::Usage;
6
7/// `input` accepts a string, an array of strings, or (arrays of) token ids.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9#[serde(untagged)]
10pub enum EmbeddingInput {
11    Text(String),
12    Texts(Vec<String>),
13    Tokens(Vec<i64>),
14    TokenArrays(Vec<Vec<i64>>),
15}
16
17impl From<&str> for EmbeddingInput {
18    fn from(text: &str) -> Self {
19        Self::Text(text.to_string())
20    }
21}
22
23impl From<String> for EmbeddingInput {
24    fn from(text: String) -> Self {
25        Self::Text(text)
26    }
27}
28
29impl From<Vec<String>> for EmbeddingInput {
30    fn from(texts: Vec<String>) -> Self {
31        Self::Texts(texts)
32    }
33}
34
35/// Request body for `POST /embeddings` (`resources/embeddings.py::create`).
36#[derive(Debug, Clone, Serialize)]
37pub struct EmbeddingRequest {
38    pub model: String,
39    pub input: EmbeddingInput,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub dimensions: Option<u32>,
42    /// `"float"` (default) or `"base64"`.
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub encoding_format: Option<String>,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub user: Option<String>,
47}
48
49impl EmbeddingRequest {
50    pub fn new(model: impl Into<String>, input: impl Into<EmbeddingInput>) -> Self {
51        Self {
52            model: model.into(),
53            input: input.into(),
54            dimensions: None,
55            encoding_format: None,
56            user: None,
57        }
58    }
59
60    pub fn dimensions(mut self, dimensions: u32) -> Self {
61        self.dimensions = Some(dimensions);
62        self
63    }
64
65    pub fn encoding_format(mut self, encoding_format: impl Into<String>) -> Self {
66        self.encoding_format = Some(encoding_format.into());
67        self
68    }
69}
70
71/// The embedding vector: floats by default, or a base64 string when
72/// `encoding_format: "base64"` was requested.
73#[derive(Debug, Clone, Serialize, Deserialize)]
74#[serde(untagged)]
75pub enum EmbeddingVector {
76    Floats(Vec<f32>),
77    Base64(String),
78}
79
80/// One embedding result.
81#[derive(Debug, Clone, Serialize, Deserialize)]
82#[non_exhaustive]
83pub struct Embedding {
84    pub index: u32,
85    pub embedding: EmbeddingVector,
86    #[serde(default)]
87    pub object: String,
88}
89
90/// Response from `POST /embeddings`.
91#[derive(Debug, Clone, Serialize, Deserialize)]
92#[non_exhaustive]
93pub struct EmbeddingResponse {
94    pub data: Vec<Embedding>,
95    pub model: String,
96    #[serde(default)]
97    pub object: String,
98    #[serde(default)]
99    pub usage: Option<Usage>,
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn input_variants_serialize() {
108        assert_eq!(
109            serde_json::to_value(EmbeddingInput::from("hi")).unwrap(),
110            serde_json::json!("hi")
111        );
112        assert_eq!(
113            serde_json::to_value(EmbeddingInput::Texts(vec!["a".into(), "b".into()])).unwrap(),
114            serde_json::json!(["a", "b"])
115        );
116        assert_eq!(
117            serde_json::to_value(EmbeddingInput::Tokens(vec![1, 2])).unwrap(),
118            serde_json::json!([1, 2])
119        );
120    }
121
122    #[test]
123    fn response_deserializes_floats() {
124        let body = r#"{
125            "object": "list",
126            "data": [{"object": "embedding", "index": 0, "embedding": [0.1, -0.2]}],
127            "model": "text-embedding-3-small",
128            "usage": {"prompt_tokens": 5, "total_tokens": 5}
129        }"#;
130        let response: EmbeddingResponse = serde_json::from_str(body).unwrap();
131        let EmbeddingVector::Floats(floats) = &response.data[0].embedding else {
132            panic!("expected float vector");
133        };
134        assert_eq!(floats.len(), 2);
135    }
136}
openai_compat/types/embeddings.rs

openai_compat/types/
embeddings.rs
