use serde::{Deserialize, Serialize};
use crate::embeddings::{Embedding, EmbeddingsRequest, EmbeddingsResponse};
use crate::value_objects::{EmbeddingsUsage, Meta};
#[derive(Debug, Serialize)]
pub struct ApiRequest {
pub input: Vec<String>,
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub output_dimension: Option<u32>,
}
#[derive(Debug, Deserialize)]
pub struct ApiResponse {
pub data: Vec<ApiEmbedding>,
pub model: String,
pub usage: ApiUsage,
}
#[derive(Debug, Deserialize)]
pub struct ApiEmbedding {
pub index: usize,
pub embedding: Vec<f32>,
}
#[derive(Debug, Deserialize)]
pub struct ApiUsage {
pub total_tokens: u32,
}
pub fn build_request(request: &EmbeddingsRequest) -> ApiRequest {
ApiRequest {
input: request.input.clone(),
model: request.model.clone(),
output_dimension: request.dimensions,
}
}
pub fn parse_response(mut response: ApiResponse) -> EmbeddingsResponse {
response.data.sort_by_key(|embedding| embedding.index);
EmbeddingsResponse {
embeddings: response
.data
.into_iter()
.map(|embedding| Embedding {
vector: embedding.embedding,
})
.collect(),
usage: EmbeddingsUsage {
prompt_tokens: response.usage.total_tokens,
},
meta: Meta {
id: None,
model: Some(response.model),
rate_limits: Vec::new(),
},
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_request_passes_dimensions_through_as_output_dimension() {
let mut request = EmbeddingsRequest::new("voyage-3-large");
request.dimensions = Some(512);
let wire_request = build_request(&request);
assert_eq!(wire_request.output_dimension, Some(512));
}
}