mempal_embed/
api.rs

1use serde::{Deserialize, Serialize};
2use serde_json::Value;
3
4use super::{EmbedError, Embedder, Result};
5
6#[derive(Debug, Clone)]
7pub struct ApiEmbedder {
8    endpoint: String,
9    model: Option<String>,
10    dimensions: usize,
11}
12
13impl ApiEmbedder {
14    pub fn new(endpoint: String, model: Option<String>, dimensions: usize) -> Self {
15        Self {
16            endpoint,
17            model,
18            dimensions,
19        }
20    }
21
22    pub fn endpoint(&self) -> &str {
23        &self.endpoint
24    }
25
26    pub fn model(&self) -> Option<&str> {
27        self.model.as_deref()
28    }
29}
30
31#[async_trait::async_trait]
32impl Embedder for ApiEmbedder {
33    async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
34        if texts.is_empty() {
35            return Ok(Vec::new());
36        }
37
38        if is_ollama_embeddings_endpoint(self.endpoint()) {
39            return self.embed_ollama_compatible(texts).await;
40        }
41
42        self.embed_openai_compatible(texts).await
43    }
44
45    fn dimensions(&self) -> usize {
46        self.dimensions
47    }
48
49    fn name(&self) -> &str {
50        "api"
51    }
52}
53
54impl ApiEmbedder {
55    async fn embed_openai_compatible(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
56        let endpoint = self.endpoint().to_string();
57        let response = reqwest::Client::new()
58            .post(self.endpoint())
59            .json(&OpenAiEmbeddingsRequest {
60                model: self.model().unwrap_or("text-embedding-3-small"),
61                input: texts,
62            })
63            .send()
64            .await
65            .map_err(|source| EmbedError::HttpRequest {
66                endpoint: endpoint.clone(),
67                source,
68            })?
69            .error_for_status()
70            .map_err(|source| EmbedError::HttpStatus {
71                endpoint: endpoint.clone(),
72                source,
73            })?
74            .json::<OpenAiEmbeddingsResponse>()
75            .await
76            .map_err(|source| EmbedError::DecodeResponse { endpoint, source })?;
77
78        let vectors = response
79            .data
80            .into_iter()
81            .map(|item| item.embedding)
82            .collect::<Vec<_>>();
83        validate_vectors(&vectors, self.dimensions())?;
84        Ok(vectors)
85    }
86
87    async fn embed_ollama_compatible(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
88        let client = reqwest::Client::new();
89        let mut vectors = Vec::with_capacity(texts.len());
90        let endpoint = self.endpoint().to_string();
91
92        for text in texts {
93            let response = client
94                .post(self.endpoint())
95                .json(&OllamaEmbeddingsRequest {
96                    model: self.model().unwrap_or("nomic-embed-text"),
97                    prompt: text,
98                })
99                .send()
100                .await
101                .map_err(|source| EmbedError::HttpRequest {
102                    endpoint: endpoint.clone(),
103                    source,
104                })?
105                .error_for_status()
106                .map_err(|source| EmbedError::HttpStatus {
107                    endpoint: endpoint.clone(),
108                    source,
109                })?
110                .json::<Value>()
111                .await
112                .map_err(|source| EmbedError::DecodeResponse {
113                    endpoint: endpoint.clone(),
114                    source,
115                })?;
116
117            let embedding = response
118                .get("embedding")
119                .and_then(Value::as_array)
120                .ok_or_else(|| {
121                    EmbedError::InvalidResponse(
122                        "embedding response missing embedding array".to_string(),
123                    )
124                })?
125                .iter()
126                .map(json_number_to_f32)
127                .collect::<Result<Vec<_>>>()?;
128            vectors.push(embedding);
129        }
130
131        validate_vectors(&vectors, self.dimensions())?;
132        Ok(vectors)
133    }
134}
135
136#[derive(Debug, Serialize)]
137struct OpenAiEmbeddingsRequest<'a> {
138    model: &'a str,
139    input: &'a [&'a str],
140}
141
142#[derive(Debug, Deserialize)]
143struct OpenAiEmbeddingsResponse {
144    data: Vec<OpenAiEmbeddingItem>,
145}
146
147#[derive(Debug, Deserialize)]
148struct OpenAiEmbeddingItem {
149    embedding: Vec<f32>,
150}
151
152#[derive(Debug, Serialize)]
153struct OllamaEmbeddingsRequest<'a> {
154    model: &'a str,
155    prompt: &'a str,
156}
157
158fn is_ollama_embeddings_endpoint(endpoint: &str) -> bool {
159    endpoint.trim_end_matches('/').ends_with("/api/embeddings")
160}
161
162fn validate_vectors(vectors: &[Vec<f32>], expected_dimensions: usize) -> Result<()> {
163    if vectors.is_empty() {
164        return Err(EmbedError::EmptyVectors);
165    }
166
167    if let Some(actual) = vectors
168        .iter()
169        .map(Vec::len)
170        .find(|length| *length != expected_dimensions)
171    {
172        return Err(EmbedError::InvalidDimensions {
173            expected: expected_dimensions,
174            actual,
175        });
176    }
177
178    Ok(())
179}
180
181fn json_number_to_f32(value: &Value) -> Result<f32> {
182    value
183        .as_f64()
184        .map(|number| number as f32)
185        .ok_or_else(|| EmbedError::InvalidResponse("embedding element was not numeric".to_string()))
186}
mempal_embed/api.rs

mempal_embed/
api.rs
