From e037e41dcec599e8affc7807b0f92075b7b21b31 Mon Sep 17 00:00:00 2001 From: DireMunchkin <1665676+DireMunchkin@users.noreply.github.com> Date: Fri, 19 Jun 2026 15:05:48 +0200 Subject: [PATCH] fix: handle base64 in text embedding response --- rust/src/types.rs | 81 +++++++++++++++++++++++++++++++++++- rust/tests/ai_integration.rs | 30 ++++++------- 2 files changed, 95 insertions(+), 16 deletions(-) diff --git a/rust/src/types.rs b/rust/src/types.rs index 6e109a9..c122ae8 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -853,11 +853,18 @@ pub struct EmbeddingResponse { pub usage: EmbeddingUsage, } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum EmbeddingValues { + Floats(Vec), + Base64(String), +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct EmbeddingData { pub object: String, pub index: i32, - pub embedding: Vec, + pub embedding: EmbeddingValues, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -1143,4 +1150,76 @@ mod tests { assert_eq!(response.access_token.as_deref(), Some("new-access")); assert_eq!(response.refresh_token, None); } + + #[test] + fn embedding_values_deserializes_float_array() { + let data: EmbeddingData = serde_json::from_value(json!({ + "object": "embedding", + "index": 0, + "embedding": [0.1, 0.2, 0.3] + })) + .unwrap(); + + assert_eq!(data.object, "embedding"); + assert_eq!(data.index, 0); + match &data.embedding { + EmbeddingValues::Floats(v) => assert_eq!(v, &[0.1, 0.2, 0.3]), + EmbeddingValues::Base64(_) => panic!("Expected Floats variant"), + } + } + + #[test] + fn embedding_values_deserializes_base64_string() { + let data: EmbeddingData = serde_json::from_value(json!({ + "object": "embedding", + "index": 0, + "embedding": "AQIDBA==" + })) + .unwrap(); + + assert_eq!(data.object, "embedding"); + assert_eq!(data.index, 0); + match &data.embedding { + EmbeddingValues::Base64(s) => assert_eq!(s, "AQIDBA=="), + EmbeddingValues::Floats(_) => panic!("Expected Base64 variant"), + } + } + + #[test] + fn embedding_response_deserializes_with_floats() { + let response: EmbeddingResponse = serde_json::from_value(json!({ + "object": "list", + "data": [{ + "object": "embedding", + "index": 0, + "embedding": [0.1, 0.2, 0.3] + }], + "model": "nomic-embed-text", + "usage": { "prompt_tokens": 5, "total_tokens": 5 } + })) + .unwrap(); + + assert_eq!(response.object, "list"); + assert_eq!(response.data.len(), 1); + assert!(matches!(response.data[0].embedding, EmbeddingValues::Floats(_))); + } + + #[test] + fn embedding_response_deserializes_with_base64() { + let response: EmbeddingResponse = serde_json::from_value(json!({ + "object": "list", + "data": [{ + "object": "embedding", + "index": 0, + "embedding": "AQIDBA==" + }], + "model": "nomic-embed-text", + "usage": { "prompt_tokens": 5, "total_tokens": 5 } + })) + .unwrap(); + + assert_eq!(response.object, "list"); + assert_eq!(response.data.len(), 1); + assert!(matches!(response.data[0].embedding, EmbeddingValues::Base64(_))); + } } diff --git a/rust/tests/ai_integration.rs b/rust/tests/ai_integration.rs index 0f35397..720e0f0 100644 --- a/rust/tests/ai_integration.rs +++ b/rust/tests/ai_integration.rs @@ -1,7 +1,7 @@ use futures::StreamExt; use opensecret::{ - ChatCompletionRequest, ChatMessage, EmbeddingInput, EmbeddingRequest, Error, Function, - OpenSecretClient, Result, Tool, + ChatCompletionRequest, ChatMessage, EmbeddingInput, EmbeddingRequest, EmbeddingValues, Error, + Function, OpenSecretClient, Result, Tool, }; use std::env; use uuid::Uuid; @@ -426,19 +426,17 @@ async fn test_create_embeddings_single_input() { assert_eq!(response.data[0].index, 0); let expected_dimensions = embedding_dimensions(); - assert_eq!( - response.data[0].embedding.len(), - expected_dimensions, - "Unexpected embedding dimensions" - ); + match &response.data[0].embedding { + EmbeddingValues::Floats(v) => assert_eq!(v.len(), expected_dimensions, "Unexpected embedding dimensions"), + EmbeddingValues::Base64(_) => panic!("Expected float embeddings, got base64"), + } // Verify usage assert!(response.usage.prompt_tokens > 0); assert!(response.usage.total_tokens > 0); println!( - "Embedding created with {} dimensions, {} tokens used", - response.data[0].embedding.len(), + "Embedding created with dimensions from response, {} tokens used", response.usage.total_tokens ); } @@ -474,11 +472,10 @@ async fn test_create_embeddings_multiple_inputs() { for (i, embedding_data) in response.data.iter().enumerate() { assert_eq!(embedding_data.object, "embedding"); assert_eq!(embedding_data.index as usize, i); - assert_eq!( - embedding_data.embedding.len(), - embedding_dimensions(), - "Unexpected embedding dimensions" - ); + match &embedding_data.embedding { + EmbeddingValues::Floats(v) => assert_eq!(v.len(), embedding_dimensions(), "Unexpected embedding dimensions"), + EmbeddingValues::Base64(_) => panic!("Expected float embeddings, got base64"), + } } // Verify usage accounts for all inputs @@ -512,7 +509,10 @@ async fn test_embeddings_from_string_conversion() { .expect("Failed to create embeddings"); assert_eq!(response.data.len(), 1); - assert_eq!(response.data[0].embedding.len(), embedding_dimensions()); + match &response.data[0].embedding { + EmbeddingValues::Floats(v) => assert_eq!(v.len(), embedding_dimensions()), + EmbeddingValues::Base64(_) => panic!("Expected float embeddings, got base64"), + } } #[tokio::test]