index.rs
Rust · 511 lines · 21.2 KB
//! Derived search index (index.db): per-section embeddings plus SQLite FTS5.
//! Rebuilt from entries/ whenever files change, so it is always safe to delete.
use crate::ai;
use crate::store::{sections, split_frontmatter, Brain};
use anyhow::{bail, Result};
use rusqlite::{params, Connection};
use serde::Serialize;
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs;
/// The built-in CPU model (with the `local-embeddings` feature).
const LOCAL_MODEL_ID: &str = "bge-base-en-v1.5";
const LOCAL_QUERY_PREFIX: &str = "Represent this sentence for searching relevant passages: ";
const CHUNK_CHARS: usize = 1400;
enum Embedder {
Unloaded,
#[cfg(feature = "local-embeddings")]
Local(Box<fastembed::TextEmbedding>),
Remote(Box<(ai::Client, String)>),
}
pub struct Index {
pub conn: Connection,
embedder: Embedder,
provider: String,
cfg: crate::store::Embeddings,
privacy: crate::store::Privacy,
timeout_secs: u64,
/// Whether the last search really compared meanings (false when embeddings were unavailable).
pub used_semantic: bool,
}
#[derive(Serialize, Clone, Debug)]
pub struct Hit {
pub id: String,
pub title: String,
pub kind: String,
pub project: String,
pub tags: String,
pub engine: String,
pub verified: bool,
pub updated: String,
pub section: String,
pub snippet: String,
pub similarity: f32,
}
#[derive(Default, Clone)]
pub struct Filter {
pub project: String,
pub kind: String,
pub tags: Vec<String>,
pub engine: String,
pub genre: String,
pub platform: String,
/// Entries to leave out entirely (retired ones), dropped before ranking so they can't use up
/// the keyword window.
pub exclude: std::collections::HashSet<String>,
}
fn to_blob(v: &[f32]) -> Vec<u8> {
v.iter().flat_map(|x| x.to_le_bytes()).collect()
}
fn from_blob(b: &[u8]) -> Vec<f32> {
b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect()
}
fn dot(a: &[f32], b: &[f32]) -> f32 {
a.iter().zip(b).map(|(x, y)| x * y).sum()
}
fn normalize(mut v: Vec<f32>) -> Vec<f32> {
let n = dot(&v, &v).sqrt().max(1e-9);
v.iter_mut().for_each(|x| *x /= n);
v
}
/// Split a section into pieces of about CHUNK_CHARS on paragraph boundaries.
fn pieces(text: &str) -> Vec<String> {
let mut out = vec![];
let mut cur = String::new();
for para in text.split("\n\n") {
if !cur.is_empty() && cur.len() + para.len() > CHUNK_CHARS {
out.push(std::mem::take(&mut cur));
}
if !cur.is_empty() {
cur.push_str("\n\n");
}
cur.push_str(para);
}
if !cur.trim().is_empty() || out.is_empty() {
out.push(cur);
}
out
}
impl Index {
pub fn open(brain: &Brain) -> Result<Index> {
let conn = Connection::open(brain.home.join("index.db"))?;
conn.busy_timeout(std::time::Duration::from_secs(30))?;
conn.execute_batch(
"PRAGMA journal_mode=WAL;
CREATE TABLE IF NOT EXISTS meta (key TEXT PRIMARY KEY, value TEXT);
CREATE TABLE IF NOT EXISTS entries (
id TEXT PRIMARY KEY, sha TEXT, title TEXT, kind TEXT, project TEXT, owner TEXT,
tags TEXT, engine TEXT, genre TEXT, platform TEXT, verified INTEGER,
updated TEXT, vec BLOB);
CREATE TABLE IF NOT EXISTS chunks (
entry_id TEXT, section TEXT, text TEXT, vec BLOB);
CREATE INDEX IF NOT EXISTS chunks_entry ON chunks(entry_id);
CREATE VIRTUAL TABLE IF NOT EXISTS fts USING fts5(
entry_id UNINDEXED, section UNINDEXED, title, body);",
)?;
let cfg = brain.config.embeddings.clone();
let provider = cfg.provider.trim().to_lowercase();
// Vectors from a different model aren't comparable, so the index remembers which made them.
let model_id = match provider.as_str() {
"local" => LOCAL_MODEL_ID.to_string(),
"none" | "" => "none".to_string(),
p => format!("{p}:{}", cfg.model),
};
let model: Option<String> = conn
.query_row("SELECT value FROM meta WHERE key='model'", [], |r| r.get(0))
.ok();
if model.as_deref() != Some(model_id.as_str()) {
conn.execute_batch("DELETE FROM entries; DELETE FROM chunks; DELETE FROM fts;")?;
conn.execute("INSERT OR REPLACE INTO meta VALUES ('model', ?1)", [&model_id])?;
}
Ok(Index {
conn,
embedder: Embedder::Unloaded,
provider,
cfg,
privacy: brain.config.privacy.clone(),
timeout_secs: 30,
used_semantic: false,
})
}
/// Whether search can match by meaning (otherwise it is keyword-only).
pub fn semantic(&self) -> bool {
!matches!(self.provider.as_str(), "none" | "")
}
fn load_embedder(&mut self) -> Result<()> {
if !matches!(self.embedder, Embedder::Unloaded) {
return Ok(());
}
self.embedder = match self.provider.as_str() {
#[cfg(feature = "local-embeddings")]
"local" => {
let dir = crate::store::app_cache_dir().join("models");
let opts = fastembed::TextInitOptions::new(fastembed::EmbeddingModel::BGEBaseENV15)
.with_cache_dir(dir)
.with_show_download_progress(false);
Embedder::Local(Box::new(
fastembed::TextEmbedding::try_new(opts)
.map_err(|e| anyhow::anyhow!("loading the local embedding model: {e}"))?,
))
}
#[cfg(not(feature = "local-embeddings"))]
"local" => bail!(
"embeddings.provider is `local`, but this build has no local model \
(rebuild with `--features local-embeddings`, or run `brain setup` to pick azure/openai/none)"
),
"azure" | "openai" => {
if self.cfg.url.is_empty() || self.cfg.model.is_empty() {
bail!("embeddings.url and embeddings.model must be set for provider `{}` (see `brain setup`)", self.provider);
}
let kind = if self.provider == "azure" { ai::Kind::Azure } else { ai::Kind::OpenAi };
let client = ai::Client::new(kind, &self.cfg.url, ai::Auth::parse(&self.cfg.auth), self.timeout_secs)
.with_token_cache(crate::store::auth_cache_path())
.with_token_hosts(&self.cfg.token_hosts)
.with_privacy(&self.privacy);
Embedder::Remote(Box::new((client, self.cfg.model.clone())))
}
other => bail!("unknown embeddings.provider `{other}` (use local, azure, openai or none)"),
};
Ok(())
}
/// Embeddings for texts; empty vectors when search is keyword-only.
pub fn embed(&mut self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() || !self.semantic() {
return Ok(vec![vec![]; texts.len()]);
}
self.load_embedder()?;
let raw = match &mut self.embedder {
#[cfg(feature = "local-embeddings")]
Embedder::Local(m) => m.embed(texts, Some(16)).map_err(|e| anyhow::anyhow!("{e}"))?,
Embedder::Remote(r) => r.0.embeddings(&r.1, texts)?,
Embedder::Unloaded => unreachable!(),
};
Ok(raw.into_iter().map(normalize).collect())
}
pub fn embed_query(&mut self, q: &str) -> Result<Vec<f32>> {
// Only the built-in bge model wants an instruction prefix on queries.
let text = if self.provider == "local" { format!("{LOCAL_QUERY_PREFIX}{q}") } else { q.to_string() };
Ok(self.embed(&[text])?.remove(0))
}
/// Whole-entry text used for duplicate detection.
pub fn doc_text(title: &str, tags: &str, body: &str) -> String {
let b: String = body.chars().take(2000).collect();
format!("{title}\n{tags}\n{b}")
}
/// Bring the index in line with entries/. Returns how many entries were (re)indexed.
pub fn sync(&mut self, brain: &Brain) -> Result<usize> {
// With embeddings on, entries indexed without a vector (embeddings were unreachable then)
// count as changed, so they get embedded once the provider works again.
let sql = if self.semantic() {
"SELECT id, CASE WHEN length(vec) = 0 THEN '' ELSE sha END FROM entries"
} else {
"SELECT id, sha FROM entries"
};
let known: HashMap<String, String> = self
.conn
.prepare(sql)?
.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
.collect::<rusqlite::Result<_>>()?;
let ids = brain.all_ids()?;
let mut changed = vec![];
for id in &ids {
let raw = fs::read_to_string(brain.entry_path(id)?)?;
let sha = hex::encode(Sha256::digest(raw.as_bytes()));
if known.get(id) != Some(&sha) {
match split_frontmatter(&raw) {
Ok((meta, body)) => changed.push((id.clone(), sha, meta, body)),
Err(e) => eprintln!("warning: skipping entry {id}: {e:#}"),
}
}
}
let gone: Vec<&String> = known.keys().filter(|k| !ids.contains(k)).collect();
if changed.is_empty() && gone.is_empty() {
return Ok(0);
}
// An entry that looks like it holds a secret is never sent to a remote embedding service
// (hand edits and git pulls skip `brain record`'s check). It stays keyword-searchable, and
// the warning repeats until someone removes the secret.
let remote = matches!(self.provider.as_str(), "azure" | "openai");
let withheld: Vec<bool> = changed
.iter()
.map(|(id, _, meta, body)| {
// Everything that could be sent: the whole frontmatter (title, tags, …) and the body.
let hit = crate::store::secret_kind(&format!("{}\n{}", serde_yaml::to_string(meta).unwrap_or_default(), body));
if let (true, Some(what)) = (remote, hit) {
eprintln!("warning: entry {id} looks like it contains {what}; not sent for embedding. Remove it (`brain edit {id}`).");
}
remote && hit.is_some()
})
.collect();
// Embed everything that changed in one batch.
let mut texts = vec![];
let mut plan = vec![]; // (entry index, section, text)
let mut skip = vec![]; // per text: withheld from the embedding service
for (i, (_, _, meta, body)) in changed.iter().enumerate() {
texts.push(Self::doc_text(&meta.title, &meta.tags.join(", "), body));
plan.push((i, String::from("<doc>"), String::new()));
for (h, t) in sections(body) {
for p in pieces(&t) {
texts.push(format!("{}\n{}\n{}", meta.title, h, p));
plan.push((i, h.clone(), p));
}
}
skip.resize(texts.len(), withheld[i]);
}
let send: Vec<String> = texts.iter().zip(&skip).filter(|(_, s)| !**s).map(|(t, _)| t.clone()).collect();
let vecs = match self.embed(&send) {
Ok(v) => {
let mut it = v.into_iter();
skip.iter().map(|s| if *s { vec![] } else { it.next().unwrap_or_default() }).collect()
}
Err(e) => {
// Keep the brain usable (keyword search) when the provider is down; retried next sync.
eprintln!("warning: embeddings unavailable ({e:#}); indexing for keyword search until they work");
vec![vec![]; texts.len()]
}
};
let tx = self.conn.transaction()?;
for (id, sha, meta, _) in &changed {
tx.execute("DELETE FROM chunks WHERE entry_id=?1", [id])?;
tx.execute("DELETE FROM fts WHERE entry_id=?1", [id])?;
tx.execute("DELETE FROM entries WHERE id=?1", [id])?;
let _ = sha;
let _ = meta;
}
for ((i, section, text), v) in plan.iter().zip(&vecs) {
let (id, sha, meta, _) = &changed[*i];
if section == "<doc>" {
tx.execute(
"INSERT INTO entries VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13)",
params![
id,
sha,
meta.title,
meta.kind,
meta.project,
if meta.owner.is_empty() { &meta.project } else { &meta.owner },
meta.tags.join(", "),
meta.engine.join(", "),
meta.genre.join(", "),
meta.platform.join(", "),
meta.verified as i32,
if meta.updated.is_empty() { &meta.created } else { &meta.updated },
to_blob(v),
],
)?;
} else {
tx.execute(
"INSERT INTO chunks VALUES (?1,?2,?3,?4)",
params![id, section, text, to_blob(v)],
)?;
tx.execute(
"INSERT INTO fts VALUES (?1,?2,?3,?4)",
params![id, section, meta.title, format!("{}\n{}\n{}", meta.tags.join(" "), section, text)],
)?;
}
}
for id in &gone {
tx.execute("DELETE FROM chunks WHERE entry_id=?1", [id])?;
tx.execute("DELETE FROM fts WHERE entry_id=?1", [id])?;
tx.execute("DELETE FROM entries WHERE id=?1", [id])?;
}
tx.commit()?;
Ok(changed.len())
}
/// Entries passing the filter: id → (title, kind, project, tags, engine, verified, updated).
fn candidates(&self, f: &Filter) -> Result<HashMap<String, Hit>> {
let mut stmt = self.conn.prepare(
"SELECT id, title, kind, project, owner, tags, engine, genre, platform, verified, updated FROM entries",
)?;
let rows = stmt.query_map([], |r| {
Ok((
Hit {
id: r.get(0)?,
title: r.get(1)?,
kind: r.get(2)?,
project: r.get(3)?,
tags: r.get(5)?,
engine: r.get(6)?,
verified: r.get::<_, i32>(9)? != 0,
updated: r.get(10)?,
section: String::new(),
snippet: String::new(),
similarity: 0.0,
},
r.get::<_, String>(4)?,
r.get::<_, String>(7)?,
r.get::<_, String>(8)?,
))
})?;
let has = |hay: &str, needle: &str| needle.is_empty() || hay.to_lowercase().contains(&needle.to_lowercase());
let mut out = HashMap::new();
for row in rows {
let (h, owner, genre, platform) = row?;
let tags: Vec<&str> = h.tags.split(", ").collect();
let ok = (has(&h.project, &f.project) || has(&owner, &f.project))
&& (f.kind.is_empty() || h.kind == f.kind)
&& f.tags.iter().all(|t| tags.contains(&t.as_str()))
&& has(&h.engine, &f.engine)
&& has(&genre, &f.genre)
&& has(&platform, &f.platform);
if ok && !f.exclude.contains(&h.id) {
out.insert(h.id.clone(), h);
}
}
Ok(out)
}
pub fn list(&self, f: &Filter) -> Result<Vec<Hit>> {
let mut v: Vec<Hit> = self.candidates(f)?.into_values().collect();
v.sort_by(|a, b| b.updated.cmp(&a.updated).then(a.id.cmp(&b.id)));
Ok(v)
}
/// Hybrid search: best-section vector similarity fused with BM25 keyword rank.
pub fn search(&mut self, queries: &[String], f: &Filter, limit: usize) -> Result<Vec<Hit>> {
let mut cands = self.candidates(f)?;
if cands.is_empty() {
return Ok(vec![]);
}
let chunks: Vec<(String, String, String, Vec<f32>)> = self
.conn
.prepare("SELECT entry_id, section, text, vec FROM chunks")?
.query_map([], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, from_blob(&r.get::<_, Vec<u8>>(3)?))))?
.filter_map(|r| r.ok())
.filter(|c| cands.contains_key(&c.0))
.collect();
let semantic = self.semantic();
self.used_semantic = false;
let mut score: HashMap<String, f32> = HashMap::new();
for (qi, q) in queries.iter().enumerate() {
// The original query counts double; --deep rephrasings only add recall.
let weight = if qi == 0 { 2.0 } else { 1.0 };
let qv = if semantic {
match self.embed_query(q) {
Ok(v) => v,
Err(e) => {
eprintln!("warning: embeddings unavailable ({e:#}); keyword search only");
vec![]
}
}
} else {
vec![]
};
self.used_semantic |= !qv.is_empty();
// Best chunk per entry.
let mut best: HashMap<&str, (f32, usize)> = HashMap::new();
for (i, c) in chunks.iter().enumerate().filter(|(_, c)| !qv.is_empty() && !c.3.is_empty()) {
let s = dot(&qv, &c.3);
let e = best.entry(c.0.as_str()).or_insert((f32::MIN, i));
if s > e.0 {
*e = (s, i);
}
}
let mut ranked: Vec<(&str, f32, usize)> = best.iter().map(|(k, v)| (*k, v.0, v.1)).collect();
ranked.sort_by(|a, b| b.1.total_cmp(&a.1));
for (rank, (id, sim, ci)) in ranked.iter().enumerate() {
*score.entry(id.to_string()).or_default() += weight / (60.0 + rank as f32);
let h = cands.get_mut(*id).unwrap();
if *sim > h.similarity {
h.similarity = *sim;
h.section = chunks[*ci].1.clone();
h.snippet = chunks[*ci].2.clone();
}
}
// Keyword rank.
let words: Vec<String> = q
.split(|c: char| !c.is_alphanumeric() && c != '_')
.filter(|w| w.len() > 1)
.map(|w| format!("\"{w}\""))
.collect();
if !words.is_empty() {
let mut stmt = self
.conn
.prepare("SELECT entry_id, section, body FROM fts WHERE fts MATCH ?1 ORDER BY bm25(fts, 0, 0, 3.0, 1.0)")?;
let hits: Vec<(String, String, String)> = stmt
.query_map([words.join(" OR ")], |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)))?
.filter_map(|r| r.ok())
.collect();
let mut seen = std::collections::HashSet::new();
let mut rank = 0;
for (id, section, body) in hits {
if cands.contains_key(&id) && seen.insert(id.clone()) {
*score.entry(id.clone()).or_default() += weight / (60.0 + rank as f32);
rank += 1;
// Keyword-only search: show the best keyword section as the match.
let h = cands.get_mut(&id).unwrap();
if qv.is_empty() && h.section.is_empty() {
h.section = section;
h.snippet = body.splitn(3, '\n').nth(2).unwrap_or(&body).to_string();
}
}
}
}
}
// Light trust tie-break: a verified entry edges ahead of an unverified one that scores the
// same (worth about one rank in the fusion), never past a clearly better match.
let mut ids: Vec<(String, f32)> = score
.into_iter()
.map(|(id, sc)| {
let bonus = if cands.get(&id).is_some_and(|h| h.verified) { 0.0005 } else { 0.0 };
(id, sc + bonus)
})
.collect();
ids.sort_by(|a, b| b.1.total_cmp(&a.1));
Ok(ids
.into_iter()
.take(limit)
.filter_map(|(id, _)| cands.remove(&id))
.collect())
}
pub fn doc_vec(&self, id: &str) -> Option<Vec<f32>> {
self.conn
.query_row("SELECT vec FROM entries WHERE id=?1", [id], |r| r.get::<_, Vec<u8>>(0))
.ok()
.map(|b| from_blob(&b))
}
/// Entries most similar to a vector (whole-entry embeddings).
pub fn nearest(&self, v: &[f32], exclude: &str, limit: usize) -> Result<Vec<(String, String, String, f32)>> {
let mut out: Vec<(String, String, String, f32)> = self
.conn
.prepare("SELECT id, title, project, vec FROM entries")?
.query_map([], |r| {
Ok((r.get(0)?, r.get(1)?, r.get(2)?, from_blob(&r.get::<_, Vec<u8>>(3)?)))
})?
.filter_map(|r| r.ok())
.filter(|r: &(String, String, String, Vec<f32>)| r.0 != exclude && !r.3.is_empty() && !v.is_empty())
.map(|(id, t, p, ev)| {
let s = dot(v, &ev);
(id, t, p, s)
})
.collect();
out.sort_by(|a, b| b.3.total_cmp(&a.3));
out.truncate(limit);
Ok(out)
}
}
