aither_openai/
embedding.rs

1use crate::{
2    client::{Config, OpenAI},
3    error::OpenAIError,
4};
5use aither_core::{EmbeddingModel, Result as CoreResult};
6use serde::{Deserialize, Serialize};
7use std::sync::Arc;
8use zenwave::{Client, client, header};
9
10impl EmbeddingModel for OpenAI {
11    fn dim(&self) -> usize {
12        self.config().embedding_dimensions
13    }
14
15    fn embed(&self, text: &str) -> impl core::future::Future<Output = CoreResult<Vec<f32>>> + Send {
16        let cfg = self.config();
17        let input = text.to_owned();
18        async move {
19            let vector = embed_once(cfg, input).await?;
20            Ok(vector)
21        }
22    }
23}
24
25async fn embed_once(cfg: Arc<Config>, input: String) -> Result<Vec<f32>, OpenAIError> {
26    let endpoint = cfg.request_url("/embeddings");
27    let mut backend = client();
28    let mut builder = backend
29        .post(endpoint)
30        .map_err(OpenAIError::Http)?
31        .header(header::AUTHORIZATION.as_str(), cfg.request_auth())
32        .map_err(OpenAIError::Http)?
33        .header(header::USER_AGENT.as_str(), "aither-openai/0.1")
34        .map_err(OpenAIError::Http)?;
35    if let Some(org) = &cfg.organization {
36        builder = builder
37            .header("OpenAI-Organization", org.clone())
38            .map_err(OpenAIError::Http)?;
39    }
40    let request = EmbeddingRequest {
41        model: &cfg.embedding_model,
42        input: &input,
43        dimensions: embedding_dimensions_for(&cfg.embedding_model, cfg.embedding_dimensions),
44    };
45    let response: EmbeddingResponse = builder
46        .json_body(&request)
47        .map_err(OpenAIError::Http)?
48        .json()
49        .await
50        .map_err(OpenAIError::Http)?;
51    response
52        .data
53        .into_iter()
54        .next()
55        .map(|item| item.embedding)
56        .ok_or_else(|| OpenAIError::Api("embedding response missing vector data".into()))
57}
58
59#[derive(Debug, Serialize)]
60struct EmbeddingRequest<'a> {
61    model: &'a str,
62    input: &'a str,
63    #[serde(skip_serializing_if = "Option::is_none")]
64    dimensions: Option<usize>,
65}
66
67#[derive(Debug, Deserialize)]
68struct EmbeddingResponse {
69    data: Vec<EmbeddingItem>,
70}
71
72#[derive(Debug, Deserialize)]
73struct EmbeddingItem {
74    embedding: Vec<f32>,
75}
76
77fn embedding_dimensions_for(model: &str, dimensions: usize) -> Option<usize> {
78    match model {
79        "text-embedding-3-large" | "text-embedding-3-small" => Some(dimensions),
80        _ => None,
81    }
82}
aither_openai/embedding.rs

aither_openai/
embedding.rs
