pub mod add;
pub mod chunk;
pub mod clients;
pub mod daemon;
pub mod embed;
pub mod filter;
pub mod fleet;
mod formats;
pub mod graph;
pub mod home;
pub mod http;
pub mod image;
pub mod index;
pub mod lock;
pub mod mcp;
pub mod portal;
pub mod proxy;
pub mod routes;
pub mod setup;
pub mod store;
pub mod upgrade;
pub mod watch;
use anyhow::{Result, bail};
use embed::Model;
use fastembed::TextEmbedding;
use filter::Filter;
use index::{Allowlist, VectorIndex};
use rusqlite::Connection;
use serde::Serialize;
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};
const BIT_WIDTH: usize = 4;
const EMBED_BATCH: usize = 8;
pub fn default_model() -> Model {
Model::Granite
}
const PENDING: &str = "";
const GENERATION: &str = "index_generation";
const RRF_K: f32 = 60.0;
type Candidate = (((bool, u64), f32), Vec<&'static str>);
const TEXT_LISTS: f32 = 3.0;
const IMAGE_FLOOR: f32 = 0.28;
pub const IMAGE_FLOOR_ENV: &str = "SEMLITH_IMAGE_FLOOR";
fn image_floor() -> f32 {
std::env::var(IMAGE_FLOOR_ENV)
.ok()
.and_then(|v| v.parse::<f32>().ok())
.unwrap_or(IMAGE_FLOOR)
}
const RANK_DEPTH: usize = 4;
const CHECKPOINT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(30);
const CHECKPOINT_SECS_ENV: &str = "SEMLITH_CHECKPOINT_SECS";
fn checkpoint_interval() -> std::time::Duration {
match std::env::var(CHECKPOINT_SECS_ENV)
.ok()
.and_then(|v| v.parse().ok())
{
Some(secs) => std::time::Duration::from_secs(secs),
None => CHECKPOINT_INTERVAL,
}
}
#[derive(Debug, Clone, Serialize)]
pub struct Hit {
pub score: f32,
pub path: String,
pub start_line: u32,
pub end_line: u32,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub lists: Vec<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub image: Option<Pixels>,
}
#[derive(Debug, Clone, Copy, serde::Serialize)]
pub struct Pixels {
pub width: u32,
pub height: u32,
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub enum FileOutcome {
#[default]
Indexing,
Unchanged,
Skipped,
Removed,
Refused,
}
impl FileOutcome {
pub fn as_str(self) -> &'static str {
match self {
Self::Indexing => "indexing",
Self::Unchanged => "unchanged",
Self::Skipped => "skipped",
Self::Removed => "removed",
Self::Refused => "refused",
}
}
}
#[derive(Debug, Default, Clone)]
pub struct Boundary {
pub roots: Option<Vec<PathBuf>>,
pub allow_secrets: bool,
}
impl Boundary {
pub fn within(roots: Vec<PathBuf>) -> Self {
Self {
roots: Some(roots),
allow_secrets: false,
}
}
pub fn refuses(&self, path: &Path) -> Option<String> {
if let Some(roots) = &self.roots
&& !filter::within_boundary(path, roots)
{
return Some(
"is outside this store's roots and outside the home directory, so \
semlith will not index it. Add it as a root first, or index it \
from the command line."
.to_string(),
);
}
if !self.allow_secrets
&& let Some(why) = filter::denied(path)
{
return Some(why.reason());
}
None
}
}
#[derive(Debug, Default, Clone, Copy)]
pub struct IndexProgress {
pub outcome: FileOutcome,
pub scanned: usize,
pub indexed: usize,
pub chunks: usize,
pub total: usize,
pub symbols: usize,
}
#[derive(Debug, Default, Clone, Serialize)]
pub struct IndexReport {
pub scanned: usize,
pub indexed: usize,
pub unchanged: usize,
pub skipped: usize,
pub removed: usize,
pub chunks: usize,
pub remaining: usize,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub refused: Vec<(String, String)>,
pub symbols: usize,
pub edges: usize,
pub images: usize,
pub stopped: bool,
pub written: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flow {
Run,
Pause,
Yield,
Stop,
}
const PAUSE_TICK: std::time::Duration = std::time::Duration::from_millis(120);
pub struct Semlith {
dir: PathBuf,
db: Connection,
index: VectorIndex,
images: VectorIndex,
model: Model,
dim: usize,
embedder: Option<TextEmbedding>,
clip: image::Clip,
generation: u64,
pub quiet: bool,
pub boundary: Boundary,
}
impl Semlith {
pub fn open(dir: impl AsRef<Path>, model: Option<Model>) -> Result<Self> {
let dir = dir.as_ref().to_path_buf();
home::secure_dir(&dir)?;
let db_path = dir.join("store.db");
let db = store::open(&db_path)?;
home::tighten_file(&db_path);
for beside in ["store.db-wal", "store.db-shm"] {
home::tighten_file(&dir.join(beside));
}
let model = match store::get_meta(&db, "model")? {
Some(existing) => {
let existing: Model = existing.parse().map_err(anyhow::Error::msg)?;
if let Some(want) = model
&& want != existing
{
bail!(
"store was built with {existing}, not {want}; \
delete {} to rebuild with a different model",
dir.display()
);
}
existing
}
None => {
let m = model.unwrap_or_else(default_model);
let _writing = store::Writing::begin(&db)?;
store::set_meta(&db, "model", &m.to_string())?;
store::set_meta(&db, store::FORMAT_KEY, &store::FORMAT_VERSION.to_string())?;
m
}
};
let dim = model.dim()?;
let sharded = store::format(&db)? >= store::SHARDED_FORMAT;
let index = VectorIndex::open(&dir, dim, BIT_WIDTH, sharded)?;
let images = VectorIndex::open(&dir.join(image::INDEX_DIR), image::DIM, BIT_WIDTH, true)?;
let generation = generation(&db)?;
Ok(Self {
dir,
db,
index,
images,
model,
dim,
embedder: None,
clip: image::Clip::default(),
generation,
quiet: false,
boundary: Boundary::default(),
})
}
pub fn open_existing(dir: impl AsRef<Path>) -> Result<Self> {
let dir = dir.as_ref();
if !dir.join("store.db").exists() {
bail!(
"{} is not a semlith store — no store.db in it; index it first",
dir.display()
);
}
Self::open(dir, None)
}
pub fn refresh(&mut self) -> Result<()> {
let current = generation(&self.db)?;
if current == self.generation {
return Ok(());
}
if !self.index.exists() {
return Ok(());
}
let was_resident = self.index.is_resident();
self.index.evict();
if was_resident {
self.index.prepare()?;
}
self.generation = current;
Ok(())
}
pub fn model(&self) -> &Model {
&self.model
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn len(&self) -> usize {
if let Some(resident) = self.index.resident_len() {
return resident;
}
store::durable_chunks(&self.db).unwrap_or(0) as usize
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn db(&self) -> &Connection {
&self.db
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn shards(&self) -> Option<(usize, usize)> {
self.index
.max_resident()
.map(|max| (self.index.shards(), max))
}
pub fn evictions(&self) -> u64 {
self.index.evictions()
}
fn embedder(&mut self) -> Result<&mut TextEmbedding> {
if self.embedder.is_none() {
self.embedder = Some(self.model.load(
model_cache_dir(),
chunk::MAX_CHARS / 2,
self.quiet,
)?);
}
Ok(self.embedder.as_mut().unwrap())
}
pub fn warm(&mut self) -> Result<()> {
self.embedder()?;
self.index.prepare()?;
Ok(())
}
pub fn warm_index(&mut self) -> Result<()> {
self.index.prepare()
}
fn embed(&mut self, texts: Vec<String>) -> Result<Vec<Vec<f32>>> {
let mut out = self
.embedder()?
.embed(texts, Some(EMBED_BATCH))
.map_err(|e| anyhow::anyhow!("{e}"))?;
for v in &mut out {
normalize(v);
}
Ok(out)
}
pub fn index_paths(
&mut self,
roots: &[PathBuf],
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.index_walk(roots, on_file)
}
pub fn index_paths_within(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.index_within_held(roots, budget, on_file)
}
pub(crate) fn index_within_held(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let deadline = std::time::Instant::now() + budget;
self.index_set(walk(roots), true, Some(deadline), None, on_file)
}
pub(crate) fn index_walk_under(
&mut self,
roots: &[PathBuf],
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(walk(roots), true, None, Some(control), on_file)
}
pub(crate) fn index_within_held_under(
&mut self,
roots: &[PathBuf],
budget: std::time::Duration,
control: &dyn Fn() -> Flow,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
let deadline = std::time::Instant::now() + budget;
self.index_set(walk(roots), true, Some(deadline), Some(control), on_file)
}
pub(crate) fn index_walk(
&mut self,
roots: &[PathBuf],
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(walk(roots), true, None, None, on_file)
}
pub(crate) fn index_changed(
&mut self,
paths: Vec<PathBuf>,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index_set(paths, false, None, None, on_file)
}
fn index_set(
&mut self,
paths: Vec<PathBuf>,
sweep: bool,
deadline: Option<std::time::Instant>,
control: Option<&dyn Fn() -> Flow>,
on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.writing(move |me| me.index_set_writing(paths, sweep, deadline, control, on_file))
}
fn writing<T>(&mut self, body: impl FnOnce(&mut Self) -> Result<T>) -> Result<T> {
store::read_only(&self.db, false)?;
let out = body(self);
let _ = store::read_only(&self.db, true);
out
}
fn index_set_writing(
&mut self,
paths: Vec<PathBuf>,
sweep: bool,
deadline: Option<std::time::Instant>,
control: Option<&dyn Fn() -> Flow>,
mut on_file: impl FnMut(&Path, IndexProgress),
) -> Result<IndexReport> {
self.index.clean();
let mut report = IndexReport::default();
let mut pending = Batch::default();
let mut written: Vec<String> = Vec::new();
let mut completed: Vec<(i64, String)> = Vec::new();
let checkpointing = matches!(self.index, index::VectorIndex::Sharded(_));
let interval = checkpoint_interval();
let mut last_checkpoint = std::time::Instant::now();
let (paths, refused): (Vec<PathBuf>, Vec<(PathBuf, String)>) = {
let mut allowed = Vec::with_capacity(paths.len());
let mut refused = Vec::new();
for path in paths {
match self.boundary.refuses(&path) {
Some(why) => refused.push((path, why)),
None => allowed.push(path),
}
}
(allowed, refused)
};
let total = paths.len() + refused.len();
for (path, why) in &refused {
report
.refused
.push((path.display().to_string(), why.clone()));
report.scanned += 1;
on_file(
path,
IndexProgress {
outcome: FileOutcome::Refused,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
}
for (seen, path) in paths.into_iter().enumerate() {
if let Some(deadline) = deadline
&& report.indexed > 0
&& std::time::Instant::now() >= deadline
{
report.remaining = total - seen;
break;
}
if let Some(ask) = control {
let mut stop = false;
let mut yielded = false;
loop {
match ask() {
Flow::Run => break,
Flow::Pause => std::thread::sleep(PAUSE_TICK),
Flow::Yield => {
yielded = true;
break;
}
Flow::Stop => {
stop = true;
break;
}
}
}
if yielded {
report.remaining = total - seen;
break;
}
if stop {
report.remaining = total - seen;
report.stopped = true;
break;
}
}
report.scanned += 1;
let key = path.to_string_lossy().into_owned();
let opened = std::fs::File::open(&path);
let measured = opened.as_ref().ok().and_then(|f| f.metadata().ok());
match measured {
Some(m) if m.is_file() && m.len() > 0 && m.len() <= chunk::MAX_FILE_BYTES => {}
_ => {
if !path.exists() {
let ids = store::delete_file(&self.db, &key)?;
if !ids.is_empty() {
for id in ids {
self.index.remove(id)?;
}
report.removed += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Removed,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
}
report.skipped += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Skipped,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
}
let read = opened.and_then(|file| {
use std::io::Read;
let mut bytes = Vec::new();
file.take(chunk::MAX_FILE_BYTES + 1)
.read_to_end(&mut bytes)
.map(|_| bytes)
});
let bytes = match read {
Ok(bytes) if bytes.len() as u64 <= chunk::MAX_FILE_BYTES => bytes,
_ => {
report.skipped += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Skipped,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
};
let hash = blake3::hash(&bytes).to_hex().to_string();
if store::file_hash(&self.db, &key)?.as_deref() == Some(hash.as_str()) {
report.unchanged += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Unchanged,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
if image::is_image(&path) {
if let Some(why) = image::too_large(&bytes) {
report.refused.push((path.display().to_string(), why));
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Refused,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
let Some((width, height)) = image::dimensions(&bytes) else {
report.skipped += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Skipped,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
};
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Indexing,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
let vector = self.clip.embed_image(&path, self.quiet)?;
for id in store::image_ids_of(&self.db, &key)? {
self.images.remove(id as u64)?;
}
for id in store::delete_file(&self.db, &key)? {
self.index.remove(id)?;
}
let file_id =
store::insert_file(&self.db, &key, PENDING, bytes.len() as u64, now())?;
let image_id = store::insert_image(&self.db, file_id, width, height)?;
self.images.add(&vector, &[image_id as u64])?;
completed.push((file_id, hash));
written.push(key.clone());
report.indexed += 1;
report.images += 1;
continue;
}
let Some(text) = chunk::extract(&path, &bytes) else {
report.skipped += 1;
continue;
};
let chunks = chunk::chunk_text(&text);
if chunks.is_empty() {
report.skipped += 1;
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Skipped,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
continue;
}
on_file(
&path,
IndexProgress {
outcome: FileOutcome::Indexing,
scanned: report.scanned,
indexed: report.indexed,
chunks: report.chunks,
total,
symbols: report.symbols,
},
);
for id in store::delete_file(&self.db, &key)? {
self.index.remove(id)?;
}
let file_id = store::insert_file(&self.db, &key, PENDING, bytes.len() as u64, now())?;
let mut spans: Vec<(u32, u32, i64)> = Vec::with_capacity(chunks.len());
for (ord, c) in chunks.iter().enumerate() {
let id =
store::insert_chunk(&self.db, file_id, ord, c.start_line, c.end_line, &c.text)?;
spans.push((c.start_line, c.end_line, id));
pending.ids.push(id as u64);
pending.texts.push(c.text.clone());
if pending.ids.len() >= EMBED_BATCH {
self.flush(&mut pending)?;
}
}
let (symbols, edges) = self.extract_graph(&path, &text, file_id, &spans)?;
report.symbols += symbols;
report.edges += edges;
completed.push((file_id, hash));
written.push(key.clone());
report.indexed += 1;
report.chunks += chunks.len();
if checkpointing && last_checkpoint.elapsed() >= interval {
self.flush(&mut pending)?;
self.checkpoint(&mut completed)?;
last_checkpoint = std::time::Instant::now();
}
}
self.flush(&mut pending)?;
report.written = written;
if sweep {
for key in store::all_paths(&self.db)? {
if !Path::new(&key).exists() {
for id in store::delete_file(&self.db, &key)? {
self.index.remove(id)?;
}
report.removed += 1;
}
}
}
if report.indexed > 0 || report.removed > 0 || !self.index.exists() {
self.save()?;
}
self.commit_hashes(&mut completed)?;
Ok(report)
}
fn checkpoint(&mut self, completed: &mut Vec<(i64, String)>) -> Result<()> {
if completed.is_empty() {
return Ok(());
}
self.save()?;
self.commit_hashes(completed)
}
fn graph_expansion(
&self,
dense: &[u64],
keyword: &[u64],
depth: usize,
filter: &Filter,
) -> Result<Vec<u64>> {
const SEEDS: usize = 8;
let seeds: Vec<u64> = dense
.iter()
.take(SEEDS)
.chain(keyword.iter().take(SEEDS))
.copied()
.collect();
if seeds.is_empty() {
return Ok(Vec::new());
}
let names = store::symbols_in_chunks(&self.db, &seeds)?;
if names.is_empty() {
return Ok(Vec::new());
}
let mut neighbours: Vec<String> = Vec::new();
for name in &names {
for end in store::edges_out(&self.db, name, &[])?
.into_iter()
.chain(store::edges_in(&self.db, name, &[])?)
{
if !neighbours.contains(&end.symbol.name) {
neighbours.push(end.symbol.name);
}
}
if neighbours.len() >= depth * 4 {
break;
}
}
let mut ids = Vec::new();
for symbol in store::symbols_by_names(&self.db, &neighbours, filter.groups())? {
let Some(chunk_id) = symbol.chunk_id else {
continue;
};
let id = chunk_id as u64;
if !ids.contains(&id) {
ids.push(id);
}
if ids.len() >= depth {
break;
}
}
Ok(ids)
}
fn extract_graph(
&self,
path: &Path,
text: &str,
file_id: i64,
spans: &[(u32, u32, i64)],
) -> Result<(usize, usize)> {
let Some(extraction) = graph::extract(path, text)? else {
return Ok((0, 0));
};
let mut ids: std::collections::HashMap<&str, i64> = std::collections::HashMap::new();
for symbol in &extraction.symbols {
let chunk_id = spans
.iter()
.find(|(start, end, _)| symbol.start_line >= *start && symbol.start_line <= *end)
.map(|(_, _, id)| *id);
let id = store::insert_symbol(&self.db, file_id, chunk_id, symbol)?;
ids.entry(symbol.name.as_str()).or_insert(id);
}
let mut written = 0;
for edge in &extraction.edges {
let Some(src) = ids.get(edge.from.as_str()) else {
continue;
};
store::insert_edge(&self.db, *src, &edge.to, &edge.kind, &edge.confidence)?;
written += 1;
}
Ok((extraction.symbols.len(), written))
}
fn commit_hashes(&mut self, completed: &mut Vec<(i64, String)>) -> Result<()> {
if completed.is_empty() {
return Ok(());
}
let tx = self.db.unchecked_transaction()?;
for (file_id, hash) in completed.iter() {
tx.execute(
"UPDATE files SET hash = ?1 WHERE id = ?2",
rusqlite::params![hash, file_id],
)?;
}
tx.commit()?;
completed.clear();
Ok(())
}
fn flush(&mut self, batch: &mut Batch) -> Result<()> {
if batch.ids.is_empty() {
return Ok(());
}
let ids = std::mem::take(&mut batch.ids);
let texts = std::mem::take(&mut batch.texts);
let vectors = self.embed(texts)?;
let flat: Vec<f32> = vectors.into_iter().flatten().collect();
self.index.add(&flat, &ids)?;
Ok(())
}
pub fn forget(&mut self, path: &Path) -> Result<(usize, usize)> {
let _lock = lock::StoreLock::acquire(&self.dir)?;
self.forget_held(path)
}
pub(crate) fn forget_held(&mut self, path: &Path) -> Result<(usize, usize)> {
self.writing(|me| me.forget_writing(path))
}
fn forget_writing(&mut self, path: &Path) -> Result<(usize, usize)> {
let key = canonical(path).to_string_lossy().into_owned();
let images = store::image_ids_of(&self.db, &key)?;
let ids = store::delete_file(&self.db, &key)?;
for id in &ids {
self.index.remove(*id)?;
}
for id in &images {
self.images.remove(*id as u64)?;
}
if !ids.is_empty() || !images.is_empty() {
self.save()?;
}
Ok((ids.len(), images.len()))
}
pub fn search(&mut self, query: &str, k: usize) -> Result<Vec<Hit>> {
self.search_filtered(query, k, &Filter::default())
}
fn allowlist(&mut self, filter: &Filter) -> Result<Allowlist> {
if filter.is_empty() {
return Ok(Allowlist::All);
}
let candidates = store::filtered_chunk_ids(&self.db, filter.groups())?;
let mut ids = Vec::with_capacity(candidates.len());
for id in candidates {
if self.index.contains(id)? {
ids.push(id);
}
}
Ok(if ids.is_empty() {
Allowlist::Empty
} else if ids.len() == self.len() {
Allowlist::All
} else {
Allowlist::Subset(ids)
})
}
pub fn matching_files(&self, filter: &Filter) -> Result<i64> {
store::matching_files(&self.db, filter.groups())
}
pub fn matching_paths(&self, filter: &Filter) -> Result<Vec<String>> {
store::filtered_paths(&self.db, filter.groups())
}
pub fn search_filtered(&mut self, query: &str, k: usize, filter: &Filter) -> Result<Vec<Hit>> {
let vector = self.embed_query(query)?;
self.search_with_vector(query, &vector, k, filter)
}
pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
Ok(self.embed(vec![self.model.query_text(query)])?.remove(0))
}
pub fn search_with_vector(
&mut self,
query: &str,
vector: &[f32],
k: usize,
filter: &Filter,
) -> Result<Vec<Hit>> {
Ok(self
.search_ranked(query, vector, k, filter)?
.into_iter()
.map(|(hit, _)| hit)
.collect())
}
pub fn search_ranked(
&mut self,
query: &str,
vector: &[f32],
k: usize,
filter: &Filter,
) -> Result<Vec<(Hit, f32)>> {
self.refresh()?;
if k == 0 || (self.is_empty() && store::image_count(&self.db)? == 0) {
return Ok(Vec::new());
}
let depth = (k * RANK_DEPTH).max(k);
let allowlist = self.allowlist(filter)?;
if matches!(allowlist, Allowlist::Empty) {
return Ok(Vec::new());
}
let (dense_scores, dense_ids) = self.index.search(vector, depth, &allowlist)?;
let keyword_ids = store::keyword_search(&self.db, query, depth, filter.groups())?;
let images = self.image_search(query, depth, filter)?;
let graph_ids = self.graph_expansion(&dense_ids, &keyword_ids, depth, filter)?;
let mut fused: Vec<((bool, u64), f32)> = Vec::new();
let mut seen: std::collections::HashMap<(bool, u64), usize> =
std::collections::HashMap::new();
let mut lists: Vec<Vec<&'static str>> = Vec::new();
let image_list: Vec<(u64, f32)> = images
.iter()
.map(|(id, similarity)| {
let weight = if *similarity >= image_floor() {
TEXT_LISTS
} else {
1.0
};
(*id, weight)
})
.collect();
for (name, is_image, ranking) in [
("image", true, image_list),
(
"vector",
false,
dense_ids.iter().map(|id| (*id, 1.0)).collect(),
),
(
"keyword",
false,
keyword_ids.iter().map(|id| (*id, 1.0)).collect(),
),
(
"graph",
false,
graph_ids.iter().map(|id| (*id, 1.0)).collect(),
),
] {
for (rank, (id, weight)) in ranking.iter().enumerate() {
let key = (is_image, *id);
let contribution = weight / (RRF_K + rank as f32 + 1.0);
match seen.get(&key) {
Some(&slot) => {
fused[slot].1 += contribution;
lists[slot].push(name);
}
None => {
seen.insert(key, fused.len());
fused.push((key, contribution));
lists.push(vec![name]);
}
}
}
}
let mut ranked: Vec<Candidate> = fused.into_iter().zip(lists).collect();
ranked.sort_by(|a, b| b.0.1.total_cmp(&a.0.1));
ranked.truncate(k);
let mut hits = Vec::with_capacity(ranked.len());
for (((is_image, id), score), found_by) in ranked {
if is_image {
if let Some(row) = store::image(&self.db, id as i64)? {
hits.push((
Hit {
score,
path: row.path,
start_line: 0,
end_line: 0,
text: String::new(),
store: None,
lists: found_by,
image: Some(Pixels {
width: row.width,
height: row.height,
}),
},
0.0,
));
}
continue;
}
if let Some(row) = store::chunk(&self.db, id)? {
let similarity = dense_ids
.iter()
.position(|d| *d == id)
.and_then(|i| dense_scores.get(i).copied())
.unwrap_or(0.0);
hits.push((
Hit {
score,
path: row.path,
start_line: row.start_line,
end_line: row.end_line,
text: row.text,
store: None,
lists: found_by,
image: None,
},
similarity,
));
}
}
Ok(hits)
}
fn image_search(
&mut self,
query: &str,
depth: usize,
filter: &Filter,
) -> Result<Vec<(u64, f32)>> {
if store::image_count(&self.db)? == 0 {
return Ok(Vec::new());
}
let allowlist = if filter.is_empty() {
index::Allowlist::All
} else {
let mut ids = Vec::new();
for id in store::filtered_image_ids(&self.db, filter.groups())? {
if self.images.contains(id as u64)? {
ids.push(id as u64);
}
}
if ids.is_empty() {
return Ok(Vec::new());
}
index::Allowlist::Subset(ids)
};
let mut vector = self.clip.embed_query(query, self.quiet)?;
normalize(&mut vector);
let (scores, ids) = self.images.search(&vector, depth, &allowlist)?;
Ok(ids
.into_iter()
.zip(scores.into_iter().chain(std::iter::repeat(0.0)))
.collect())
}
pub fn image_count(&self) -> Result<i64> {
store::image_count(&self.db)
}
pub fn stats(&self) -> Result<(i64, i64, i64)> {
store::stats(&self.db)
}
pub fn save(&mut self) -> Result<()> {
self.index.save()?;
self.images.save()?;
self.generation = generation(&self.db)? + 1;
store::set_meta(&self.db, GENERATION, &self.generation.to_string())?;
Ok(())
}
}
#[derive(Default)]
struct Batch {
ids: Vec<u64>,
texts: Vec<String>,
}
pub(crate) fn normalize(v: &mut [f32]) {
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in v {
*x /= norm;
}
}
}
fn generation(db: &Connection) -> Result<u64> {
Ok(store::get_meta(db, GENERATION)?
.and_then(|v| v.parse().ok())
.unwrap_or(0))
}
fn now() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0)
}
pub fn human_bytes(bytes: i64) -> String {
const KB: f64 = 1024.0;
let b = bytes as f64;
if b >= KB * KB {
format!("{:.1} MB", b / (KB * KB))
} else if b >= KB {
format!("{:.0} KB", b / KB)
} else {
format!("{bytes} B")
}
}
pub fn canonical(path: &Path) -> PathBuf {
std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
}
pub const MODEL_CACHE_ENV: &str = "SEMLITH_MODEL_CACHE";
pub fn model_cache_dir() -> PathBuf {
if let Ok(dir) = std::env::var(MODEL_CACHE_ENV) {
return PathBuf::from(dir);
}
let base = std::env::var("HOME")
.map(PathBuf::from)
.unwrap_or_else(|_| PathBuf::from("."));
base.join(".cache").join("semlith").join("models")
}
fn walk(roots: &[PathBuf]) -> Vec<PathBuf> {
let mut out = Vec::new();
let mut seen = std::collections::HashSet::new();
for root in roots {
let mut builder = ignore::WalkBuilder::new(root);
builder
.hidden(true)
.git_ignore(true)
.git_global(true)
.git_exclude(true)
.parents(true)
.require_git(false)
.filter_entry(|e| e.file_name() != ".semlith");
for result in builder.build() {
let entry = match result {
Ok(e) => e,
Err(e) => {
eprintln!("semlith: skipping unreadable path: {e}");
continue;
}
};
if !entry.file_type().is_some_and(|t| t.is_file()) {
continue;
}
let path = canonical(entry.path());
if seen.insert(path.clone()) {
out.push(path);
}
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn normalize_gives_unit_length() {
let mut v = vec![3.0, 4.0];
normalize(&mut v);
assert!((v.iter().map(|x| x * x).sum::<f32>() - 1.0).abs() < 1e-6);
let mut z = vec![0.0, 0.0];
normalize(&mut z);
assert_eq!(z, vec![0.0, 0.0]);
}
#[test]
fn one_large_file_does_not_queue_more_than_a_batch() {
let text = "a line of perfectly ordinary text\n".repeat(4000);
let chunks = chunk::chunk_text(&text);
assert!(
chunks.len() > EMBED_BATCH * 4,
"test file is too small to prove anything: {} chunks",
chunks.len()
);
let mut queued = 0usize;
let mut high_water = 0usize;
for _ in &chunks {
queued += 1;
high_water = high_water.max(queued);
if queued >= EMBED_BATCH {
queued = 0;
}
}
assert_eq!(
high_water, EMBED_BATCH,
"queue grew past one batch within a single file"
);
}
#[test]
fn a_new_store_records_the_default_model_and_reopens_on_it() {
let dir = tempdir();
let created = Semlith::open(&dir, None).unwrap();
assert_eq!(*created.model(), default_model());
assert_eq!(created.dim(), 384);
drop(created);
let reopened = Semlith::open(&dir, None).unwrap();
assert_eq!(*reopened.model(), default_model());
}
#[test]
fn an_existing_store_keeps_its_own_model() {
use fastembed::EmbeddingModel;
let dir = tempdir2();
let old = Model::Builtin(EmbeddingModel::BGESmallENV15);
let created = Semlith::open(&dir, Some(old.clone())).unwrap();
assert_eq!(*created.model(), old);
drop(created);
let reopened = Semlith::open(&dir, None).unwrap();
assert_eq!(*reopened.model(), old);
assert!(Semlith::open(&dir, Some(Model::Granite)).is_err());
}
fn tempdir() -> PathBuf {
let d = std::env::temp_dir().join(format!("semlith-test-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
fn tempdir2() -> PathBuf {
let d = std::env::temp_dir().join(format!("semlith-test-b-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
}