Merge pull request #4 from spiceai/jeadie/24-11-08/base64

simplify embedding format in response

Author: Jeadie

async-openai/src/embedding.rs

@@ -1,10 +1,7 @@
 use crate::{
     config::Config,
     error::OpenAIError,
-    types::{
-        CreateBase64EmbeddingResponse, CreateEmbeddingRequest, CreateEmbeddingResponse,
-        EncodingFormat,
-    },
+    types::{CreateEmbeddingRequest, CreateEmbeddingResponse, EncodingFormat},
     Client,
 };
 
@@ -25,6 +22,16 @@ impl<'c, C: Config> Embeddings<'c, C> {
     pub async fn create(
         &self,
         request: CreateEmbeddingRequest,
+    ) -> Result<CreateEmbeddingResponse, OpenAIError> {
+        self.client.post("/embeddings", request).await
+    }
+
+    /// Creates an embedding vector representing the input text.
+    ///
+    /// The response will contain the embedding in float-vector format.
+    pub async fn create_float(
+        &self,
+        request: CreateEmbeddingRequest,
     ) -> Result<CreateEmbeddingResponse, OpenAIError> {
         if matches!(request.encoding_format, Some(EncodingFormat::Base64)) {
             return Err(OpenAIError::InvalidArgument(
@@ -40,10 +47,10 @@ impl<'c, C: Config> Embeddings<'c, C> {
     pub async fn create_base64(
         &self,
         request: CreateEmbeddingRequest,
-    ) -> Result<CreateBase64EmbeddingResponse, OpenAIError> {
+    ) -> Result<CreateEmbeddingResponse, OpenAIError> {
         if !matches!(request.encoding_format, Some(EncodingFormat::Base64)) {
             return Err(OpenAIError::InvalidArgument(
-                "When encoding_format is not base64, use Embeddings::create".into(),
+                "When encoding_format is not base64, use Embeddings::create_float".into(),
             ));
         }
 
@@ -166,7 +173,7 @@ mod tests {
             .encoding_format(EncodingFormat::Base64)
             .build()
             .unwrap();
-        let b64_response = client.embeddings().create(b64_request).await;
+        let b64_response = client.embeddings().create_float(b64_request).await;
         assert!(matches!(b64_response, Err(OpenAIError::InvalidArgument(_))));
     }
 
@@ -196,8 +203,8 @@ mod tests {
             .input(INPUT)
             .build()
             .unwrap();
-        let response = client.embeddings().create(request).await.unwrap();
-        let embedding = response.data.into_iter().next().unwrap().embedding;
+        let response = client.embeddings().create_float(request).await.unwrap();
+        let embedding: Vec<f32> = response.data.into_iter().next().unwrap().embedding.into();
 
         assert_eq!(b64_embedding.len(), embedding.len());
         for (b64, normal) in b64_embedding.iter().zip(embedding.iter()) {

async-openai/src/types/embedding.rs

@@ -53,42 +53,79 @@ pub struct CreateEmbeddingRequest {
     pub dimensions: Option<u32>,
 }
 
-/// Represents an embedding vector returned by embedding endpoint.
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
-pub struct Embedding {
-    /// The index of the embedding in the list of embeddings.
-    pub index: u32,
-    /// The object type, which is always "embedding".
-    pub object: String,
-    /// The embedding vector, which is a list of floats. The length of vector
-    /// depends on the model as listed in the [embedding guide](https://platform.openai.com/docs/guides/embeddings).
-    pub embedding: Vec<f32>,
+#[serde(untagged)]
+pub enum EmbeddingVector {
+    Float(Vec<f32>),
+    Base64(String),
 }
 
-#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
-pub struct Base64EmbeddingVector(pub String);
-
-impl From<Base64EmbeddingVector> for Vec<f32> {
-    fn from(value: Base64EmbeddingVector) -> Self {
-        let bytes = general_purpose::STANDARD
-            .decode(value.0)
-            .expect("openai base64 encoding to be valid");
-        let chunks = bytes.chunks_exact(4);
-        chunks
-            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
-            .collect()
+impl From<EmbeddingVector> for Vec<f32> {
+    fn from(val: EmbeddingVector) -> Self {
+        match val {
+            EmbeddingVector::Float(v) => v,
+            EmbeddingVector::Base64(s) => {
+                let bytes = general_purpose::STANDARD
+                    .decode(s)
+                    .expect("openai base64 encoding to be valid");
+                let chunks = bytes.chunks_exact(4);
+                chunks
+                    .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
+                    .collect::<Vec<f32>>()
+            }
+        }
+    }
+}
+
+/// Converts an embedding vector to a base64-encoded string.
+impl From<EmbeddingVector> for String {
+    fn from(val: EmbeddingVector) -> Self {
+        match val {
+            EmbeddingVector::Float(v) => {
+                let mut bytes = Vec::with_capacity(v.len() * 4);
+                for f in v {
+                    bytes.extend_from_slice(&f.to_le_bytes());
+                }
+                general_purpose::STANDARD.encode(&bytes)
+            }
+            EmbeddingVector::Base64(s) => s,
+        }
+    }
+}
+
+impl EmbeddingVector {
+    pub fn is_empty(&self) -> bool {
+        match self {
+            EmbeddingVector::Float(v) => v.is_empty(),
+
+            // Don't use .len() to avoid decoding the base64 string
+            EmbeddingVector::Base64(v) => v.is_empty(),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        match self {
+            EmbeddingVector::Float(v) => v.len(),
+            EmbeddingVector::Base64(v) => {
+                let bytes = general_purpose::STANDARD
+                    .decode(v)
+                    .expect("openai base64 encoding to be valid");
+                bytes.len() / 4
+            }
+        }
     }
 }
 
-/// Represents an base64-encoded embedding vector returned by embedding endpoint.
+/// Represents an embedding vector returned by embedding endpoint.
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
-pub struct Base64Embedding {
+pub struct Embedding {
     /// The index of the embedding in the list of embeddings.
     pub index: u32,
     /// The object type, which is always "embedding".
     pub object: String,
-    /// The embedding vector, encoded in base64.
-    pub embedding: Base64EmbeddingVector,
+    /// The embedding vector, which is a list of floats. The length of vector
+    /// depends on the model as listed in the [embedding guide](https://platform.openai.com/docs/guides/embeddings).
+    pub embedding: EmbeddingVector,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
@@ -109,14 +146,3 @@ pub struct CreateEmbeddingResponse {
     /// The usage information for the request.
     pub usage: EmbeddingUsage,
 }
-
-#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)]
-pub struct CreateBase64EmbeddingResponse {
-    pub object: String,
-    /// The name of the model used to generate the embedding.
-    pub model: String,
-    /// The list of embeddings generated by the model.
-    pub data: Vec<Base64Embedding>,
-    /// The usage information for the request.
-    pub usage: EmbeddingUsage,
-}