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.
14///
15/// `embed` takes `&self`, and the trait requires `Sync`, because an embedder is
16/// a *client* of a remote service, not a piece of mutable state. Every caller
17/// in this workspace shares one instance across threads (a database's writer,
18/// the napi binding's libuv workers, the MCP worker pool), and a `&mut self`
19/// signature forced every one of them to put a `Mutex` in front of it. That
20/// mutex serialized the HTTP round trips: four concurrent recalls against a
21/// 300 ms provider took 1200 ms, with the provider seeing one request at a
22/// time. With `&self` they take 300 ms and the provider sees four.
23///
24/// An implementation that genuinely needs mutable state (a local cache, a
25/// rate-limit budget) brings its own interior mutability, which is the right
26/// place for it: only that implementation knows what may overlap and what may
27/// not. [`OpenAiCompatEmbedder`] needs none — a `ureq::Agent` is a
28/// connection-pool handle built for concurrent use.
29pub trait Embedder: Send + Sync {
30    /// Vector dimension this embedder produces. `0` disables the vector
31    /// layer (the engine is fully functional without it).
32    fn dim(&self) -> usize;
33
34    /// Embeds every text, one vector per input, in input order.
35    ///
36    /// Called concurrently from several threads. An implementation that keeps
37    /// state must guard it itself.
38    ///
39    /// # Errors
40    ///
41    /// [`HostError::Embed`] describing the transport or response
42    /// problem.
43    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError>;
44}
45
46/// The no-op embedder: dimension 0, never called by the database (a
47/// structural-only memory).
48#[derive(Clone, Copy, Debug, Default)]
49pub struct NullEmbedder;
50
51impl Embedder for NullEmbedder {
52    fn dim(&self) -> usize {
53        0
54    }
55
56    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
57        Ok(vec![Vec::new(); texts.len()])
58    }
59}
60
61/// One embedder handed to several databases.
62///
63/// [`crate::DatabaseBuilder::embedder`] takes ownership, which is right for one
64/// database and wrong for a workspace: a hundred chats do not want a hundred
65/// HTTP clients pointed at the same endpoint. Each database gets its own
66/// `SharedEmbedder` over one shared provider instead.
67///
68/// A plain refcount, with no lock in it. Sharing an [`Embedder`] needs nothing
69/// more, because `embed` takes `&self`; this type used to hold a `Mutex` and
70/// that mutex was the only reason concurrent embedding serialized. Cloning is
71/// an atomic increment, so handing one out per database costs nothing.
72#[derive(Clone)]
73pub struct SharedEmbedder(std::sync::Arc<dyn Embedder>);
74
75impl SharedEmbedder {
76    /// Wraps `inner` so it can be cloned into many databases.
77    pub fn new(inner: Box<dyn Embedder>) -> Self {
78        Self(std::sync::Arc::from(inner))
79    }
80}
81
82impl std::fmt::Debug for SharedEmbedder {
83    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84        f.debug_struct("SharedEmbedder")
85            .field("dim", &self.0.dim())
86            .finish()
87    }
88}
89
90impl Embedder for SharedEmbedder {
91    fn dim(&self) -> usize {
92        self.0.dim()
93    }
94
95    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
96        self.0.embed(texts)
97    }
98}
99
100/// An `/v1/embeddings` client for any OpenAI-compatible server.
101#[derive(Debug)]
102pub struct OpenAiCompatEmbedder {
103    url: String,
104    model: String,
105    api_key: Option<String>,
106    dim: usize,
107    agent: ureq::Agent,
108}
109
110impl OpenAiCompatEmbedder {
111    /// Creates a client for `base_url` (e.g. `https://api.openai.com/v1`
112    /// or `http://localhost:11434/v1`), a model name and the expected
113    /// dimension. The dimension is explicit — no startup probe request,
114    /// and a server disagreeing with it is a typed error, not a silently
115    /// reconfigured database.
116    pub fn new(base_url: &str, model: &str, dim: usize) -> Self {
117        Self {
118            url: format!("{}/embeddings", base_url.trim_end_matches('/')),
119            model: model.to_string(),
120            api_key: None,
121            dim,
122            agent: ureq::Agent::new_with_defaults(),
123        }
124    }
125
126    /// Adds a bearer API key (OpenAI et al.; local servers usually need
127    /// none).
128    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
129        self.api_key = Some(key.into());
130        self
131    }
132}
133
134impl Embedder for OpenAiCompatEmbedder {
135    fn dim(&self) -> usize {
136        self.dim
137    }
138
139    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
140        if texts.is_empty() {
141            return Ok(Vec::new());
142        }
143        let body = serde_json::json!({ "model": self.model, "input": texts });
144        let mut request = self.agent.post(&self.url);
145        if let Some(key) = &self.api_key {
146            request = request.header("Authorization", &format!("Bearer {key}"));
147        }
148        let mut response = request
149            .send_json(&body)
150            .map_err(|e| HostError::Embed(format!("request to {}: {e}", self.url)))?;
151        let value: serde_json::Value = response
152            .body_mut()
153            .read_json()
154            .map_err(|e| HostError::Embed(format!("response body: {e}")))?;
155
156        // { "data": [ { "index": i, "embedding": [f32...] }, ... ] } —
157        // placed by the `index` field, per the contract (providers may
158        // reorder).
159        let data = value
160            .get("data")
161            .and_then(|d| d.as_array())
162            .ok_or_else(|| HostError::Embed("response has no data array".into()))?;
163        if data.len() != texts.len() {
164            return Err(HostError::Embed(format!(
165                "expected {} embeddings, got {}",
166                texts.len(),
167                data.len()
168            )));
169        }
170        let mut out = vec![Vec::new(); texts.len()];
171        for item in data {
172            let index = item
173                .get("index")
174                .and_then(|i| i.as_u64())
175                .ok_or_else(|| HostError::Embed("embedding without an index".into()))?
176                as usize;
177            let raw = item
178                .get("embedding")
179                .and_then(|e| e.as_array())
180                .ok_or_else(|| HostError::Embed("embedding is not an array".into()))?;
181            if index >= out.len() || !out[index].is_empty() {
182                return Err(HostError::Embed(format!("bad embedding index {index}")));
183            }
184            if raw.len() != self.dim {
185                return Err(HostError::Embed(format!(
186                    "dimension mismatch: server sent {}, configured {}",
187                    raw.len(),
188                    self.dim
189                )));
190            }
191            let mut v = Vec::with_capacity(raw.len());
192            for x in raw {
193                v.push(
194                    x.as_f64().ok_or_else(|| {
195                        HostError::Embed("embedding component is not a number".into())
196                    })? as f32,
197                );
198            }
199            out[index] = v;
200        }
201        Ok(out)
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    use std::sync::atomic::{AtomicUsize, Ordering};
210
211    /// Counts its calls, so a test can tell one shared provider from several
212    /// independent ones. State behind an atomic because `embed` takes `&self`
213    /// — the arrangement the trait asks a stateful implementation to make.
214    struct Counting(AtomicUsize);
215    impl Embedder for Counting {
216        fn dim(&self) -> usize {
217            3
218        }
219        fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
220            let total = self.0.fetch_add(texts.len(), Ordering::Relaxed) + texts.len();
221            Ok(vec![vec![total as f32; 3]; texts.len()])
222        }
223    }
224
225    /// Blocks for a moment and records how many calls were inside `embed` at
226    /// once, which is what a mutex in front of it would hold at one.
227    struct Overlapping {
228        inside: AtomicUsize,
229        peak: AtomicUsize,
230    }
231    impl Embedder for Overlapping {
232        fn dim(&self) -> usize {
233            1
234        }
235        fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
236            let now = self.inside.fetch_add(1, Ordering::SeqCst) + 1;
237            self.peak.fetch_max(now, Ordering::SeqCst);
238            std::thread::sleep(std::time::Duration::from_millis(50));
239            self.inside.fetch_sub(1, Ordering::SeqCst);
240            Ok(vec![vec![0.0]; texts.len()])
241        }
242    }
243
244    #[test]
245    fn clones_of_a_shared_embedder_reach_the_same_provider() {
246        let shared = SharedEmbedder::new(Box::new(Counting(AtomicUsize::new(0))));
247        let a = shared.clone();
248        let b = shared.clone();
249
250        assert_eq!(a.dim(), 3);
251        assert_eq!(format!("{shared:?}"), "SharedEmbedder { dim: 3 }");
252
253        // Two databases' worth of handles, one counter behind them: the second
254        // call sees the first one's effect.
255        assert_eq!(a.embed(&["x"]).unwrap(), vec![vec![1.0; 3]]);
256        assert_eq!(b.embed(&["y", "z"]).unwrap(), vec![vec![3.0; 3]; 2]);
257    }
258
259    #[test]
260    fn concurrent_callers_are_inside_the_provider_at_the_same_time() {
261        // The invariant the `&self` signature exists for. Under the old
262        // `&mut self` trait this could not be written at all: the `Mutex` that
263        // a shared embedder needed held `peak` at 1, and four concurrent
264        // recalls against a slow provider cost four round trips.
265        let provider = std::sync::Arc::new(Overlapping {
266            inside: AtomicUsize::new(0),
267            peak: AtomicUsize::new(0),
268        });
269        let shared = SharedEmbedder(provider.clone());
270
271        std::thread::scope(|scope| {
272            for _ in 0..4 {
273                let handle = shared.clone();
274                scope.spawn(move || handle.embed(&["question"]).unwrap());
275            }
276        });
277
278        assert!(
279            provider.peak.load(Ordering::SeqCst) > 1,
280            "callers serialized: peak concurrency was {}",
281            provider.peak.load(Ordering::SeqCst)
282        );
283    }
284
285    #[test]
286    fn the_null_embedder_produces_one_empty_vector_per_text() {
287        let null = NullEmbedder;
288        assert_eq!(null.dim(), 0);
289        assert_eq!(null.embed(&["a", "b"]).unwrap(), vec![Vec::<f32>::new(); 2]);
290    }
291}
plugmem_host/embedder.rs

plugmem_host/
embedder.rs
