use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use super::provider::{EmbeddingProvider, EmbeddingResponse, EmbeddingUsage};
use crate::error::{IndexerError, Result};
pub struct OpenAIProvider {
client: Client,
endpoint: String,
model: String,
api_key: String,
}
#[derive(Serialize)]
struct OpenAIEmbedRequest {
input: Vec<String>,
model: String,
}
#[derive(Deserialize)]
struct OpenAIEmbedResponse {
data: Vec<OpenAIEmbedData>,
model: String,
usage: OpenAIUsage,
}
#[derive(Deserialize)]
struct OpenAIEmbedData {
embedding: Vec<f32>,
#[allow(dead_code)]
index: usize,
#[allow(dead_code)]
object: String,
}
#[derive(Deserialize)]
struct OpenAIUsage {
prompt_tokens: u32,
total_tokens: u32,
}
impl OpenAIProvider {
pub fn new(endpoint: String, model: String, api_key: String) -> Self {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(60)) .build()
.unwrap_or_else(|_| Client::new());
Self {
client,
endpoint,
model,
api_key,
}
}
}
#[async_trait]
impl EmbeddingProvider for OpenAIProvider {
fn model_name(&self) -> &str {
&self.model
}
fn embedding_dimension(&self) -> usize {
match self.model.as_str() {
"text-embedding-3-large" => 3072,
_ => 1536,
}
}
async fn generate_embeddings(&self, texts: Vec<String>) -> Result<EmbeddingResponse> {
let request = OpenAIEmbedRequest {
input: texts,
model: self.model.clone(),
};
let response = self
.client
.post(format!("{}/v1/embeddings", self.endpoint))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&request)
.send()
.await
.map_err(|e| IndexerError::embedding(format!("Failed to send OpenAI request: {e}")))?;
if !response.status().is_success() {
let status = response.status();
return Err(IndexerError::embedding(format!(
"OpenAI API returned error: {status}"
)));
}
let openai_response: OpenAIEmbedResponse = response.json().await.map_err(|e| {
IndexerError::embedding(format!("Failed to parse OpenAI response: {e}"))
})?;
let embeddings = openai_response
.data
.into_iter()
.map(|data| data.embedding)
.collect();
Ok(EmbeddingResponse {
embeddings,
model: openai_response.model,
usage: Some(EmbeddingUsage {
prompt_tokens: Some(openai_response.usage.prompt_tokens),
total_tokens: Some(openai_response.usage.total_tokens),
}),
})
}
async fn health_check(&self) -> Result<bool> {
let test_request = OpenAIEmbedRequest {
input: vec!["test".to_string()],
model: self.model.clone(),
};
let response = self
.client
.post(format!("{}/v1/embeddings", self.endpoint))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&test_request)
.send()
.await;
match response {
Ok(resp) => Ok(resp.status().is_success()),
Err(_) => Ok(false),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[tokio::test]
async fn test_new_provider() {
let provider = OpenAIProvider::new(
"https://api.openai.com".to_string(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
assert_eq!(provider.model_name(), "text-embedding-3-small");
assert_eq!(provider.embedding_dimension(), 1536);
}
#[tokio::test]
async fn test_embedding_dimensions() {
let small_provider = OpenAIProvider::new(
"https://api.openai.com".to_string(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
assert_eq!(small_provider.embedding_dimension(), 1536);
let large_provider = OpenAIProvider::new(
"https://api.openai.com".to_string(),
"text-embedding-3-large".to_string(),
"test-key".to_string(),
);
assert_eq!(large_provider.embedding_dimension(), 3072);
let ada_provider = OpenAIProvider::new(
"https://api.openai.com".to_string(),
"text-embedding-ada-002".to_string(),
"test-key".to_string(),
);
assert_eq!(ada_provider.embedding_dimension(), 1536);
}
#[tokio::test]
async fn test_generate_embeddings_success() {
let mock_server = MockServer::start().await;
let response_body = r#"{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0
},
{
"object": "embedding",
"embedding": [0.4, 0.5, 0.6],
"index": 1
}
],
"model": "text-embedding-3-small",
"usage": {
"prompt_tokens": 10,
"total_tokens": 10
}
}"#;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.and(header("authorization", "Bearer test-key"))
.and(header("content-type", "application/json"))
.respond_with(ResponseTemplate::new(200).set_body_string(response_body))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider
.generate_embeddings(vec!["hello".to_string(), "world".to_string()])
.await;
assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.embeddings.len(), 2);
assert_eq!(response.embeddings[0], vec![0.1, 0.2, 0.3]);
assert_eq!(response.embeddings[1], vec![0.4, 0.5, 0.6]);
assert_eq!(response.model, "text-embedding-3-small");
assert!(response.usage.is_some());
assert_eq!(response.usage.unwrap().total_tokens, Some(10));
}
#[tokio::test]
async fn test_generate_embeddings_single_text() {
let mock_server = MockServer::start().await;
let response_body = r#"{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0
}
],
"model": "text-embedding-3-small",
"usage": {
"prompt_tokens": 5,
"total_tokens": 5
}
}"#;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(200).set_body_string(response_body))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider
.generate_embeddings(vec!["hello world".to_string()])
.await;
assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.embeddings.len(), 1);
assert_eq!(response.embeddings[0], vec![0.1, 0.2, 0.3]);
}
#[tokio::test]
async fn test_generate_embeddings_api_error() {
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(
ResponseTemplate::new(401)
.set_body_string(r#"{"error": {"message": "Invalid API key"}}"#),
)
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"invalid-key".to_string(),
);
let result = provider.generate_embeddings(vec!["test".to_string()]).await;
assert!(result.is_err());
let error = result.unwrap_err();
assert!(error.to_string().contains("OpenAI API returned error"));
}
#[tokio::test]
async fn test_generate_embeddings_invalid_json() {
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(200).set_body_string("invalid json"))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider.generate_embeddings(vec!["test".to_string()]).await;
assert!(result.is_err());
let error = result.unwrap_err();
assert!(error
.to_string()
.contains("Failed to parse OpenAI response"));
}
#[tokio::test]
async fn test_health_check_success() {
let mock_server = MockServer::start().await;
let response_body = r#"{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0
}
],
"model": "text-embedding-3-small",
"usage": {
"prompt_tokens": 1,
"total_tokens": 1
}
}"#;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(200).set_body_string(response_body))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider.health_check().await;
assert!(result.is_ok());
assert!(result.unwrap());
}
#[tokio::test]
async fn test_health_check_failure() {
let mock_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(401))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"invalid-key".to_string(),
);
let result = provider.health_check().await;
assert!(result.is_ok());
assert!(!result.unwrap());
}
#[tokio::test]
async fn test_health_check_network_error() {
let provider = OpenAIProvider::new(
"http://invalid-url-that-does-not-exist:9999".to_string(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider.health_check().await;
assert!(result.is_ok());
assert!(!result.unwrap());
}
#[tokio::test]
async fn test_request_headers_and_body() {
let mock_server = MockServer::start().await;
let response_body = r#"{
"object": "list",
"data": [
{
"object": "embedding",
"embedding": [0.1, 0.2, 0.3],
"index": 0
}
],
"model": "text-embedding-ada-002",
"usage": {
"prompt_tokens": 8,
"total_tokens": 8
}
}"#;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.and(header("authorization", "Bearer secret-key"))
.and(header("content-type", "application/json"))
.respond_with(ResponseTemplate::new(200).set_body_string(response_body))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-ada-002".to_string(),
"secret-key".to_string(),
);
let result = provider
.generate_embeddings(vec!["The food was delicious".to_string()])
.await;
assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.model, "text-embedding-ada-002");
}
#[tokio::test]
async fn test_empty_embeddings_list() {
let mock_server = MockServer::start().await;
let response_body = r#"{
"object": "list",
"data": [],
"model": "text-embedding-3-small",
"usage": {
"prompt_tokens": 0,
"total_tokens": 0
}
}"#;
Mock::given(method("POST"))
.and(path("/v1/embeddings"))
.respond_with(ResponseTemplate::new(200).set_body_string(response_body))
.mount(&mock_server)
.await;
let provider = OpenAIProvider::new(
mock_server.uri(),
"text-embedding-3-small".to_string(),
"test-key".to_string(),
);
let result = provider.generate_embeddings(vec!["test".to_string()]).await;
assert!(result.is_ok());
let response = result.unwrap();
assert_eq!(response.embeddings.len(), 0);
assert_eq!(response.model, "text-embedding-3-small");
}
}