Vector search

use vectors.qweruse vectors_add, vectors_search from vectors      // Quillo 0.6: only these

// "More like this" for a shop: products as small embeddings (from any embedding model, here// four made-up numbers: sporty, cosy, techy, outdoorsy), searched by similarity.////   quillo run vectors.qwe      then open http://localhost:8080/like?id=tentuse "vectors.qwer"let products = [    { id: "tent", name: "Two-person tent", v: [0.2, 0.1, 0.1, 0.95] },    { id: "boots", name: "Hiking boots", v: [0.6, 0.1, 0.0, 0.8] },    { id: "blanket", name: "Wool blanket", v: [0.0, 0.95, 0.0, 0.3] },    { id: "watch", name: "GPS watch", v: [0.7, 0.0, 0.8, 0.5] },    { id: "laptop", name: "Laptop", v: [0.0, 0.2, 0.98, 0.0] },]for p in products { vectors_add("products", p.id, p.v, { name: p.name })? }route GET /like ? { id: str } {    let p = find(products, (p) => p.id == id) ?? reply 404 "no such product"    vectors_search("products", p.v, 4, 0.5)?        | filter((m) => m.id != id)        | map((m) => { name: m.metadata.name, score: round(m.score, 3) })}route GET /search ? { sporty: float = 0.0, cosy: float = 0.0, techy: float = 0.0, outdoorsy: float = 0.0 } =    vectors_search("products", [sporty, cosy, techy, outdoorsy], 2, -1.0)? | map((m) => m.metadata.name)route GET /save {    let n = vectors_save("products", "data/products.vectors.json")?    vectors_remove("products", "laptop")    { saved: n, after_remove: vectors_count("products"), loaded: vectors_load("copy", "data/products.vectors.json")? }}

//! Nearest neighbours among embeddings, kept in memory by collection name: "more like this", semantic search and recommendations, with cosine similarity, saved to a file and loaded back.//! version: 1.0.0//!//! [dependencies]use serde_json::{json, Value};use std::collections::HashMap;use std::sync::{Mutex, OnceLock};const MAX_DIMS: usize = 8192;const MAX_ITEMS: usize = 200_000;struct Item {    id: String,    v: Vec<f32>,    meta: Value,}#[derive(Default)]struct Coll {    dims: usize,    items: Vec<Item>,    at: HashMap<String, usize>,}fn all() -> std::sync::MutexGuard<'static, HashMap<String, Coll>> {    static C: OnceLock<Mutex<HashMap<String, Coll>>> = OnceLock::new();    C.get_or_init(|| Mutex::new(HashMap::new())).lock().unwrap_or_else(|e| e.into_inner())}/// The vector scaled to length 1, so cosine similarity is a dot product.fn unit(v: &[f64]) -> Result<Vec<f32>, String> {    if v.is_empty() || v.len() > MAX_DIMS {        return Err(format!("a vector has 1 to {} numbers, not {}", MAX_DIMS, v.len()));    }    let n = v.iter().map(|x| x * x).sum::<f64>().sqrt();    if n == 0.0 || !n.is_finite() {        return Err("a vector of zeros (or not finite) has no direction".into());    }    Ok(v.iter().map(|x| (x / n) as f32).collect())}fn add(c: &mut Coll, id: &str, v: Vec<f32>, meta: Value) -> Result<(), String> {    if c.items.is_empty() {        c.dims = v.len();    } else if v.len() != c.dims {        return Err(format!("this collection's vectors have {} numbers, not {}", c.dims, v.len()));    }    match c.at.get(id) {        Some(&i) => c.items[i] = Item { id: id.to_string(), v, meta },        None => {            if c.items.len() >= MAX_ITEMS {                return Err(format!("a collection holds at most {} vectors", MAX_ITEMS));            }            c.at.insert(id.to_string(), c.items.len());            c.items.push(Item { id: id.to_string(), v, meta });        }    }    Ok(())}/// Adds a vector (an embedding, any length up to 8,192, the same for the whole collection) under `id`, with `metadata` given back by searches. The same `id` again replaces it.pub fn vectors_add(collection: &str, id: &str, vector: Vec<f64>, metadata: serde_json::Value) -> Result<i64, String> {    let v = unit(&vector)?;    let mut all = all();    let c = all.entry(collection.to_string()).or_default();    add(c, id, v, metadata)?;    Ok(c.items.len() as i64)}/// The `k` vectors most like `vector` (1 to 1,000), best first: `[{ id, score, metadata }]`, the score being cosine similarity from -1 to 1 (1: same direction)./// `min_score` drops weaker matches (-1 keeps all). An unknown collection gives an empty list.pub fn vectors_search(collection: &str, vector: Vec<f64>, k: i64, min_score: f64) -> Result<Vec<Value>, String> {    let q = unit(&vector)?;    let all = all();    let Some(c) = all.get(collection) else { return Ok(vec![]) };    if !c.items.is_empty() && q.len() != c.dims {        return Err(format!("this collection's vectors have {} numbers, not {}", c.dims, q.len()));    }    let mut scored: Vec<(f32, usize)> = c.items.iter().enumerate().map(|(i, it)| (it.v.iter().zip(&q).map(|(a, b)| a * b).sum::<f32>(), i)).filter(|(s, _)| *s as f64 >= min_score).collect();    scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));    Ok(scored.into_iter().take(k.clamp(1, 1000) as usize).map(|(s, i)| json!({ "id": c.items[i].id, "score": (s as f64 * 1e6).round() / 1e6, "metadata": c.items[i].meta })).collect())}/// Removes the vector under `id`, and says whether it was there.pub fn vectors_remove(collection: &str, id: &str) -> bool {    let mut all = all();    let Some(c) = all.get_mut(collection) else { return false };    let Some(i) = c.at.remove(id) else { return false };    c.items.swap_remove(i);    if i < c.items.len() {        let moved = c.items[i].id.clone();        c.at.insert(moved, i);    }    true}/// How many vectors a collection holds (0 for an unknown one).pub fn vectors_count(collection: &str) -> i64 {    all().get(collection).map(|c| c.items.len() as i64).unwrap_or(0)}/// Writes a collection to a JSON file (folders made as needed), to load after a restart. Gives the number written.pub fn vectors_save(collection: &str, path: &str) -> Result<i64, String> {    let all = all();    let items: Vec<Value> = all.get(collection).map(|c| c.items.iter().map(|it| json!({ "id": it.id, "v": it.v, "m": it.meta })).collect()).unwrap_or_default();    if let Some(dir) = std::path::Path::new(path).parent() {        if !dir.as_os_str().is_empty() {            std::fs::create_dir_all(dir).map_err(|e| e.to_string())?;        }    }    std::fs::write(path, Value::Array(items.clone()).to_string()).map_err(|e| format!("{}: {}", path, e))?;    Ok(items.len() as i64)}/// Loads a file written by `vectors_save` into a collection (replacing what it held), and gives the number loaded.pub fn vectors_load(collection: &str, path: &str) -> Result<i64, String> {    let text = std::fs::read_to_string(path).map_err(|e| format!("{}: {}", path, e))?;    let items: Vec<Value> = serde_json::from_str(&text).map_err(|e| e.to_string())?;    let mut c = Coll::default();    for it in items {        let v: Vec<f64> = serde_json::from_value(it["v"].clone()).map_err(|e| e.to_string())?;        add(&mut c, it["id"].as_str().unwrap_or(""), unit(&v)?, it["m"].clone())?;    }    let n = c.items.len() as i64;    all().insert(collection.to_string(), c);    Ok(n)}