plugmem_host/
embedder.rs

1//! The embedder contract and its implementations.
2//!
3//! The engine takes ready vectors; computing them is the host's job.
4//! One HTTP client covers the whole OpenAI-compatible ecosystem —
5//! OpenAI itself, Ollama (`http://localhost:11434/v1`), LM Studio,
6//! vLLM, llama.cpp-server — because they all speak the same
7//! `/v1/embeddings` shape; a provider-specific client would be a second
8//! implementation of the same JSON (records the decision).
9
10use crate::error::HostError;
11
12/// Turns texts into embedding vectors. Batched by design — providers
13/// price and perform far better on batches.
14pub trait Embedder: Send {
15    /// Vector dimension this embedder produces. `0` disables the vector
16    /// layer (the engine is fully functional without it).
17    fn dim(&self) -> usize;
18
19    /// Embeds every text, one vector per input, in input order.
20    ///
21    /// # Errors
22    ///
23    /// [`HostError::Embed`] describing the transport or response
24    /// problem.
25    fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError>;
26}
27
28/// The no-op embedder: dimension 0, never called by the database (a
29/// structural-only memory).
30#[derive(Clone, Copy, Debug, Default)]
31pub struct NullEmbedder;
32
33impl Embedder for NullEmbedder {
34    fn dim(&self) -> usize {
35        0
36    }
37
38    fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
39        Ok(vec![Vec::new(); texts.len()])
40    }
41}
42
43/// One embedder handed to several databases.
44///
45/// [`crate::DatabaseBuilder::embedder`] takes ownership, which is right for one
46/// database and wrong for a workspace: a hundred chats do not want a hundred
47/// HTTP clients pointed at the same endpoint. Each database gets its own
48/// `SharedEmbedder` over one shared provider instead.
49///
50/// Calls serialize on the inner mutex. That is not a compromise — the provider
51/// behind it is a single service, and the host already serializes embedding
52/// within a database. What matters is that the wait happens *outside* the
53/// engine lock, which is still true: the database embeds before it takes its
54/// own lock.
55#[derive(Clone)]
56pub struct SharedEmbedder(std::sync::Arc<std::sync::Mutex<Box<dyn Embedder>>>);
57
58impl SharedEmbedder {
59    /// Wraps `inner` so it can be cloned into many databases.
60    pub fn new(inner: Box<dyn Embedder>) -> Self {
61        Self(std::sync::Arc::new(std::sync::Mutex::new(inner)))
62    }
63
64    /// The inner provider. A panic in one caller's `embed` leaves the provider
65    /// itself intact (it is an HTTP client, not a half-mutated structure), so a
66    /// poisoned mutex is recovered rather than propagated — the same rule the
67    /// engine lock follows.
68    fn inner(&self) -> std::sync::MutexGuard<'_, Box<dyn Embedder>> {
69        self.0.lock().unwrap_or_else(|e| e.into_inner())
70    }
71}
72
73impl std::fmt::Debug for SharedEmbedder {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        f.debug_struct("SharedEmbedder")
76            .field("dim", &self.dim())
77            .finish()
78    }
79}
80
81impl Embedder for SharedEmbedder {
82    fn dim(&self) -> usize {
83        self.inner().dim()
84    }
85
86    fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
87        self.inner().embed(texts)
88    }
89}
90
91/// An `/v1/embeddings` client for any OpenAI-compatible server.
92#[derive(Debug)]
93pub struct OpenAiCompatEmbedder {
94    url: String,
95    model: String,
96    api_key: Option<String>,
97    dim: usize,
98    agent: ureq::Agent,
99}
100
101impl OpenAiCompatEmbedder {
102    /// Creates a client for `base_url` (e.g. `https://api.openai.com/v1`
103    /// or `http://localhost:11434/v1`), a model name and the expected
104    /// dimension. The dimension is explicit — no startup probe request,
105    /// and a server disagreeing with it is a typed error, not a silently
106    /// reconfigured database.
107    pub fn new(base_url: &str, model: &str, dim: usize) -> Self {
108        Self {
109            url: format!("{}/embeddings", base_url.trim_end_matches('/')),
110            model: model.to_string(),
111            api_key: None,
112            dim,
113            agent: ureq::Agent::new_with_defaults(),
114        }
115    }
116
117    /// Adds a bearer API key (OpenAI et al.; local servers usually need
118    /// none).
119    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
120        self.api_key = Some(key.into());
121        self
122    }
123}
124
125impl Embedder for OpenAiCompatEmbedder {
126    fn dim(&self) -> usize {
127        self.dim
128    }
129
130    fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
131        if texts.is_empty() {
132            return Ok(Vec::new());
133        }
134        let body = serde_json::json!({ "model": self.model, "input": texts });
135        let mut request = self.agent.post(&self.url);
136        if let Some(key) = &self.api_key {
137            request = request.header("Authorization", &format!("Bearer {key}"));
138        }
139        let mut response = request
140            .send_json(&body)
141            .map_err(|e| HostError::Embed(format!("request to {}: {e}", self.url)))?;
142        let value: serde_json::Value = response
143            .body_mut()
144            .read_json()
145            .map_err(|e| HostError::Embed(format!("response body: {e}")))?;
146
147        // { "data": [ { "index": i, "embedding": [f32...] }, ... ] } —
148        // placed by the `index` field, per the contract (providers may
149        // reorder).
150        let data = value
151            .get("data")
152            .and_then(|d| d.as_array())
153            .ok_or_else(|| HostError::Embed("response has no data array".into()))?;
154        if data.len() != texts.len() {
155            return Err(HostError::Embed(format!(
156                "expected {} embeddings, got {}",
157                texts.len(),
158                data.len()
159            )));
160        }
161        let mut out = vec![Vec::new(); texts.len()];
162        for item in data {
163            let index = item
164                .get("index")
165                .and_then(|i| i.as_u64())
166                .ok_or_else(|| HostError::Embed("embedding without an index".into()))?
167                as usize;
168            let raw = item
169                .get("embedding")
170                .and_then(|e| e.as_array())
171                .ok_or_else(|| HostError::Embed("embedding is not an array".into()))?;
172            if index >= out.len() || !out[index].is_empty() {
173                return Err(HostError::Embed(format!("bad embedding index {index}")));
174            }
175            if raw.len() != self.dim {
176                return Err(HostError::Embed(format!(
177                    "dimension mismatch: server sent {}, configured {}",
178                    raw.len(),
179                    self.dim
180                )));
181            }
182            let mut v = Vec::with_capacity(raw.len());
183            for x in raw {
184                v.push(
185                    x.as_f64().ok_or_else(|| {
186                        HostError::Embed("embedding component is not a number".into())
187                    })? as f32,
188                );
189            }
190            out[index] = v;
191        }
192        Ok(out)
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use super::*;
199
200    /// Counts its calls, so a test can tell one shared provider from several
201    /// independent ones.
202    struct Counting(usize);
203    impl Embedder for Counting {
204        fn dim(&self) -> usize {
205            3
206        }
207        fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
208            self.0 += texts.len();
209            Ok(vec![vec![self.0 as f32; 3]; texts.len()])
210        }
211    }
212
213    #[test]
214    fn clones_of_a_shared_embedder_reach_the_same_provider() {
215        let shared = SharedEmbedder::new(Box::new(Counting(0)));
216        let mut a = shared.clone();
217        let mut b = shared.clone();
218
219        assert_eq!(a.dim(), 3);
220        assert_eq!(format!("{shared:?}"), "SharedEmbedder { dim: 3 }");
221
222        // Two databases' worth of handles, one counter behind them: the second
223        // call sees the first one's effect.
224        assert_eq!(a.embed(&["x"]).unwrap(), vec![vec![1.0; 3]]);
225        assert_eq!(b.embed(&["y", "z"]).unwrap(), vec![vec![3.0; 3]; 2]);
226    }
227
228    #[test]
229    fn a_poisoned_provider_is_recovered_rather_than_propagated() {
230        let shared = SharedEmbedder::new(Box::new(Counting(0)));
231        let poisoner = shared.clone();
232        // Panic while holding the lock, exactly as a failing provider would.
233        let _ = std::thread::spawn(move || {
234            let _guard = poisoner.inner();
235            panic!("provider blew up");
236        })
237        .join();
238
239        let mut after = shared.clone();
240        assert_eq!(after.embed(&["still works"]).unwrap().len(), 1);
241    }
242
243    #[test]
244    fn the_null_embedder_produces_one_empty_vector_per_text() {
245        let mut null = NullEmbedder;
246        assert_eq!(null.dim(), 0);
247        assert_eq!(null.embed(&["a", "b"]).unwrap(), vec![Vec::<f32>::new(); 2]);
248    }
249}
plugmem_host/embedder.rs

plugmem_host/
embedder.rs
