index.rs
Rust · 821 lines · 38.0 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, Mind};
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;
/// At most this many texts are sent for embedding in one sync; the rest follow on later runs.
const MAX_TEXTS_PER_SYNC: usize = 2000;
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();
// A paragraph longer than a piece (a pasted log) is cut, so no single text is oversized.
let paras: Vec<String> = text
.split("\n\n")
.flat_map(|p| {
let chars: Vec<char> = p.chars().collect();
if chars.len() <= CHUNK_CHARS {
vec![p.to_string()]
} else {
chars.chunks(CHUNK_CHARS).map(|c| c.iter().collect()).collect()
}
})
.collect();
for para in ¶s {
let para = para.as_str();
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 {
/// The search index. It is only a cache of the entries: a corrupt one is deleted and rebuilt.
pub fn open(mind: &Mind) -> Result<Index> {
match Self::open_db(mind) {
Err(e) if is_corrupt(&e) => {
eprintln!("warning: the search index was damaged ({e:#}); rebuilding it from the entries");
for f in ["index.db", "index.db-wal", "index.db-shm"] {
let _ = fs::remove_file(mind.home.join(f));
}
Self::open_db(mind)
}
r => r,
}
}
fn open_db(mind: &Mind) -> Result<Index> {
let conn = Connection::open(mind.home.join("index.db"))?;
conn.busy_timeout(std::time::Duration::from_secs(30))?;
conn.execute_batch("PRAGMA journal_mode=WAL;")?;
Self::with_conn(conn, mind, mind.config.embeddings.clone())
}
/// The saved index, read-only, for when it can't be written (a sandbox such as Codex's, a
/// read-only folder): searches use it as it was last saved. None if there is no usable one.
pub fn open_readonly(mind: &Mind) -> Option<Index> {
use rusqlite::OpenFlags;
let path = mind.home.join("index.db");
if !path.is_file() {
return None;
}
let conn = Connection::open_with_flags(immutable_uri(&path), OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_URI | OpenFlags::SQLITE_OPEN_NO_MUTEX).ok()?;
let ix = Self::with_conn_readonly(conn, mind).ok()?;
// Only if it was made by the model configured now (vectors from another aren't comparable).
let made_by: Option<String> = ix.conn.query_row("SELECT value FROM meta WHERE key='model'", [], |r| r.get(0)).ok();
(made_by.as_deref() == Some(model_id(&mind.config.embeddings).as_str())).then_some(ix)
}
/// A temporary keyword-only index in memory, built from the entry files, for when there is no
/// saved index this process can use.
pub fn open_memory(mind: &Mind) -> Result<Index> {
let cfg = crate::store::Embeddings { provider: "none".into(), ..mind.config.embeddings.clone() };
Self::with_conn(Connection::open_in_memory()?, mind, cfg)
}
fn with_conn_readonly(conn: Connection, mind: &Mind) -> Result<Index> {
let cfg = mind.config.embeddings.clone();
Ok(Index {
conn,
embedder: Embedder::Unloaded,
provider: cfg.provider.trim().to_lowercase(),
cfg,
privacy: mind.config.privacy.clone(),
timeout_secs: 30,
used_semantic: false,
})
}
fn with_conn(conn: Connection, mind: &Mind, cfg: crate::store::Embeddings) -> Result<Index> {
conn.execute_batch(
"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 provider = cfg.provider.trim().to_lowercase();
let model_id = model_id(&cfg);
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: mind.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");
// The first use downloads the model: say so, or it looks like a hang.
if fs::read_dir(&dir).map_or(true, |mut d| d.next().is_none()) {
eprintln!("Downloading the local search model (about 420 MB, once) into {} …", dir.display());
}
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 `mind 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 `mind 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, mind: &Mind) -> 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 mut known: HashMap<String, String> = self
.conn
.prepare(sql)?
.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
.collect::<rusqlite::Result<_>>()?;
// Entries the provider refused for this exact content stay keyword-only until they change
// (else one oversized or rejected entry would be re-sent, with everything else, every run).
if self.semantic() {
let failed: Vec<(String, String)> = self
.conn
.prepare("SELECT substr(key, 14), value FROM meta WHERE key LIKE 'embed-failed:%'")?
.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?
.collect::<rusqlite::Result<_>>()?;
for (id, sha) in failed {
if known.get(&id).is_some_and(|k| k.is_empty()) {
known.insert(id, sha);
}
}
}
let ids = mind.all_ids()?;
let mut changed = vec![];
// Entries that can't be read or parsed now (a merge conflict, a bad edit) leave the index,
// so search doesn't keep showing their old text.
let mut broken: Vec<String> = vec![];
for id in &ids {
// An unreadable entry (not UTF-8, permissions) is skipped, like bad frontmatter: one
// file mustn't stop search, record and the hooks.
let raw = match fs::read_to_string(mind.entry_path(id)?) {
Ok(r) => r,
Err(e) => {
eprintln!("warning: skipping entry {id}: {e}");
broken.push(id.clone());
continue;
}
};
let sha = hex::encode(Sha256::digest(raw.as_bytes()));
if known.get(id) != Some(&sha) {
match split_frontmatter(&raw) {
// A conflict in the body (frontmatter intact) must not put both versions in search.
Ok((_, body)) if crate::store::has_conflict_markers(&body) => {
eprintln!("warning: skipping entry {id}: it has git conflict markers (resolve the merge)");
broken.push(id.clone());
}
Ok((meta, body)) => changed.push((id.clone(), sha, meta, body)),
Err(e) => {
eprintln!("warning: skipping entry {id}: {e:#}");
broken.push(id.clone());
}
}
}
}
let gone: Vec<&String> = known.keys().filter(|k| !ids.contains(k) || broken.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 `mind 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 (`mind 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]);
}
// Embedded in groups of entries, at most MAX_TEXTS_PER_SYNC texts per run (the rest are
// keyword-only until the next run). A group that fails is retried entry by entry: if a
// one-word probe works, the provider is up and that entry is refused (remembered, not
// retried until it changes); if the probe fails, the provider is down (all retried later).
let mut vecs: Vec<Vec<f32>> = vec![vec![]; texts.len()];
let mut failed_entries: Vec<usize> = vec![];
let by_entry: Vec<Vec<usize>> = (0..changed.len())
.map(|i| (0..texts.len()).filter(|&t| plan[t].0 == i && !skip[t]).collect())
.collect();
let mut sent = 0usize;
let mut down = None;
'groups: for group in by_entry.chunks(16) {
if down.is_some() || sent >= MAX_TEXTS_PER_SYNC {
break;
}
let idx: Vec<usize> = group.iter().flatten().copied().collect();
let batch: Vec<String> = idx.iter().map(|&t| texts[t].clone()).collect();
sent += batch.len();
match self.embed(&batch) {
Ok(v) => idx.iter().zip(v).for_each(|(&t, v)| vecs[t] = v),
Err(e) if !is_refusal(&e) => {
down = Some(e);
break 'groups;
}
Err(_) => {
for one in group {
let batch: Vec<String> = one.iter().map(|&t| texts[t].clone()).collect();
match self.embed(&batch) {
Ok(v) => one.iter().zip(v).for_each(|(&t, v)| vecs[t] = v),
// Only a refusal of the request itself (400/413/422) is about this
// entry; a rate limit, server error or timeout means try again later.
Err(e) if is_refusal(&e) => {
if let Some(&t) = one.first() {
failed_entries.push(plan[t].0);
}
}
Err(e) => {
down = Some(e);
break 'groups;
}
}
}
}
}
}
if let Some(e) = &down {
// Keep the mind usable (keyword search) when the provider is down; retried next sync.
eprintln!("warning: embeddings unavailable ({e:#}); indexing for keyword search until they work");
}
let tx = self.conn.transaction()?;
for (i, (id, sha, _, _)) in changed.iter().enumerate() {
let key = format!("embed-failed:{id}");
if failed_entries.contains(&i) {
eprintln!("warning: the embedding service refused entry {id}; it is keyword-only until it changes");
tx.execute("INSERT OR REPLACE INTO meta (key, value) VALUES (?1, ?2)", params![key, sha])?;
} else {
tx.execute("DELETE FROM meta WHERE key=?1", [key])?;
}
}
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())
}
/// Entries that look relevant to a chat message, by keywords alone (fast: no model, no sync),
/// for the prompt hook to point an agent at before it starts. Strict, so it stays quiet unless
/// an entry shares at least two meaningful words with the message (a third of them for longer
/// messages). Best first: most words shared, then keyword rank.
pub fn suggest(&self, message: &str, limit: usize) -> Result<Vec<(String, String)>> {
let words = meaningful_words(message);
if words.len() < 2 {
return Ok(vec![]);
}
let need = 2.max(words.len().div_ceil(3));
// A word in more than 1 in 20 entries ("build", "test" in a mind of coding notes) says little
// about which entry fits: an entry is only suggested if it shares at least one rarer word.
let total: usize = self.conn.query_row("SELECT COUNT(*) FROM entries", [], |r| r.get::<_, i64>(0)).unwrap_or(0).max(0) as usize;
let rare_below = 3.max(total / 20);
let mut shared: HashMap<String, usize> = HashMap::new();
let mut in_title: HashMap<String, usize> = HashMap::new();
let mut has_rare: std::collections::HashSet<String> = Default::default();
let ids_for = |q: String| -> Result<Vec<String>> {
// Bounded, so a huge (or planted) index can't make every message slow.
let mut stmt = self.conn.prepare("SELECT DISTINCT entry_id FROM fts WHERE fts MATCH ?1 LIMIT 500")?;
let v = stmt.query_map([q], |r| r.get(0))?.filter_map(|r| r.ok()).collect();
Ok(v)
};
for w in &words {
let ids = ids_for(format!("\"{w}\""))?;
let rare = ids.len() <= rare_below;
for id in ids {
if rare {
has_rare.insert(id.clone());
}
*shared.entry(id).or_default() += 1;
}
for id in ids_for(format!("title : \"{w}\""))? {
*in_title.entry(id).or_default() += 1;
}
}
// Distinctive enough: a rare shared word, or two shared words in the entry's title
// ("deploy studios" and an entry titled "Deploy studios.coinopolis.io").
shared.retain(|id, _| has_rare.contains(id) || in_title.get(id).copied().unwrap_or(0) >= 2);
let mut rank: HashMap<String, usize> = HashMap::new();
let q = words.iter().map(|w| format!("\"{w}\"")).collect::<Vec<_>>().join(" OR ");
let mut stmt = self.conn.prepare("SELECT entry_id FROM fts WHERE fts MATCH ?1 ORDER BY bm25(fts, 0, 0, 3.0, 1.0) LIMIT 500")?;
for (i, id) in stmt.query_map([q], |r| r.get::<_, String>(0))?.filter_map(|r| r.ok()).enumerate() {
rank.entry(id).or_insert(i);
}
// Only entries about as close as the best one (within a word of it): a weaker one riding on a
// single common word ("hangs") is noise next to a strong match.
let best = shared.values().copied().max().unwrap_or(0);
let mut ids: Vec<(String, usize)> = shared.into_iter().filter(|(_, n)| *n >= need.max(best.saturating_sub(1))).collect();
ids.sort_by(|a, b| b.1.cmp(&a.1).then(rank.get(&a.0).cmp(&rank.get(&b.0))));
let mut out = vec![];
for (id, _) in ids.into_iter().take(limit) {
if let Ok(t) = self.conn.query_row("SELECT title FROM entries WHERE id=?1", [&id], |r| r.get::<_, String>(0)) {
out.push((id, t));
}
}
Ok(out)
}
/// The saved index read-only whatever model made it (keyword search only needs its text).
pub fn open_for_keywords(mind: &Mind) -> Option<Index> {
use rusqlite::OpenFlags;
let path = mind.home.join("index.db");
if !path.is_file() {
return None;
}
// A plain read-only open sees the newest changes (the -wal file); where even that can't work
// (a read-only folder with no -shm file), the saved file as it is.
let plain = Connection::open_with_flags(&path, OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX)
.ok()
.filter(|c| c.busy_timeout(std::time::Duration::from_millis(500)).is_ok() && c.query_row("SELECT 1 FROM entries LIMIT 1", [], |_| Ok(())).is_ok());
let conn = match plain {
Some(c) => c,
None => Connection::open_with_flags(immutable_uri(&path), OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_URI | OpenFlags::SQLITE_OPEN_NO_MUTEX).ok()?,
};
Self::with_conn_readonly(conn, mind).ok()
}
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)
}
}
/// A SQLite error meaning the file is damaged or isn't a database at all.
/// A `file:` URI that opens `path` as immutable: no locks and no -wal/-shm files (they can't be
/// created where the folder isn't writable). Absolute first (a relative MIND_HOME would otherwise
/// turn into a path from the root).
fn immutable_uri(path: &std::path::Path) -> String {
let abs = std::path::absolute(path).unwrap_or_else(|_| path.to_path_buf());
let p = abs.to_string_lossy().replace('\\', "/").replace('%', "%25").replace('?', "%3f").replace('#', "%23");
if p.starts_with('/') { format!("file:{p}?immutable=1") } else { format!("file:/{p}?immutable=1") }
}
/// The words of a chat message worth matching on: four letters or more, not common words, each once.
pub(crate) fn meaningful_words(message: &str) -> Vec<String> {
const COMMON: &[&str] = &[
"this", "that", "these", "those", "there", "their", "them", "they", "then", "than", "with", "without", "from",
"into", "onto", "about", "after", "before", "what", "when", "where", "which", "while", "who", "whom", "whose",
"have", "having", "been", "being", "were", "will", "would", "could", "should", "shall", "might", "must", "does",
"doing", "done", "make", "made", "need", "needs", "want", "wants", "like", "just", "also", "only", "some", "more",
"most", "much", "many", "very", "your", "yours", "ours", "here", "please", "thanks", "okay", "sure", "yeah", "know",
"think", "look", "looks", "work", "works", "working", "using", "used", "uses", "code", "file", "files", "thing",
"things", "into", "over", "under", "again", "still", "even", "each", "every", "other", "another", "same", "right",
"good", "great", "better", "best", "change", "changes", "check", "help", "show", "tell", "give", "take", "keep",
"let's", "lets", "dont", "don't", "can't", "cant", "isn't", "it's", "that's", "what's", "there's", "how's", "now",
"today", "something", "anything", "everything", "nothing", "because", "since", "until", "through", "between",
"build", "builds", "test", "tests", "error", "errors", "issue", "issues", "problem", "fail", "fails", "failing",
"failed", "broken", "update", "fixed", "implement", "feature", "project", "function", "class", "method", "line",
"lines", "output", "input", "value", "data", "run", "runs", "running", "start", "stop", "load", "save", "open",
"close", "write", "read", "find", "found", "remove", "delete", "create", "setup", "config", "version", "issue",
// Conversation, not topic: rare in notes, so they'd count as distinctive.
"else", "anyway", "maybe", "already", "really", "quick", "quickly", "ready", "next", "last", "first", "later",
"soon", "kind", "sort", "stuff", "whatever", "mean", "means", "seems", "seem", "feel", "guess", "wonder",
"should've", "would've", "could've", "also", "either", "neither", "both", "else's", "around", "enough",
];
let mut out: Vec<String> = vec![];
for w in message.split(|c: char| !c.is_alphanumeric() && c != '\'' && c != '_' && c != '-') {
let w = w.trim_matches(|c: char| c == '\'' || c == '-').to_lowercase();
if w.chars().count() >= 4 && !COMMON.contains(&w.as_str()) && !w.contains('\'') && !out.contains(&w) {
out.push(w);
}
if out.len() >= 12 {
break;
}
}
out
}
/// Which model made an index's vectors. Vectors from a different model aren't comparable, so the
/// index remembers it.
fn model_id(cfg: &crate::store::Embeddings) -> String {
match cfg.provider.trim().to_lowercase().as_str() {
"local" => LOCAL_MODEL_ID.to_string(),
"none" | "" => "none".to_string(),
// The endpoint too: two Azure resources can both call their deployment `embeddings`.
p => format!("{p}:{}:{}", cfg.url.trim().trim_end_matches('/'), cfg.model),
}
}
/// Whether an error is the index (or its folder) not being writable: a sandbox such as Codex's
/// that blocks writes outside the project, a read-only disk, missing permissions.
pub(crate) fn is_unwritable(e: &anyhow::Error) -> bool {
e.chain().any(|c| {
c.downcast_ref::<rusqlite::Error>().is_some_and(|e| {
matches!(
e.sqlite_error_code(),
Some(rusqlite::ErrorCode::CannotOpen | rusqlite::ErrorCode::ReadOnly | rusqlite::ErrorCode::PermissionDenied)
)
}) || c.downcast_ref::<std::io::Error>().is_some_and(|e| {
matches!(e.kind(), std::io::ErrorKind::PermissionDenied | std::io::ErrorKind::ReadOnlyFilesystem)
})
})
}
pub(crate) fn is_corrupt(e: &anyhow::Error) -> bool {
e.chain().any(|c| {
c.downcast_ref::<rusqlite::Error>().is_some_and(|e| {
matches!(e.sqlite_error_code(), Some(rusqlite::ErrorCode::DatabaseCorrupt | rusqlite::ErrorCode::NotADatabase))
})
})
}
/// Whether an embeddings error is the service refusing this input (too long, invalid, filtered),
/// rather than being busy or unreachable.
fn is_refusal(e: &anyhow::Error) -> bool {
let m = format!("{e:#}");
["HTTP 400", "HTTP 413", "HTTP 422"].iter().any(|c| m.contains(c))
}
