gqls/semantic/
mod.rs

1//! Semantic (embedding) search over schema records.
2//!
3//! Embeds `path + description + type` per record and the query with a local
4//! `all-MiniLM-L6-v2` model (pipeline borrowed from ae), compresses to 64-d
5//! Matryoshka vectors, and ranks by cosine. Falls back to a deterministic hash
6//! embedder when the model can't be fetched, so it always runs.
7//!
8//! v0 embeds every record per invocation (fine at schema scale). A persistent
9//! embedding cache keyed by schema hash — like ae's SQLite store — is the
10//! follow-up for very large / federated schemas; the embed and rank steps are
11//! kept separable here so that cache drops in cleanly.
12
13mod cache;
14mod embed;
15mod mrl;
16
17use crate::model::{Kind, SchemaRecord};
18use embed::default_embedder;
19use mrl::{compress_matryoshka_vector, cosine_similarity};
20
21/// Embed the query and each record, rank by cosine. Returns `(score, record)`
22/// pairs, best first — the caller formats them exactly like fuzzy hits, so
23/// `--json` / `--ndjson` work identically in both modes.
24pub fn search<'a>(
25    query: &str,
26    records: &'a [SchemaRecord],
27    kind: Option<Kind>,
28    limit: usize,
29    model: Option<&str>,
30    refresh: bool,
31) -> Vec<(f64, &'a SchemaRecord)> {
32    let embedder = default_embedder(model);
33    eprintln!("gqls: semantic search via {} embeddings", embedder.kind());
34
35    // Per-record vectors are the expensive part; cache them keyed by schema
36    // content + embedder + model, so editing the schema (or switching model)
37    // changes the key and re-embeds automatically. `--refresh` forces a miss.
38    let cache_path = cache::path(records, embedder.kind(), model);
39    let cached = if refresh {
40        None
41    } else {
42        cache_path
43            .as_deref()
44            .and_then(|p| cache::load(p, records.len()))
45    };
46    let vectors = match cached {
47        Some(v) => {
48            if let Some(p) = cache_path.as_deref() {
49                cache::touch(p); // LRU: mark this schema as recently used
50            }
51            v
52        }
53        None => {
54            use rayon::prelude::*;
55            use std::io::IsTerminal;
56            use std::sync::atomic::{AtomicUsize, Ordering};
57
58            let total = records.len();
59            eprintln!(
60                "gqls: embedding {total} records (one-time, then cached; a large schema can take ~a minute)…"
61            );
62            let done = AtomicUsize::new(0);
63
64            // Per-worker embedder/session: the session is behind a Mutex, so a
65            // shared one would serialize the parallel pass. The model is a cached
66            // file, so every worker resolves the same kind as the main embedder
67            // above (no mixed onnx/hash vectors). Order-preserving `collect`, so
68            // vectors stay index-aligned with `records`.
69            let v: Vec<Vec<f32>> = std::thread::scope(|scope| {
70                let show_progress = std::io::stderr().is_terminal() && total > 500;
71                if show_progress {
72                    scope.spawn(|| loop {
73                        std::thread::sleep(std::time::Duration::from_millis(300));
74                        let d = done.load(Ordering::Relaxed);
75                        eprint!("\rgqls: embedded {d}/{total}…    ");
76                        if d >= total {
77                            eprintln!();
78                            break;
79                        }
80                    });
81                }
82                records
83                    .par_iter()
84                    .map_init(
85                        || default_embedder(model),
86                        |emb, r| {
87                            let out = compress_matryoshka_vector(&emb.embed(&record_text(r)));
88                            done.fetch_add(1, Ordering::Relaxed);
89                            out
90                        },
91                    )
92                    .collect()
93            });
94            if let Some(p) = cache_path.as_deref() {
95                cache::store(p, &v);
96                cache::prune(cache::max_files()); // evict least-recently-used
97            }
98            v
99        }
100    };
101
102    let query_vec = compress_matryoshka_vector(&embedder.embed(query));
103    let mut hits: Vec<(f64, &SchemaRecord)> = records
104        .iter()
105        .zip(&vectors)
106        .filter(|(r, _)| kind.is_none_or(|k| r.kind == k))
107        .map(|(r, v)| (cosine_similarity(&query_vec, v) as f64, r))
108        .collect();
109
110    hits.sort_by(|a, b| b.0.total_cmp(&a.0));
111    hits.truncate(limit);
112    hits
113}
114
115/// The natural-language signal a semantic query matches against: the path plus
116/// the human description and the type.
117fn record_text(r: &SchemaRecord) -> String {
118    let mut s = r.path.clone();
119    if let Some(d) = &r.description {
120        s.push_str(" — ");
121        s.push_str(d);
122    }
123    if let Some(t) = &r.type_ref {
124        s.push_str(" : ");
125        s.push_str(t);
126    }
127    s
128}
129
130/// Delete all cached embedding vector files; returns how many were removed.
131pub fn clear_cache() -> usize {
132    cache::clear()
133}
134
135/// Whether the schema's vectors are already cached (either embedder kind), so a
136/// semantic search would be warm. Checked without loading a model.
137pub fn is_cached(records: &[SchemaRecord], model: Option<&str>) -> bool {
138    cache::exists(records, "onnx", model) || cache::exists(records, "hash", model)
139}
140
141/// Embed + cache every record's vector without running a real query — pre-warms
142/// the cache so the first search is instant. Returns the record count.
143pub fn warm(records: &[SchemaRecord], model: Option<&str>, refresh: bool) -> usize {
144    let _ = search("", records, None, 0, model, refresh);
145    records.len()
146}
gqls/semantic/mod.rs

gqls/semantic/
mod.rs
