dynamo_protocols/types/
embeddings.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3//
4// Based on https://github.com/64bit/async-openai/ by Himanshu Neema
5// Original Copyright (c) 2022 Himanshu Neema
6// Licensed under MIT License (see ATTRIBUTIONS-Rust.md)
7//
8// Modifications Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES.
9// Licensed under Apache 2.0
10
11//! Dynamo-owned overrides of the upstream `async-openai` embedding response
12//! types.
13//!
14//! Upstream's `Embedding.embedding` is typed `Vec<f32>`, which can't
15//! represent the `encoding_format=base64` shape from the OpenAI spec
16//! (where each per-input embedding is a base64-encoded little-endian
17//! f32 byte string). We own the narrowest subtree -- [`Embedding`] and
18//! [`CreateEmbeddingResponse`] -- so the wire-level value can be either a
19//! JSON array of floats or a JSON string.
20//!
21//! Everything else (request, [`EmbeddingInput`], [`EmbeddingUsage`],
22//! [`EncodingFormat`], the builder types, the base64-specific upstream
23//! types like `Base64Embedding`) is re-exported via the glob below;
24//! the local [`Embedding`] and [`CreateEmbeddingResponse`] shadow their
25//! upstream namesakes. This preserves the previous full re-export
26//! surface for downstream crates that may have been importing those
27//! other symbols.
28
29use serde::{Deserialize, Serialize};
30
31#[allow(unused_imports)]
32pub use async_openai::types::embeddings::*;
33
34/// Per-input embedding payload.
35///
36/// The OpenAI `/v1/embeddings` spec returns either a vector of floats
37/// (`encoding_format="float"`, the default) or a base64 string encoding
38/// raw little-endian `f32` bytes (`encoding_format="base64"`).
39///
40/// Serializes as an untagged union: a JSON array for [`Float`] and a JSON
41/// string for [`Base64`], matching the OpenAI wire format.
42///
43/// [`Float`]: EmbeddingVector::Float
44/// [`Base64`]: EmbeddingVector::Base64
45#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
46#[serde(untagged)]
47pub enum EmbeddingVector {
48    Float(Vec<f32>),
49    Base64(String),
50}
51
52impl EmbeddingVector {
53    /// Number of float dimensions if this is the `Float` variant. Returns
54    /// `None` for `Base64` -- the count is encoded in the byte length and
55    /// only resolved on decode.
56    pub fn dim(&self) -> Option<usize> {
57        match self {
58            EmbeddingVector::Float(v) => Some(v.len()),
59            EmbeddingVector::Base64(_) => None,
60        }
61    }
62}
63
64impl From<Vec<f32>> for EmbeddingVector {
65    fn from(v: Vec<f32>) -> Self {
66        EmbeddingVector::Float(v)
67    }
68}
69
70/// A single embedding object inside an embeddings response.
71///
72/// Mirrors `async_openai::types::embeddings::Embedding` except for the
73/// `embedding` field, which is widened to [`EmbeddingVector`].
74#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
75pub struct Embedding {
76    pub index: u32,
77    pub object: String,
78    pub embedding: EmbeddingVector,
79}
80
81/// Top-level `/v1/embeddings` response.
82///
83/// Mirrors `async_openai::types::embeddings::CreateEmbeddingResponse`
84/// except the `data` field carries the Dynamo-owned [`Embedding`].
85#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
86pub struct CreateEmbeddingResponse {
87    pub object: String,
88    pub model: String,
89    pub data: Vec<Embedding>,
90    pub usage: EmbeddingUsage,
91}
92
93#[cfg(test)]
94mod tests {
95    use super::*;
96
97    #[test]
98    fn float_variant_serializes_as_array() {
99        let v = EmbeddingVector::Float(vec![0.1, 0.2, 0.3]);
100        let json = serde_json::to_string(&v).unwrap();
101        assert_eq!(json, "[0.1,0.2,0.3]");
102    }
103
104    #[test]
105    fn base64_variant_serializes_as_string() {
106        let v = EmbeddingVector::Base64("AAACOQ==".to_string());
107        let json = serde_json::to_string(&v).unwrap();
108        assert_eq!(json, "\"AAACOQ==\"");
109    }
110
111    #[test]
112    fn deserializes_array_as_float() {
113        let v: EmbeddingVector = serde_json::from_str("[0.1, 0.2, 0.3]").unwrap();
114        assert_eq!(v, EmbeddingVector::Float(vec![0.1, 0.2, 0.3]));
115    }
116
117    #[test]
118    fn deserializes_string_as_base64() {
119        let v: EmbeddingVector = serde_json::from_str("\"AAACOQ==\"").unwrap();
120        assert_eq!(v, EmbeddingVector::Base64("AAACOQ==".to_string()));
121    }
122
123    #[test]
124    fn round_trips_via_embedding_struct() {
125        let e = Embedding {
126            index: 0,
127            object: "embedding".to_string(),
128            embedding: EmbeddingVector::Float(vec![1.0, 2.0]),
129        };
130        let json = serde_json::to_string(&e).unwrap();
131        let parsed: Embedding = serde_json::from_str(&json).unwrap();
132        assert_eq!(e, parsed);
133
134        let e = Embedding {
135            index: 1,
136            object: "embedding".to_string(),
137            embedding: EmbeddingVector::Base64("AAACOQ==".to_string()),
138        };
139        let json = serde_json::to_string(&e).unwrap();
140        let parsed: Embedding = serde_json::from_str(&json).unwrap();
141        assert_eq!(e, parsed);
142    }
143}
dynamo_protocols/types/embeddings.rs

dynamo_protocols/types/
embeddings.rs
