use crate::embed::Model;
use crate::filter::Filter;
use crate::{Hit, Semlith, canonical, chunk, model_cache_dir};
use anyhow::{Context, Result, bail};
use fastembed::TextEmbedding;
use std::collections::VecDeque;
use std::path::{Path, PathBuf};
pub struct Fleet {
members: Vec<Member>,
embedders: Vec<(Model, TextEmbedding)>,
tokenizer: Option<tokenizers::Tokenizer>,
embeds: usize,
pub quiet: bool,
}
struct Member {
label: String,
store: Semlith,
}
impl Fleet {
pub fn empty() -> Self {
Self {
members: Vec::new(),
embedders: Vec::new(),
tokenizer: None,
embeds: 0,
quiet: true,
}
}
pub fn name_from(&mut self, named: &[(PathBuf, String)]) {
for member in &mut self.members {
if let Some(name) = name_for(member.store.dir(), named) {
member.label = name;
}
}
}
pub fn open(dirs: &[PathBuf]) -> Result<Self> {
if dirs.is_empty() {
bail!("no store given");
}
let mut keys: Vec<PathBuf> = Vec::new();
let mut members = Vec::new();
for dir in dirs {
let store = Semlith::open_existing(dir)?;
let key = canonical(dir);
if keys.contains(&key) {
continue;
}
keys.push(key);
members.push(Member {
label: String::new(),
store,
});
}
for (member, key) in members.iter_mut().zip(&keys) {
member.label = label(key);
}
disambiguate(&mut members, &keys);
Ok(Self {
members,
embedders: Vec::new(),
tokenizer: None,
embeds: 0,
quiet: false,
})
}
pub fn len(&self) -> usize {
self.members.len()
}
pub fn is_empty(&self) -> bool {
self.members.is_empty()
}
pub fn each(&self) -> impl Iterator<Item = (&str, &Semlith)> {
self.members.iter().map(|m| (m.label.as_str(), &m.store))
}
pub fn labels(&self) -> Vec<&str> {
self.members.iter().map(|m| m.label.as_str()).collect()
}
pub fn chunks(&self) -> usize {
self.members.iter().map(|m| m.store.len()).sum()
}
pub fn files(&self) -> Result<i64> {
let mut total = 0;
for member in &self.members {
total += member.store.stats()?.0;
}
Ok(total)
}
pub fn query_embeds(&self) -> usize {
self.embeds
}
pub fn warm(&mut self) -> Result<()> {
for i in 0..self.members.len() {
let model = self.members[i].store.model().clone();
self.embedder(&model)?;
}
Ok(())
}
pub fn matching_files(&self, filter: &Filter) -> Result<i64> {
let mut total = 0;
for member in &self.members {
total += member.store.matching_files(filter)?;
}
Ok(total)
}
pub fn search(&mut self, query: &str, k: usize) -> Result<Vec<Hit>> {
self.search_filtered(query, k, &Filter::default())
}
pub fn search_filtered(&mut self, query: &str, k: usize, filter: &Filter) -> Result<Vec<Hit>> {
self.search_in(None, query, k, filter)
}
pub fn search_in(
&mut self,
only: Option<&[String]>,
query: &str,
k: usize,
filter: &Filter,
) -> Result<Vec<Hit>> {
self.search_preferring(only, query, k, filter, crate::Prefer::default())
}
pub fn search_preferring(
&mut self,
only: Option<&[String]>,
query: &str,
k: usize,
filter: &Filter,
prefer: crate::Prefer,
) -> Result<Vec<Hit>> {
let chosen = self.chosen(only)?;
let label_hits = self.members.len() > 1;
let mut vectors: Vec<(Model, Vec<f32>)> = Vec::new();
let mut queues: Vec<Vec<(Hit, f32)>> = Vec::new();
for i in chosen {
let model = self.members[i].store.model().clone();
let vector = match vectors.iter().find(|(m, _)| *m == model) {
Some((_, v)) => v.clone(),
None => {
let v = self.embed_query(&model, query)?;
vectors.push((model, v.clone()));
v
}
};
let hits = self.members[i]
.store
.search_preferring(query, &vector, k, filter, prefer)?;
let label = self.members[i].label.clone();
queues.push(
hits.into_iter()
.map(|(mut h, similarity)| {
if label_hits {
h.store = Some(label.clone());
}
(h, similarity)
})
.collect(),
);
}
Ok(merge(queues, k))
}
pub fn paths_in(
&self,
only: Option<&[String]>,
filter: &Filter,
limit: usize,
) -> Result<(Vec<String>, usize)> {
let chosen = self.chosen(only)?;
let label = self.members.len() > 1;
let mut out = Vec::new();
let mut total = 0;
for i in chosen {
for path in self.members[i].store.matching_paths(filter)? {
total += 1;
if out.len() < limit {
out.push(if label {
format!("[{}] {path}", self.members[i].label)
} else {
path
});
}
}
}
let left_out = total - out.len();
Ok((out, left_out))
}
pub fn symbols_in(
&self,
only: Option<&[String]>,
name: &str,
limit: usize,
) -> Result<Vec<crate::store::SymbolRow>> {
self.graph_in(only, |store| {
crate::store::symbols_named(store.db(), name, limit)
})
}
pub fn pattern_in(
&self,
only: Option<&[String]>,
language: &str,
source: &str,
filter: &crate::filter::Filter,
) -> Result<crate::pattern::Matches> {
let chosen = self.chosen(only)?;
let label_rows = self.members.len() > 1;
let mut out = crate::pattern::Matches {
language: language.trim().to_ascii_lowercase(),
matches: Vec::new(),
files: 0,
truncated: false,
};
for i in chosen {
let part = crate::pattern::run(self.members[i].store.db(), language, source, filter)?;
out.files += part.files;
out.truncated |= part.truncated;
for mut found in part.matches {
if label_rows {
found.store = Some(self.members[i].label.clone());
}
out.matches.push(found);
}
}
Ok(out)
}
pub fn read_in(
&self,
only: Option<&[String]>,
target: &crate::Target,
filter: &crate::filter::Filter,
) -> Result<Option<crate::Read>> {
let chosen = self.chosen(only)?;
let label_rows = self.members.len() > 1;
for i in chosen {
let Some(found) = self.members[i].store.read(target, filter)? else {
continue;
};
return Ok(Some(match found {
crate::Read::One(mut span) => {
if label_rows {
span.store = Some(self.members[i].label.clone());
}
crate::Read::One(span)
}
crate::Read::Choose(mut rows) => {
if label_rows {
for row in &mut rows {
row.store = Some(self.members[i].label.clone());
}
}
crate::Read::Choose(rows)
}
}));
}
Ok(None)
}
pub fn evidence_in(
&self,
only: Option<&[String]>,
name: &str,
kinds: &[String],
limit: usize,
all: bool,
) -> Result<crate::graph::Evidence> {
let mut merged: Option<crate::graph::Evidence> = None;
for part in self.graph_each(only, |s| {
crate::graph::evidence(s.db(), name, kinds, limit, all)
})? {
match &mut merged {
None => merged = Some(part),
Some(into) => {
into.definitions.extend(part.definitions);
into.callers.extend(part.callers);
into.callees.extend(part.callees);
into.ego.extend(part.ego);
}
}
}
Ok(merged.unwrap_or_else(|| crate::graph::Evidence {
name: name.to_string(),
definitions: Vec::new(),
callers: Vec::new(),
callees: Vec::new(),
ego: Vec::new(),
}))
}
pub fn neighbours_in(
&self,
only: Option<&[String]>,
name: &str,
kinds: &[String],
all: bool,
) -> Result<crate::graph::Neighbours> {
let callers = self.graph_in(only, |s| crate::store::edges_in(s.db(), name, kinds))?;
let callees = self.graph_in(only, |s| crate::store::edges_out(s.db(), name, kinds))?;
let unresolved =
self.graph_in(only, |s| crate::store::unresolved_out(s.db(), name, kinds))?;
Ok(crate::graph::Neighbours {
callers,
callees: if all {
callees
} else {
crate::graph::collapse(callees)
},
hidden: if all { 0 } else { unresolved.len() },
unresolved: if all { unresolved } else { Vec::new() },
})
}
pub fn counter(&self) -> crate::ledger::Counter<'_> {
self.tokenizer
.as_ref()
.or_else(|| self.members.iter().find_map(|m| m.store.tokenizer()))
.map(crate::ledger::Counter::Model)
.unwrap_or(crate::ledger::Counter::Chars4)
}
pub fn path_in(
&self,
only: Option<&[String]>,
from: &str,
to: &str,
depth: u32,
all_edges: bool,
) -> Result<Option<crate::graph::Chain>> {
let mut best: Option<crate::graph::Chain> = None;
for i in self.chosen(only)? {
if let Some(chain) =
crate::graph::shortest_path(self.members[i].store.db(), from, to, depth, all_edges)?
&& best
.as_ref()
.is_none_or(|b| chain.steps.len() < b.steps.len())
{
best = Some(chain);
}
}
Ok(best)
}
fn graph_each<T>(
&self,
only: Option<&[String]>,
read: impl Fn(&Semlith) -> Result<T>,
) -> Result<Vec<T>> {
let chosen = self.chosen(only)?;
let mut out = Vec::new();
for i in chosen {
out.push(read(&self.members[i].store)?);
}
Ok(out)
}
fn graph_in<T: Labelled>(
&self,
only: Option<&[String]>,
read: impl Fn(&Semlith) -> Result<Vec<T>>,
) -> Result<Vec<T>> {
let chosen = self.chosen(only)?;
let label_rows = self.members.len() > 1;
let mut out = Vec::new();
for i in chosen {
let label = self.members[i].label.clone();
for mut row in read(&self.members[i].store)? {
if label_rows {
row.label(&label);
}
out.push(row);
}
}
Ok(out)
}
pub fn writable(&mut self, only: &[String]) -> Result<&mut Semlith> {
let i = match only {
[] if self.members.len() == 1 => 0,
[] => bail!(
"this server has {} stores open, so a write has to name one: {}",
self.members.len(),
self.labels().join(", ")
),
[name] => self
.members
.iter()
.position(|m| m.label == *name)
.with_context(|| {
format!(
"no store called {name} is open; these are: {}",
self.labels().join(", ")
)
})?,
names => bail!(
"a write goes to one store, not {}: {}",
names.len(),
names.join(", ")
),
};
Ok(&mut self.members[i].store)
}
fn chosen(&self, only: Option<&[String]>) -> Result<Vec<usize>> {
let Some(names) = only.filter(|n| !n.is_empty()) else {
return Ok((0..self.members.len()).collect());
};
let mut chosen = Vec::new();
for name in names {
match self.members.iter().position(|m| m.label == *name) {
Some(i) => {
if !chosen.contains(&i) {
chosen.push(i);
}
}
None => bail!(
"no store called {name} is open; these are: {}",
self.labels().join(", ")
),
}
}
Ok(chosen)
}
fn embed_query(&mut self, model: &Model, query: &str) -> Result<Vec<f32>> {
let text = model.query_text(query);
let i = self.embedder(model)?;
let mut out = self.embedders[i]
.1
.embed(vec![text], Some(1))
.map_err(|e| anyhow::anyhow!("{e}"))?;
self.embeds += 1;
let mut vector = out.remove(0);
crate::normalize(&mut vector);
Ok(vector)
}
fn embedder(&mut self, model: &Model) -> Result<usize> {
if let Some(i) = self.embedders.iter().position(|(m, _)| m == model) {
return Ok(i);
}
let cache = model_cache_dir()?;
if self.tokenizer.is_none() {
self.tokenizer = model.tokenizer(&cache);
}
let loaded = model.load(cache, chunk::MAX_CHARS / 2, self.quiet)?;
self.embedders.push((model.clone(), loaded));
Ok(self.embedders.len() - 1)
}
}
fn merge(queues: Vec<Vec<(Hit, f32)>>, k: usize) -> Vec<Hit> {
let mut queues: Vec<VecDeque<(Hit, f32)>> = queues.into_iter().map(VecDeque::from).collect();
let mut out = Vec::new();
while out.len() < k {
let mut best: Option<(usize, f32, f32)> = None;
for (i, queue) in queues.iter().enumerate() {
let Some((hit, similarity)) = queue.front() else {
continue;
};
let better = match best {
None => true,
Some((_, score, sim)) => {
hit.score > score || (hit.score == score && *similarity > sim)
}
};
if better {
best = Some((i, hit.score, *similarity));
}
}
let Some((i, _, _)) = best else { break };
let (hit, _) = queues[i].pop_front().expect("the winner has a head");
out.push(hit);
}
out
}
fn label(dir: &Path) -> String {
let own = dir
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_default();
if (own.is_empty() || own.starts_with('.'))
&& let Some(parent) = dir
.parent()
.and_then(Path::file_name)
.map(|n| n.to_string_lossy().into_owned())
.filter(|n| !n.is_empty())
{
return parent;
}
if own.is_empty() {
dir.display().to_string()
} else {
own
}
}
fn name_for(dir: &Path, named: &[(PathBuf, String)]) -> Option<String> {
let dir = canonical(dir);
named
.iter()
.find(|(at, _)| canonical(at) == dir)
.map(|(_, name)| name.clone())
.filter(|name| !name.is_empty())
}
fn disambiguate(members: &mut [Member], keys: &[PathBuf]) {
let labels: Vec<String> = members.iter().map(|m| m.label.clone()).collect();
for (i, member) in members.iter_mut().enumerate() {
if labels
.iter()
.enumerate()
.any(|(j, l)| j != i && *l == labels[i])
{
member.label = keys[i].display().to_string();
}
}
}
pub trait Labelled {
fn label(&mut self, store: &str);
}
impl Labelled for crate::store::SymbolRow {
fn label(&mut self, store: &str) {
self.store = Some(store.to_string());
}
}
impl Labelled for crate::store::EdgeEnd {
fn label(&mut self, store: &str) {
self.symbol.store = Some(store.to_string());
}
}
impl Labelled for crate::store::Unresolved {
fn label(&mut self, _store: &str) {}
}
#[cfg(test)]
mod tests {
use super::name_for;
#[test]
fn a_daemon_name_replaces_the_label_derived_from_the_directory() {
let dir = std::env::temp_dir()
.join("semlith-fleet-name-test")
.join("store");
let named = vec![(dir.clone(), "work".to_string())];
assert_eq!(name_for(&dir, &named).as_deref(), Some("work"));
let elsewhere = std::env::temp_dir().join("semlith-fleet-name-other");
assert_eq!(name_for(&elsewhere, &named), None);
let blank = vec![(dir.clone(), String::new())];
assert_eq!(name_for(&dir, &blank), None);
}
use super::*;
fn hit(path: &str, score: f32) -> Hit {
Hit {
score,
path: path.into(),
start_line: 1,
end_line: 1,
text: String::new(),
store: None,
lists: vec!["vector"],
image: None,
fresh: true,
symbol: None,
symbol_kind: None,
provenance: None,
}
}
#[test]
fn equal_scores_are_decided_by_similarity_not_by_store_order() {
let unrelated = vec![(hit("bread.md", 0.0328), 0.31)];
let answer = vec![(hit("ownership.md", 0.0328), 0.62)];
let out = merge(vec![unrelated, answer], 5);
assert_eq!(out[0].path, "ownership.md");
assert_eq!(out.len(), 2, "nothing may be dropped, only reordered");
}
#[test]
fn a_stores_own_order_survives_the_merge() {
let one = vec![
(hit("first.md", 0.0328), 0.10),
(hit("second.md", 0.0328), 0.90),
];
let out = merge(vec![one], 5);
assert_eq!(
out.iter().map(|h| h.path.as_str()).collect::<Vec<_>>(),
vec!["first.md", "second.md"]
);
}
#[test]
fn k_bounds_the_merged_list_not_each_store() {
let a = vec![(hit("a1", 0.9), 0.9), (hit("a2", 0.8), 0.8)];
let b = vec![(hit("b1", 0.7), 0.7), (hit("b2", 0.6), 0.6)];
let out = merge(vec![a, b], 3);
assert_eq!(
out.iter().map(|h| h.path.as_str()).collect::<Vec<_>>(),
vec!["a1", "a2", "b1"]
);
}
#[test]
fn a_store_is_named_after_the_directory_holding_it() {
assert_eq!(label(Path::new("/home/dev/api/.semlith")), "api");
assert_eq!(label(Path::new("/home/dev/stores/api-store")), "api-store");
assert_eq!(label(Path::new("/.semlith")), ".semlith");
}
#[test]
fn stores_that_would_share_a_label_get_their_paths_instead() {
let keys = [
PathBuf::from("/a/api/.semlith"),
PathBuf::from("/b/api/.semlith"),
PathBuf::from("/b/cli/.semlith"),
];
let mut labels: Vec<String> = keys.iter().map(|k| label(k)).collect();
assert_eq!(labels, vec!["api", "api", "cli"]);
let original = labels.clone();
for i in 0..labels.len() {
if original
.iter()
.enumerate()
.any(|(j, l)| j != i && *l == original[i])
{
labels[i] = keys[i].display().to_string();
}
}
assert_eq!(
labels,
vec!["/a/api/.semlith", "/b/api/.semlith", "cli"],
"a label that names two stores is worse than a long one"
);
}
}