emb: lines into vectors, against a server that owns the model

crates/emb/src/lib.rs

@@ -0,0 +1,275 @@
+//! Lines in, vectors out, against an HTTP embedding endpoint.
+//!
+//! The one tool here that touches the network, and it does so for a reason:
+//! an embedding model is half a gigabyte of weights and a GPU, which belongs
+//! behind a server rather than inside a filter. What this owns is everything
+//! around the call — batching, the defensive truncation, the retry, and the
+//! assertion that the vectors coming back are the width the index expects.
+//!
+//! The request shape is the one Ollama's `/api/embed` uses, which several
+//! other servers also speak.
+
+use std::time::Duration;
+
+use serde::{Deserialize, Serialize};
+
+/// How many texts go in one request.
+pub const BATCH: usize = 32;
+
+/// Tokens a single input is cut to before it is sent.
+///
+/// A defensive cap, not a model limit. A server returns an error for an
+/// input past its context length and that error fails the whole batch, so one
+/// oversized passage costs thirty-one good ones. Counted with a tokeniser
+/// that is not the model's — it is within ten or fifteen per cent, which is
+/// all a safety net needs to be.
+pub const MAX_INPUT_TOKENS: usize = 1800;
+
+/// How long to wait for a response.
+///
+/// Generous, because a cold server loading a model onto a GPU can take
+/// minutes before it answers at all.
+pub const TIMEOUT: Duration = Duration::from_secs(600);
+
+/// How many times a failed batch is tried again.
+pub const ATTEMPTS: u32 = 3;
+
+/// What the endpoint is asked.
+#[derive(Debug, Serialize)]
+struct Request<'a> {
+    model: &'a str,
+    input: &'a [String],
+}
+
+/// What it answers.
+#[derive(Debug, Deserialize)]
+struct Response {
+    embeddings: Vec<Vec<f64>>,
+}
+
+/// Why a batch could not be embedded.
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    /// The server could not be reached, or did not answer well.
+    #[error("{0}")]
+    Http(String),
+    /// The answer was not the shape this expects.
+    #[error("the server's answer was not an embedding: {0}")]
+    Shape(String),
+    /// The vectors are not the width the caller asked for.
+    #[error("vectors are {found} wide, expected {expected}")]
+    Width {
+        /// What came back.
+        found: usize,
+        /// What was asked for.
+        expected: usize,
+    },
+    /// The server answered with fewer vectors than it was given texts.
+    #[error("sent {sent} texts and got {received} vectors back")]
+    Count {
+        /// How many texts went out.
+        sent: usize,
+        /// How many vectors came back.
+        received: usize,
+    },
+}
+
+/// Where and how to embed.
+#[derive(Debug, Clone)]
+pub struct Endpoint {
+    /// The server's base URL, without a trailing slash.
+    pub host: String,
+    /// The model to ask for.
+    pub model: String,
+    /// The width the vectors must be, when the caller knows it.
+    pub dim: Option<usize>,
+    /// Texts per request.
+    pub batch: usize,
+    /// Tokens per text, or `None` to send them whole.
+    pub max_input_tokens: Option<usize>,
+    /// How long to wait for one request.
+    pub timeout: Duration,
+    /// How many times to try a failing batch.
+    pub attempts: u32,
+}
+
+impl Endpoint {
+    /// An endpoint with the usual settings.
+    pub fn new(host: impl Into<String>, model: impl Into<String>) -> Self {
+        Self {
+            host: host.into().trim_end_matches('/').to_string(),
+            model: model.into(),
+            dim: None,
+            batch: BATCH,
+            max_input_tokens: Some(MAX_INPUT_TOKENS),
+            timeout: TIMEOUT,
+            attempts: ATTEMPTS,
+        }
+    }
+
+    /// The URL a request goes to.
+    pub fn url(&self) -> String {
+        format!("{}/api/embed", self.host)
+    }
+
+    /// Cut a text to the input cap, if there is one.
+    pub fn truncate(&self, text: &str) -> String {
+        match self.max_input_tokens {
+            None | Some(0) => text.to_string(),
+            Some(limit) => {
+                let tokens = textkit::encode(text);
+                if tokens.len() <= limit {
+                    text.to_string()
+                } else {
+                    textkit::decode_lossy(&tokens[..limit])
+                }
+            }
+        }
+    }
+}
+
+/// Embed one batch, retrying a failure.
+///
+/// Waits one second before the second attempt and two before the third, which
+/// is long enough for a server that is loading a model and short enough that
+/// a server that is simply down does not hold up a pipeline for a minute.
+pub fn embed_batch(
+    agent: &ureq::Agent,
+    endpoint: &Endpoint,
+    texts: &[String],
+) -> Result<Vec<Vec<f64>>, Error> {
+    let sent: Vec<String> = texts.iter().map(|text| endpoint.truncate(text)).collect();
+    let mut last = Error::Http("no attempt was made".into());
+
+    for attempt in 0..endpoint.attempts.max(1) {
+        match post(agent, endpoint, &sent) {
+            Ok(vectors) => return check(endpoint, &sent, vectors),
+            Err(e @ Error::Width { .. } | e @ Error::Count { .. }) => return Err(e),
+            Err(e) => last = e,
+        }
+        std::thread::sleep(Duration::from_secs(1 << attempt));
+    }
+    Err(last)
+}
+
+fn post(
+    agent: &ureq::Agent,
+    endpoint: &Endpoint,
+    texts: &[String],
+) -> Result<Vec<Vec<f64>>, Error> {
+    let request = Request {
+        model: &endpoint.model,
+        input: texts,
+    };
+    let mut response = agent
+        .post(endpoint.url())
+        .send_json(&request)
+        .map_err(|e| Error::Http(e.to_string()))?;
+    let body: Response = response
+        .body_mut()
+        .read_json()
+        .map_err(|e| Error::Shape(e.to_string()))?;
+    Ok(body.embeddings)
+}
+
+fn check(
+    endpoint: &Endpoint,
+    sent: &[String],
+    vectors: Vec<Vec<f64>>,
+) -> Result<Vec<Vec<f64>>, Error> {
+    if vectors.len() != sent.len() {
+        return Err(Error::Count {
+            sent: sent.len(),
+            received: vectors.len(),
+        });
+    }
+    if let (Some(expected), Some(first)) = (endpoint.dim, vectors.first())
+        && first.len() != expected
+    {
+        return Err(Error::Width {
+            found: first.len(),
+            expected,
+        });
+    }
+    Ok(vectors)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn endpoint() -> Endpoint {
+        Endpoint::new("http://example.invalid/", "a-model")
+    }
+
+    #[test]
+    fn a_trailing_slash_does_not_double_up() {
+        assert_eq!(endpoint().url(), "http://example.invalid/api/embed");
+        assert_eq!(
+            Endpoint::new("http://example.invalid", "m").url(),
+            "http://example.invalid/api/embed"
+        );
+    }
+
+    #[test]
+    fn a_short_text_is_sent_whole() {
+        let text = "the mill stood at the northern end";
+        assert_eq!(endpoint().truncate(text), text);
+    }
+
+    #[test]
+    fn a_long_text_is_cut_to_the_cap() {
+        let mut endpoint = endpoint();
+        endpoint.max_input_tokens = Some(4);
+        let cut = endpoint.truncate(&"word ".repeat(100));
+        assert_eq!(textkit::count(&cut), 4);
+    }
+
+    #[test]
+    fn the_cap_can_be_turned_off() {
+        let mut endpoint = endpoint();
+        endpoint.max_input_tokens = None;
+        let long = "word ".repeat(5000);
+        assert_eq!(endpoint.truncate(&long), long);
+    }
+
+    #[test]
+    fn a_width_that_is_not_the_one_asked_for_is_refused() {
+        let mut endpoint = endpoint();
+        endpoint.dim = Some(1024);
+        let sent = vec!["a".to_string()];
+        let e = check(&endpoint, &sent, vec![vec![0.0; 768]]).expect_err("should refuse");
+        assert!(matches!(
+            e,
+            Error::Width {
+                found: 768,
+                expected: 1024
+            }
+        ));
+    }
+
+    #[test]
+    fn a_missing_vector_is_refused() {
+        let sent = vec!["a".to_string(), "b".to_string()];
+        let e = check(&endpoint(), &sent, vec![vec![0.0; 8]]).expect_err("should refuse");
+        assert!(matches!(
+            e,
+            Error::Count {
+                sent: 2,
+                received: 1
+            }
+        ));
+    }
+
+    #[test]
+    fn vectors_of_the_right_shape_come_straight_back() {
+        let mut endpoint = endpoint();
+        endpoint.dim = Some(3);
+        let sent = vec!["a".to_string()];
+        let vectors = vec![vec![1.0, 2.0, 3.0]];
+        assert_eq!(
+            check(&endpoint, &sent, vectors.clone()).expect("should pass"),
+            vectors
+        );
+    }
+}