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

crates/emb/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "emb"
+description = "Lines in, vectors out, against an HTTP embedding endpoint."
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+catena-cli.workspace = true
+clap.workspace = true
+corpus-core.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+textkit.workspace = true
+thiserror.workspace = true
+ureq.workspace = true
+
+[dev-dependencies]
+serde_json.workspace = true
+textkit.workspace = true
+
+[lints]
+workspace = true

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
+        );
+    }
+}

crates/emb/src/main.rs

@@ -0,0 +1,137 @@
+//! Lines in, vectors out, against an HTTP embedding endpoint.
+
+use std::io::Write;
+use std::process::ExitCode;
+use std::time::Duration;
+
+use clap::Parser;
+use corpus_core::exit;
+use emb::Endpoint;
+
+/// Embed each line of the input and write its vector.
+///
+/// One JSON array per line, in the order the lines arrived, so the output
+/// lines up with the input line for line. Empty lines are sent as they are:
+/// dropping them would make the two disagree.
+///
+/// The only tool here that uses the network. It never loads a model — the
+/// model belongs behind a server, where it can be shared and where its GPU
+/// is somebody else's problem.
+#[derive(Parser)]
+#[command(
+    version,
+    about,
+    after_help = "\
+Exit codes: 0 every line was embedded, 65 the server answered something that
+was not an embedding, 69 the server could not be reached, 74 the input could
+not be read.
+
+Example:
+
+    jq -r .text chunks.ndjson | emb --model bge-m3 --dim 1024 > vectors.ndjson"
+)]
+struct Args {
+    /// File of texts, one per line. Reads stdin when absent.
+    path: Option<std::path::PathBuf>,
+
+    /// The server to ask.
+    #[arg(long, default_value = "http://localhost:11434", value_name = "URL")]
+    host: String,
+
+    /// The model to ask for.
+    #[arg(long, value_name = "NAME")]
+    model: String,
+
+    /// Refuse vectors that are not this wide.
+    #[arg(long, value_name = "WIDTH")]
+    dim: Option<usize>,
+
+    /// Texts per request.
+    #[arg(long, default_value_t = emb::BATCH, value_name = "COUNT")]
+    batch: usize,
+
+    /// Cut each text to this many tokens before sending it. Zero sends them
+    /// whole.
+    #[arg(long, default_value_t = emb::MAX_INPUT_TOKENS, value_name = "TOKENS")]
+    max_input_tokens: usize,
+
+    /// How long to wait for one request.
+    #[arg(long, default_value_t = emb::TIMEOUT.as_secs(), value_name = "SECONDS")]
+    timeout: u64,
+
+    /// How many times to try a failing request.
+    #[arg(long, default_value_t = emb::ATTEMPTS, value_name = "COUNT")]
+    attempts: u32,
+}
+
+fn main() -> ExitCode {
+    let args: Args = match catena_cli::parse() {
+        Ok(args) => args,
+        Err(code) => return code,
+    };
+    if args.batch == 0 {
+        eprintln!("emb: --batch must be above zero");
+        return ExitCode::from(exit::USAGE);
+    }
+
+    let endpoint = Endpoint {
+        dim: args.dim,
+        batch: args.batch,
+        max_input_tokens: Some(args.max_input_tokens),
+        timeout: Duration::from_secs(args.timeout),
+        attempts: args.attempts,
+        ..Endpoint::new(&args.host, &args.model)
+    };
+
+    let text = match catena_cli::read_text(args.path.as_deref()) {
+        Ok(text) => text,
+        Err((code, why)) => {
+            eprintln!("emb: {why}");
+            return ExitCode::from(code);
+        }
+    };
+    // A trailing newline ends the last line rather than starting an empty one.
+    let lines: Vec<String> = text
+        .strip_suffix('\n')
+        .unwrap_or(&text)
+        .split('\n')
+        .map(str::to_string)
+        .collect();
+    if text.is_empty() {
+        return ExitCode::from(exit::OK);
+    }
+
+    let agent: ureq::Agent = ureq::Agent::config_builder()
+        .timeout_global(Some(endpoint.timeout))
+        .build()
+        .into();
+
+    let mut out = std::io::BufWriter::new(std::io::stdout().lock());
+    for batch in lines.chunks(endpoint.batch) {
+        let vectors = match emb::embed_batch(&agent, &endpoint, batch) {
+            Ok(vectors) => vectors,
+            Err(e) => {
+                eprintln!("emb: {e}");
+                return ExitCode::from(match e {
+                    emb::Error::Http(_) => exit::UNAVAILABLE,
+                    _ => exit::DATA,
+                });
+            }
+        };
+        for vector in vectors {
+            let line = serde_json::to_string(&vector).expect("a vector serialises");
+            match writeln!(out, "{line}").and_then(|()| out.flush()) {
+                Ok(()) => {}
+                // A consumer that stopped reading is not a failure.
+                Err(e) if e.kind() == std::io::ErrorKind::BrokenPipe => {
+                    return ExitCode::from(exit::OK);
+                }
+                Err(e) => {
+                    eprintln!("emb: {e}");
+                    return ExitCode::from(exit::IO);
+                }
+            }
+        }
+    }
+    catena_cli::finish("emb", out)
+}

crates/emb/tests/cli.rs

@@ -0,0 +1,212 @@
+//! The tool against a server, without needing one.
+//!
+//! The server here is twenty lines of `TcpListener` that answers whatever the
+//! test wants it to. That is enough to check the things that actually go
+//! wrong: the request shape, the batching, the retry, and what happens when
+//! the answer is not what was asked for.
+
+use std::io::{BufRead, BufReader, Read, Write};
+use std::net::{TcpListener, TcpStream};
+use std::process::{Command, Output};
+use std::sync::mpsc;
+use std::thread;
+
+/// What a request asked for, as the test wants to see it.
+struct Asked {
+    model: String,
+    inputs: Vec<String>,
+}
+
+/// Answer `replies` in turn, one per request, and report what was asked.
+fn server(replies: Vec<(u16, String)>) -> (String, mpsc::Receiver<Asked>) {
+    let listener = TcpListener::bind("127.0.0.1:0").expect("a free port");
+    let host = format!("http://{}", listener.local_addr().expect("an address"));
+    let (sender, receiver) = mpsc::channel();
+
+    thread::spawn(move || {
+        for (at, connection) in listener.incoming().enumerate() {
+            let Ok(mut connection) = connection else {
+                break;
+            };
+            if let Some(asked) = serve(&mut connection, replies.get(at)) {
+                let _ = sender.send(asked);
+            }
+            if at + 1 >= replies.len() {
+                break;
+            }
+        }
+    });
+
+    (host, receiver)
+}
+
+fn serve(connection: &mut TcpStream, reply: Option<&(u16, String)>) -> Option<Asked> {
+    let mut reader = BufReader::new(connection.try_clone().expect("a clone"));
+    let mut length = 0usize;
+    loop {
+        let mut line = String::new();
+        if reader.read_line(&mut line).ok()? == 0 {
+            return None;
+        }
+        if let Some(value) = line.to_lowercase().strip_prefix("content-length:") {
+            length = value.trim().parse().unwrap_or(0);
+        }
+        if line == "\r\n" || line == "\n" {
+            break;
+        }
+    }
+    let mut body = vec![0; length];
+    reader.read_exact(&mut body).ok()?;
+    let request: serde_json::Value = serde_json::from_slice(&body).ok()?;
+
+    let (status, reply) = reply.cloned().unwrap_or((500, String::new()));
+    let response = format!(
+        "HTTP/1.1 {status} X\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{reply}",
+        reply.len()
+    );
+    let _ = connection.write_all(response.as_bytes());
+    let _ = connection.flush();
+
+    Some(Asked {
+        model: request["model"].as_str().unwrap_or_default().to_string(),
+        inputs: request["input"]
+            .as_array()
+            .map(|inputs| {
+                inputs
+                    .iter()
+                    .map(|input| input.as_str().unwrap_or_default().to_string())
+                    .collect()
+            })
+            .unwrap_or_default(),
+    })
+}
+
+fn vectors(count: usize, width: usize) -> String {
+    let one = format!("[{}]", vec!["0.5"; width].join(","));
+    format!("{{\"embeddings\":[{}]}}", vec![one; count].join(","))
+}
+
+fn run(args: &[&str], stdin: &str) -> Output {
+    use std::process::Stdio;
+    let mut child = Command::new(env!("CARGO_BIN_EXE_emb"))
+        .args(args)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
+        .expect("the binary should run");
+    child
+        .stdin
+        .take()
+        .expect("stdin was piped")
+        .write_all(stdin.as_bytes())
+        .expect("the input should write");
+    child.wait_with_output().expect("the binary should finish")
+}
+
+#[test]
+fn every_line_comes_back_as_a_vector() {
+    let (host, asked) = server(vec![(200, vectors(2, 4))]);
+    let out = run(
+        &["--host", &host, "--model", "a-model"],
+        "the mill stood\nat the northern end\n",
+    );
+
+    assert_eq!(
+        out.status.code(),
+        Some(0),
+        "{:?}",
+        String::from_utf8_lossy(&out.stderr)
+    );
+    let lines: Vec<&str> = std::str::from_utf8(&out.stdout)
+        .expect("utf-8")
+        .lines()
+        .collect();
+    assert_eq!(lines.len(), 2);
+    assert_eq!(lines[0], "[0.5,0.5,0.5,0.5]");
+
+    let asked = asked.recv().expect("a request");
+    assert_eq!(asked.model, "a-model");
+    assert_eq!(asked.inputs, ["the mill stood", "at the northern end"]);
+}
+
+#[test]
+fn the_lines_are_sent_in_batches() {
+    let (host, asked) = server(vec![(200, vectors(2, 2)), (200, vectors(1, 2))]);
+    let out = run(
+        &["--host", &host, "--model", "m", "--batch", "2"],
+        "one\ntwo\nthree\n",
+    );
+
+    assert_eq!(out.status.code(), Some(0));
+    assert_eq!(out.stdout.lines().count(), 3);
+    assert_eq!(asked.recv().expect("a request").inputs.len(), 2);
+    assert_eq!(asked.recv().expect("a request").inputs.len(), 1);
+}
+
+#[test]
+fn a_vector_of_the_wrong_width_is_refused() {
+    let (host, _asked) = server(vec![(200, vectors(1, 768))]);
+    let out = run(&["--host", &host, "--model", "m", "--dim", "1024"], "one\n");
+
+    assert_eq!(out.status.code(), Some(65), "bad data must exit 65");
+    assert!(out.stdout.is_empty(), "a bad answer must write nothing");
+    let said = String::from_utf8_lossy(&out.stderr);
+    assert!(said.contains("768"), "{said}");
+}
+
+#[test]
+fn a_server_that_is_not_there_is_reported_as_unavailable() {
+    // Port 1 on the loopback: nothing is listening and nothing will be.
+    let out = run(
+        &[
+            "--host",
+            "http://127.0.0.1:1",
+            "--model",
+            "m",
+            "--attempts",
+            "1",
+        ],
+        "one\n",
+    );
+    assert_eq!(out.status.code(), Some(69), "unreachable must exit 69");
+}
+
+#[test]
+fn a_failed_request_is_tried_again() {
+    let (host, asked) = server(vec![(503, "{}".into()), (200, vectors(1, 2))]);
+    let out = run(
+        &["--host", &host, "--model", "m", "--attempts", "2"],
+        "one\n",
+    );
+
+    assert_eq!(
+        out.status.code(),
+        Some(0),
+        "{:?}",
+        String::from_utf8_lossy(&out.stderr)
+    );
+    assert_eq!(out.stdout.lines().count(), 1);
+    assert_eq!(asked.recv().expect("a first request").inputs, ["one"]);
+    assert_eq!(asked.recv().expect("a second request").inputs, ["one"]);
+}
+
+#[test]
+fn a_long_line_is_cut_before_it_is_sent() {
+    let (host, asked) = server(vec![(200, vectors(1, 2))]);
+    let long = "word ".repeat(100);
+    run(
+        &["--host", &host, "--model", "m", "--max-input-tokens", "5"],
+        &format!("{long}\n"),
+    );
+
+    let asked = asked.recv().expect("a request");
+    assert_eq!(textkit::count(&asked.inputs[0]), 5);
+}
+
+#[test]
+fn no_input_asks_the_server_nothing() {
+    let out = run(&["--host", "http://127.0.0.1:1", "--model", "m"], "");
+    assert_eq!(out.status.code(), Some(0));
+    assert!(out.stdout.is_empty());
+}