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/// An `/v1/embeddings` client for any OpenAI-compatible server.
44#[derive(Debug)]
45pub struct OpenAiCompatEmbedder {
46    url: String,
47    model: String,
48    api_key: Option<String>,
49    dim: usize,
50    agent: ureq::Agent,
51}
52
53impl OpenAiCompatEmbedder {
54    /// Creates a client for `base_url` (e.g. `https://api.openai.com/v1`
55    /// or `http://localhost:11434/v1`), a model name and the expected
56    /// dimension. The dimension is explicit — no startup probe request,
57    /// and a server disagreeing with it is a typed error, not a silently
58    /// reconfigured database.
59    pub fn new(base_url: &str, model: &str, dim: usize) -> Self {
60        Self {
61            url: format!("{}/embeddings", base_url.trim_end_matches('/')),
62            model: model.to_string(),
63            api_key: None,
64            dim,
65            agent: ureq::Agent::new_with_defaults(),
66        }
67    }
68
69    /// Adds a bearer API key (OpenAI et al.; local servers usually need
70    /// none).
71    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
72        self.api_key = Some(key.into());
73        self
74    }
75}
76
77impl Embedder for OpenAiCompatEmbedder {
78    fn dim(&self) -> usize {
79        self.dim
80    }
81
82    fn embed(&mut self, texts: &[&str]) -> Result<Vec<Vec<f32>>, HostError> {
83        if texts.is_empty() {
84            return Ok(Vec::new());
85        }
86        let body = serde_json::json!({ "model": self.model, "input": texts });
87        let mut request = self.agent.post(&self.url);
88        if let Some(key) = &self.api_key {
89            request = request.header("Authorization", &format!("Bearer {key}"));
90        }
91        let mut response = request
92            .send_json(&body)
93            .map_err(|e| HostError::Embed(format!("request to {}: {e}", self.url)))?;
94        let value: serde_json::Value = response
95            .body_mut()
96            .read_json()
97            .map_err(|e| HostError::Embed(format!("response body: {e}")))?;
98
99        // { "data": [ { "index": i, "embedding": [f32...] }, ... ] } —
100        // placed by the `index` field, per the contract (providers may
101        // reorder).
102        let data = value
103            .get("data")
104            .and_then(|d| d.as_array())
105            .ok_or_else(|| HostError::Embed("response has no data array".into()))?;
106        if data.len() != texts.len() {
107            return Err(HostError::Embed(format!(
108                "expected {} embeddings, got {}",
109                texts.len(),
110                data.len()
111            )));
112        }
113        let mut out = vec![Vec::new(); texts.len()];
114        for item in data {
115            let index = item
116                .get("index")
117                .and_then(|i| i.as_u64())
118                .ok_or_else(|| HostError::Embed("embedding without an index".into()))?
119                as usize;
120            let raw = item
121                .get("embedding")
122                .and_then(|e| e.as_array())
123                .ok_or_else(|| HostError::Embed("embedding is not an array".into()))?;
124            if index >= out.len() || !out[index].is_empty() {
125                return Err(HostError::Embed(format!("bad embedding index {index}")));
126            }
127            if raw.len() != self.dim {
128                return Err(HostError::Embed(format!(
129                    "dimension mismatch: server sent {}, configured {}",
130                    raw.len(),
131                    self.dim
132                )));
133            }
134            let mut v = Vec::with_capacity(raw.len());
135            for x in raw {
136                v.push(
137                    x.as_f64().ok_or_else(|| {
138                        HostError::Embed("embedding component is not a number".into())
139                    })? as f32,
140                );
141            }
142            out[index] = v;
143        }
144        Ok(out)
145    }
146}
plugmem_host/embedder.rs

plugmem_host/
embedder.rs
