use crate::Semlith;
use crate::fleet::Fleet;
use crate::store;
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
pub const CHARS4: &str = "chars4";
pub const OFF_ENV: &str = "SEMLITH_LEDGER";
pub fn enabled() -> bool {
off_reason().is_none()
}
pub fn env_on() -> bool {
!matches!(
std::env::var(OFF_ENV).as_deref(),
Ok("0") | Ok("off") | Ok("false")
)
}
static SESSION_OFF: AtomicBool = AtomicBool::new(false);
static PAUSED: AtomicBool = AtomicBool::new(false);
pub fn set_session_off(off: bool) {
SESSION_OFF.store(off, Ordering::Relaxed);
}
pub fn set_paused(paused: bool) {
PAUSED.store(paused, Ordering::Relaxed);
}
pub fn off_reason() -> Option<&'static str> {
if !env_on() {
Some("env")
} else if SESSION_OFF.load(Ordering::Relaxed) {
Some("flag")
} else if PAUSED.load(Ordering::Relaxed) {
Some("paused")
} else {
None
}
}
static UNRECORDED: std::sync::RwLock<Vec<PathBuf>> = std::sync::RwLock::new(Vec::new());
pub fn set_unrecorded(dirs: Vec<PathBuf>) {
*UNRECORDED.write().unwrap_or_else(|e| e.into_inner()) =
dirs.iter().map(|d| crate::canonical(d)).collect();
}
fn write(db: &rusqlite::Connection, row: &store::NewRetrieval<'_>) -> anyhow::Result<()> {
let off = UNRECORDED.read().unwrap_or_else(|e| e.into_inner());
if !off.is_empty()
&& let Some(dir) = db.path().and_then(|p| Path::new(p).parent())
&& off.contains(&crate::canonical(dir))
{
return Ok(());
}
drop(off);
store::record_counted(db, row, store::BASIS_TEXT)
}
pub fn recording_state() -> serde_json::Value {
let reason = off_reason();
serde_json::json!({ "on": reason.is_none(), "reason": reason })
}
pub fn estimate(text: &str) -> i64 {
text.len().div_ceil(4) as i64
}
#[derive(Clone, Copy)]
pub enum Counter<'a> {
Model(&'a tokenizers::Tokenizer),
Chars4,
}
impl Counter<'_> {
pub fn count(&self, text: &str) -> i64 {
match self {
Counter::Model(tokenizer) => match tokenizer.encode(text, false) {
Ok(encoded) => encoded.len() as i64,
Err(_) => estimate(text),
},
Counter::Chars4 => estimate(text),
}
}
pub fn count_bytes(&self, bytes: u64) -> i64 {
bytes.div_ceil(4) as i64
}
pub fn label(&self) -> &'static str {
match self {
Counter::Model(_) => MODEL,
Counter::Chars4 => CHARS4,
}
}
}
pub const MODEL: &str = "model";
fn query_id() -> String {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
let at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_nanos())
.unwrap_or(0);
format!("{at:x}-{:x}", SEQ.fetch_add(1, Ordering::Relaxed))
}
fn whole_of(db: &rusqlite::Connection, counter: &Counter<'_>, path: &str) -> i64 {
let p = Path::new(path);
let ext = p
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_ascii_lowercase();
if crate::image::is_image(p) {
return 0;
}
if ext == "pdf" || crate::formats::handles(&ext) {
if let Ok(Some(text)) = store::held_text(db, path) {
return counter.count(&text);
}
}
std::fs::metadata(p)
.map(|m| counter.count_bytes(m.len()))
.unwrap_or(0)
}
pub fn search(
fleet: &Fleet,
who: &Who<'_>,
query: &str,
hits: &[crate::Hit],
elapsed: std::time::Duration,
) {
let counter = fleet.counter();
if !enabled() {
return;
}
let id = query_id();
let labelled = hits.iter().any(|h| h.store.is_some());
let mut first = true;
for (label, store) in fleet.each() {
let mine: Vec<&crate::Hit> = if labelled {
hits.iter()
.filter(|h| h.store.as_deref() == Some(label))
.collect()
} else if first {
hits.iter().collect()
} else {
Vec::new()
};
first = false;
if mine.is_empty() && !hits.is_empty() {
continue;
}
let excerpt: i64 = mine.iter().map(|h| counter.count(&h.text)).sum();
let paths: std::collections::BTreeSet<&str> =
mine.iter().map(|h| h.path.as_str()).collect();
let whole: i64 = paths
.iter()
.map(|p| whole_of(store.db(), &counter, p))
.sum();
let stale = mine.iter().filter(|h| !h.fresh).count() as i64;
let _ = write(
store.db(),
&store::NewRetrieval {
client: who.client,
session: who.session,
client_version: who.version,
tool: "search",
query,
hits: mine.len() as i64,
micros: elapsed.as_micros() as i64,
excerpt_tokens: excerpt,
whole_file_tokens: whole,
stale_hits: stale,
tokenizer: counter.label(),
query_id: &id,
},
);
if hits.is_empty() {
break;
}
}
}
pub fn reply(
fleet: &Fleet,
who: &Who<'_>,
tool: &str,
query: &str,
body: &str,
scope: &[String],
elapsed: std::time::Duration,
) {
if !enabled() {
return;
}
let counter = fleet.counter();
let id = query_id();
let write = |db: &rusqlite::Connection, paths: &BTreeSet<String>, text: &str| {
let whole: i64 = paths.iter().map(|p| whole_of(db, &counter, p)).sum();
let _ = write(
db,
&store::NewRetrieval {
client: who.client,
session: who.session,
client_version: who.version,
tool,
query,
hits: paths.len() as i64,
micros: elapsed.as_micros() as i64,
excerpt_tokens: counter.count(text),
whole_file_tokens: whole,
stale_hits: text.matches("ยท stale").count() as i64,
tokenizer: counter.label(),
query_id: &id,
},
);
};
let members: Vec<(&str, &Semlith)> = fleet.each().collect();
let parts = shares(&members, &fleet.roots(), body);
if parts.is_empty() {
let searched = fleet
.selected(Some(scope))
.ok()
.and_then(|chosen| chosen.into_iter().next())
.or_else(|| members.first().copied());
if let Some((_, first)) = searched {
write(first.db(), &BTreeSet::new(), body);
}
return;
}
for share in &parts {
write(members[share.member].1.db(), &share.paths, &share.text);
}
}
#[derive(Debug)]
struct Share {
member: usize,
paths: BTreeSet<String>,
text: String,
}
fn shares(members: &[(&str, &Semlith)], roots: &[PathBuf], body: &str) -> Vec<Share> {
let mut out: Vec<Share> = Vec::new();
let mut current: Option<usize> = None;
let mut preamble = String::new();
for line in body.lines() {
if let Some((member, path)) = heading(members, roots, line) {
let at = match out.iter().position(|s| s.member == member) {
Some(at) => at,
None => {
out.push(Share {
member,
paths: BTreeSet::new(),
text: String::new(),
});
out.len() - 1
}
};
out[at].paths.insert(path);
current = Some(at);
}
let text = match current {
Some(at) => &mut out[at].text,
None => &mut preamble,
};
text.push_str(line);
text.push('\n');
}
if let Some(first) = out.first_mut() {
first.text.insert_str(0, &preamble);
}
out
}
fn heading(members: &[(&str, &Semlith)], roots: &[PathBuf], line: &str) -> Option<(usize, String)> {
if members.is_empty() || line.starts_with(char::is_whitespace) {
return None;
}
let mut rest = line.trim_end();
if let Some((n, tail)) = rest.strip_prefix('[').and_then(|r| r.split_once("] "))
&& !n.is_empty()
&& n.bytes().all(|b| b.is_ascii_digit())
{
rest = tail;
}
let mut named = None;
for (i, (label, _)) in members.iter().enumerate() {
let tail = rest
.strip_prefix('[')
.and_then(|r| r.strip_prefix(*label))
.and_then(|r| r.strip_prefix("] "))
.or_else(|| rest.strip_prefix(*label).and_then(|r| r.strip_prefix(' ')));
if let Some(tail) = tail {
named = Some(i);
rest = tail;
break;
}
}
let path = match rest
.char_indices()
.find(|&(i, c)| c == ':' && rest[i + 1..].starts_with(|d: char| d.is_ascii_digit()))
{
Some((i, _)) => &rest[..i],
None => rest.split(" \u{2014} image").next().unwrap_or(rest),
};
if path.is_empty() || (named.is_none() && path.contains(char::is_whitespace)) {
return None;
}
let given = Path::new(path);
let candidates: Vec<PathBuf> = if given.is_absolute() {
vec![given.to_path_buf()]
} else {
roots
.iter()
.map(|r| r.join(given))
.chain([given.to_path_buf()])
.collect()
};
let files: Vec<String> = candidates
.iter()
.filter(|c| c.is_file())
.map(|c| crate::canonical(c).to_string_lossy().into_owned())
.collect();
let held = |m: usize, file: &str| {
store::file_hash(members[m].1.db(), file)
.ok()
.flatten()
.is_some()
};
let (member, file) = match named {
Some(m) => (
m,
files
.iter()
.find(|f| held(m, f.as_str()))
.or(files.first())?,
),
None => files
.iter()
.find_map(|f| {
(0..members.len())
.find(|&m| held(m, f.as_str()))
.map(|m| (m, f))
})
.or_else(|| files.first().map(|f| (0, f)))?,
};
Some((member, file.clone()))
}
pub fn graph(
fleet: &Fleet,
who: &Who<'_>,
tool: &str,
name: &str,
reply: &str,
found: bool,
elapsed: std::time::Duration,
) {
if !enabled() {
return;
}
let counter = fleet.counter();
let members: Vec<(&str, &Semlith)> = fleet.each().collect();
let costs: Vec<i64> = members
.iter()
.map(|(_, store)| match store::grep_cost(store.db(), name) {
Ok(bytes) => counter.count_bytes(bytes as u64),
Err(_) => 0,
})
.collect();
let at = costs.iter().position(|c| *c > 0).unwrap_or(0);
let Some((_, store)) = members.get(at) else {
return;
};
let whole: i64 = costs.iter().sum();
let _ = write(
store.db(),
&store::NewRetrieval {
client: who.client,
session: who.session,
client_version: who.version,
tool,
query: name,
hits: i64::from(found && whole > 0),
micros: elapsed.as_micros() as i64,
excerpt_tokens: counter.count(reply),
whole_file_tokens: whole,
stale_hits: 0,
tokenizer: counter.label(),
query_id: &query_id(),
},
);
}
pub fn refusal(fleet: &Fleet, who: &Who<'_>, tool: &str) {
if !enabled() {
return;
}
let Some((_, store)) = fleet.each().next() else {
return;
};
let _ = write(
store.db(),
&store::NewRetrieval {
client: who.client,
session: who.session,
client_version: who.version,
tool: store::REFUSED_TOOL,
query: tool,
hits: 0,
micros: 0,
excerpt_tokens: 0,
whole_file_tokens: 0,
stale_hits: 0,
tokenizer: fleet.counter().label(),
query_id: &query_id(),
},
);
}
pub fn brief(
fleet: &Fleet,
who: &Who<'_>,
question: &str,
brief: &crate::brief::Brief,
body: &str,
elapsed: std::time::Duration,
) {
if !enabled() {
return;
}
let counter = fleet.counter();
let paths: std::collections::BTreeSet<&str> =
brief.spans.iter().map(|s| s.path.as_str()).collect();
if let Some((_, store)) = fleet.each().next() {
let whole: i64 = paths
.iter()
.map(|p| whole_of(store.db(), &counter, p))
.sum();
let _ = write(
store.db(),
&store::NewRetrieval {
client: who.client,
session: who.session,
client_version: who.version,
tool: "brief",
query: question,
hits: paths.len() as i64,
micros: elapsed.as_micros() as i64,
excerpt_tokens: counter.count(body),
whole_file_tokens: whole,
stale_hits: 0,
tokenizer: counter.label(),
query_id: &query_id(),
},
);
}
}
pub const RAW_READ: &str = "raw-read";
pub fn raw_read(fleet: &Fleet, client: &str, session: &str, path: &str) -> bool {
if !enabled() {
return false;
}
let counter = fleet.counter();
for (_, store) in fleet.each() {
if !store::holds_path(store.db(), path).unwrap_or(false) {
continue;
}
let canonical = crate::canonical(Path::new(path));
let whole = whole_of(store.db(), &counter, &canonical.to_string_lossy());
let recorded = write(
store.db(),
&store::NewRetrieval {
client_version: "",
client,
session,
tool: RAW_READ,
query: path,
hits: 0,
micros: 0,
excerpt_tokens: 0,
whole_file_tokens: whole,
stale_hits: 0,
tokenizer: counter.label(),
query_id: &query_id(),
},
);
return recorded.is_ok();
}
false
}
pub fn acceptance(
db: &rusqlite::Connection,
a: &store::Acceptance,
action: &str,
) -> anyhow::Result<()> {
if !enabled() {
return Ok(());
}
let query = serde_json::json!({
"path": crate::plain(&a.path),
"class": a.class,
"mode": a.mode,
"confidence": a.confidence,
})
.to_string();
store::record_retrieval(
db,
&store::NewRetrieval {
client_version: "",
client: &a.source,
session: "",
tool: action,
query: &query,
hits: 0,
micros: 0,
excerpt_tokens: 0,
whole_file_tokens: 0,
stale_hits: 0,
tokenizer: CHARS4,
query_id: &query_id(),
},
)
}
#[derive(Debug, Clone, Copy)]
pub struct Who<'a> {
pub client: &'a str,
pub session: &'a str,
pub version: &'a str,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Window {
pub since: i64,
pub until: Option<i64>,
}
impl Window {
pub const ALL: Window = Window {
since: 0,
until: None,
};
}
pub fn is_retrieval(tool: &str) -> bool {
!matches!(
tool,
RAW_READ | store::NOTE_TOOL | store::REFUSED_TOOL | "accept" | "revoke"
)
}
pub fn is_human(client: &str) -> bool {
matches!(client, "cli" | "portal")
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct SessionFigures {
pub session: String,
pub client: String,
pub client_version: String,
pub first: i64,
pub last: i64,
pub queries: i64,
pub zero_hit: i64,
pub excerpt_tokens: i64,
pub whole_file_tokens: i64,
pub net_tokens: i64,
pub model: Option<String>,
pub stores: Vec<String>,
pub excerpt_counter: &'static str,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct StoreFigures {
pub store: String,
pub queries: i64,
pub credited: i64,
pub zero_hit: i64,
pub excerpt_tokens: i64,
pub whole_file_tokens: i64,
pub net_tokens: i64,
pub refunds: i64,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct Tier {
pub excerpt: &'static str,
pub whole: &'static str,
pub word: &'static str,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct ClientFigures {
pub client: String,
pub version: String,
pub queries: i64,
pub credited: i64,
pub net_tokens: i64,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct GapFigures {
pub query: String,
pub asked: i64,
pub last: i64,
pub stores: Vec<String>,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct HumanFigures {
pub queries: i64,
pub net_tokens: i64,
}
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct Figures {
pub window: (i64, Option<i64>),
pub queries: i64,
pub credited: i64,
pub zero_hit: i64,
pub coverage: i64,
pub excerpt_tokens: i64,
pub whole_file_tokens: i64,
pub net_tokens: i64,
pub ratio: Option<f64>,
pub clients: Vec<String>,
pub by_client: std::collections::BTreeMap<String, i64>,
pub client_figures: Vec<ClientFigures>,
pub sessions: i64,
pub refunds: i64,
pub timed: i64,
pub p50_micros: i64,
pub p95_micros: i64,
pub tier: Tier,
pub human: HumanFigures,
pub stores: Vec<StoreFigures>,
pub per_session: Vec<SessionFigures>,
pub gaps: Vec<GapFigures>,
}
pub fn figures(
members: &[(&str, &rusqlite::Connection)],
window: Window,
humans: bool,
) -> anyhow::Result<Figures> {
use std::collections::{BTreeMap, BTreeSet};
let mut all: Vec<(String, store::LedgerFact)> = Vec::new();
for (label, db) in members {
for fact in store::ledger_facts(db, window.since, window.until)? {
all.push((label.to_string(), fact));
}
}
let key = |label: &str, f: &store::LedgerFact| {
if f.query_id.is_empty() {
format!("{label}:row:{}", f.id)
} else {
f.query_id.clone()
}
};
let mut out = Figures {
window: (window.since, window.until),
..Default::default()
};
let mut keys: BTreeSet<String> = BTreeSet::new();
let mut credited: BTreeSet<String> = BTreeSet::new();
let mut human_keys: BTreeSet<String> = BTreeSet::new();
let mut by_client: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
let mut micros: BTreeMap<String, i64> = BTreeMap::new();
let mut excerpt_counters: BTreeSet<&'static str> = BTreeSet::new();
let mut whole_bases: BTreeSet<&'static str> = BTreeSet::new();
let mut stores: BTreeMap<String, (BTreeSet<String>, BTreeSet<String>, StoreFigures)> =
BTreeMap::new();
type Conversation = (
SessionFigures,
BTreeSet<String>,
BTreeSet<String>,
BTreeMap<String, i64>,
);
let mut sessions: BTreeMap<(String, String), Conversation> = BTreeMap::new();
let mut questions: BTreeMap<String, (String, bool, i64, BTreeSet<String>)> = BTreeMap::new();
type Client = (String, BTreeSet<String>, BTreeSet<String>, i64);
let mut clients: BTreeMap<String, Client> = BTreeMap::new();
for (label, _) in members {
stores.entry(label.to_string()).or_default().2.store = label.to_string();
}
for (label, f) in &all {
if f.tool == RAW_READ {
out.refunds += 1;
stores.entry(label.clone()).or_default().2.refunds += 1;
continue;
}
if !is_retrieval(&f.tool) {
continue;
}
let k = key(label, f);
if is_human(&f.client) && !humans {
human_keys.insert(k);
if f.hits > 0 {
out.human.net_tokens += f.whole_file_tokens - f.excerpt_tokens;
}
continue;
}
keys.insert(k.clone());
by_client
.entry(f.client.clone())
.or_default()
.insert(k.clone());
if f.micros > 0 {
let m = micros.entry(k.clone()).or_default();
*m = (*m).max(f.micros);
}
let q = questions
.entry(k.clone())
.or_insert_with(|| (f.query.clone(), false, 0, BTreeSet::new()));
q.1 |= f.hits > 0;
q.2 = q.2.max(f.at);
q.3.insert(label.clone());
let c = clients.entry(f.client.clone()).or_default();
if !f.client_version.is_empty() {
c.0 = f.client_version.clone();
}
c.1.insert(k.clone());
if f.hits > 0 {
c.2.insert(k.clone());
c.3 += f.whole_file_tokens - f.excerpt_tokens;
}
let st = stores.entry(label.clone()).or_default();
st.0.insert(k.clone());
let conversation = sessions
.entry((f.client.clone(), f.session.clone()))
.or_insert_with(|| {
(
SessionFigures {
session: f.session.clone(),
client: f.client.clone(),
first: f.at,
..Default::default()
},
BTreeSet::new(),
BTreeSet::new(),
BTreeMap::new(),
)
});
conversation.0.first = conversation.0.first.min(f.at);
conversation.0.last = conversation.0.last.max(f.at);
if !f.client_version.is_empty() {
conversation.0.client_version = f.client_version.clone();
}
conversation.1.insert(k.clone());
if !conversation.0.stores.contains(label) {
conversation.0.stores.push(label.clone());
}
if let Some(model) = f.usage.as_ref().and_then(|u| u.model.clone()) {
*conversation.3.entry(model).or_default() += 1;
}
if f.hits > 0 {
credited.insert(k.clone());
st.1.insert(k.clone());
conversation.2.insert(k.clone());
out.excerpt_tokens += f.excerpt_tokens;
out.whole_file_tokens += f.whole_file_tokens;
out.net_tokens += f.whole_file_tokens - f.excerpt_tokens;
st.2.excerpt_tokens += f.excerpt_tokens;
st.2.whole_file_tokens += f.whole_file_tokens;
st.2.net_tokens += f.whole_file_tokens - f.excerpt_tokens;
conversation.0.excerpt_tokens += f.excerpt_tokens;
conversation.0.whole_file_tokens += f.whole_file_tokens;
conversation.0.net_tokens += f.whole_file_tokens - f.excerpt_tokens;
excerpt_counters.insert(if f.tokenizer == MODEL { MODEL } else { CHARS4 });
whole_bases.insert(if f.whole_basis == store::BASIS_TEXT {
store::BASIS_TEXT
} else {
store::BASIS_BYTES
});
conversation.0.excerpt_counter = if f.tokenizer == MODEL { MODEL } else { CHARS4 };
}
}
out.queries = keys.len() as i64;
out.credited = credited.len() as i64;
out.zero_hit = out.queries - out.credited;
out.coverage = percent(out.credited, out.queries);
out.ratio =
(out.excerpt_tokens > 0).then(|| out.whole_file_tokens as f64 / out.excerpt_tokens as f64);
out.clients = by_client.keys().cloned().collect();
out.by_client = by_client
.into_iter()
.map(|(c, k)| (c, k.len() as i64))
.collect();
out.sessions = sessions.len() as i64;
let mut spent: Vec<i64> = micros.into_values().collect();
spent.sort_unstable();
out.timed = spent.len() as i64;
out.client_figures = clients
.into_iter()
.map(|(client, (version, keys, credited, net))| ClientFigures {
client,
version,
queries: keys.len() as i64,
credited: credited.len() as i64,
net_tokens: net,
})
.collect();
out.p50_micros = quantile(&spent, 0.50);
out.p95_micros = quantile(&spent, 0.95);
let one = |set: &BTreeSet<&'static str>, none: &'static str| match set.len() {
0 => none,
1 => *set.iter().next().unwrap(),
_ => "mixed",
};
out.tier = Tier {
excerpt: one(&excerpt_counters, MODEL),
whole: one(&whole_bases, store::BASIS_TEXT),
word: if out.credited > 0 && excerpt_counters.iter().all(|c| *c == MODEL) {
"measured"
} else {
"modelled"
},
};
out.human.queries = human_keys.len() as i64;
out.stores = stores
.into_values()
.map(|(keys, credited, mut st)| {
st.queries = keys.len() as i64;
st.credited = credited.len() as i64;
st.zero_hit = st.queries - st.credited;
st
})
.collect();
let mut per_session: Vec<SessionFigures> = sessions
.into_values()
.map(|(mut s, keys, credited, models)| {
s.queries = keys.len() as i64;
s.zero_hit = (keys.len() - credited.len()) as i64;
s.model = models
.into_iter()
.max_by_key(|(m, n)| (*n, std::cmp::Reverse(m.clone())))
.map(|(m, _)| m);
if s.excerpt_counter.is_empty() {
s.excerpt_counter = MODEL;
}
s
})
.collect();
per_session.sort_by_key(|s| std::cmp::Reverse(s.last));
out.per_session = per_session;
let mut gaps: BTreeMap<String, GapFigures> = BTreeMap::new();
for (query, found, last, asked_of) in questions.into_values() {
if found {
continue;
}
let gap = gaps.entry(query.clone()).or_insert_with(|| GapFigures {
query,
..Default::default()
});
gap.asked += 1;
gap.last = gap.last.max(last);
for store in asked_of {
if !gap.stores.contains(&store) {
gap.stores.push(store);
}
}
}
let mut gaps: Vec<GapFigures> = gaps.into_values().collect();
gaps.sort_by(|a, b| b.asked.cmp(&a.asked).then_with(|| a.query.cmp(&b.query)));
out.gaps = gaps;
Ok(out)
}
pub fn percent(part: i64, whole: i64) -> i64 {
if whole <= 0 {
return 0;
}
(part * 200 + whole) / (whole * 2)
}
fn quantile(sorted: &[i64], q: f64) -> i64 {
if sorted.is_empty() {
return 0;
}
sorted[((sorted.len() as f64 - 1.0) * q).round() as usize]
}
#[cfg(test)]
mod tests {
use super::*;
fn store_holding(dir: &Path, files: &[PathBuf]) -> Semlith {
std::fs::create_dir_all(dir).unwrap();
let s = Semlith::open(dir, None).unwrap();
store::read_only(s.db(), false).unwrap();
for f in files {
s.db()
.execute(
"INSERT INTO files (path, hash, bytes, indexed_at) VALUES (?1, 'h', 1, 0)",
[crate::canonical(f).to_string_lossy().into_owned()],
)
.unwrap();
}
store::read_only(s.db(), true).unwrap();
s
}
fn canon(p: &Path) -> String {
crate::canonical(p).to_string_lossy().into_owned()
}
fn two_roots(tmp: &Path) -> (PathBuf, PathBuf) {
let (a, b) = (tmp.join("a"), tmp.join("b"));
for root in [&a, &b] {
std::fs::create_dir_all(root.join("src")).unwrap();
std::fs::write(root.join("src").join("lib.rs"), "fn x() {}\n").unwrap();
}
std::fs::write(a.join("notes.md"), "# notes\n").unwrap();
(a, b)
}
#[test]
fn a_fleet_reply_is_split_by_the_store_each_heading_names() {
let tmp = tempfile::tempdir().unwrap();
let (a, b) = two_roots(tmp.path());
let alpha = store_holding(
&tmp.path().join("alpha"),
&[a.join("src").join("lib.rs"), a.join("notes.md")],
);
let beta = store_holding(&tmp.path().join("beta"), &[b.join("src").join("lib.rs")]);
let members = [("alpha", &alpha), ("beta", &beta)];
let roots = [a.clone(), b.clone()];
let body = "roots a, b\n\
identifier ยท keyword\n\
beta src/lib.rs\n\
\x20 1-1 x fn ยท stale | fn x() {}\n\
alpha src/lib.rs\n\
\x20 1-1 x fn | fn x() {}\n\
alpha notes.md\n\
\x20 1-1 | # notes";
let got = shares(&members, &roots, body);
assert_eq!(got.len(), 2, "{got:?}");
assert_eq!(got[0].member, 1, "beta answered first: {got:?}");
assert_eq!(
got[0].paths,
BTreeSet::from([canon(&b.join("src").join("lib.rs"))]),
"beta's src/lib.rs is the one under beta's root"
);
assert!(got[0].text.contains("ยท stale"));
assert!(
got[0].text.starts_with("roots a, b\n"),
"the lines above the first heading go to the first store: {:?}",
got[0].text
);
assert_eq!(got[1].member, 0);
assert_eq!(
got[1].paths,
BTreeSet::from([
canon(&a.join("src").join("lib.rs")),
canon(&a.join("notes.md"))
])
);
assert!(!got[1].text.contains("stale"));
assert_eq!(
got[0].text.len() + got[1].text.len(),
body.len() + 1,
"every line of the reply is in exactly one share"
);
}
#[test]
fn every_reply_shape_names_its_store_and_prose_does_not() {
let tmp = tempfile::tempdir().unwrap();
let (a, b) = two_roots(tmp.path());
let alpha = store_holding(
&tmp.path().join("alpha"),
&[a.join("src").join("lib.rs"), a.join("notes.md")],
);
let beta = store_holding(&tmp.path().join("beta"), &[b.join("src").join("lib.rs")]);
let members = [("alpha", &alpha), ("beta", &beta)];
let roots = [a.clone(), b.clone()];
let abs = crate::plain(&a.join("notes.md").to_string_lossy());
let body = format!(
"[beta] src/lib.rs:1-1 x @1 [vector]\n\
\x20 fn x() {{}}\n\
[1] alpha {abs}:1-1 (score 0.500 via vector)\n\
# notes\n\
see src/lib.rs for the rest"
);
let got = shares(&members, &roots, &body);
let by = |m: usize| got.iter().find(|s| s.member == m).map(|s| s.paths.clone());
assert_eq!(
by(1),
Some(BTreeSet::from([canon(&b.join("src").join("lib.rs"))]))
);
assert_eq!(by(0), Some(BTreeSet::from([canon(&a.join("notes.md"))])));
let one = [("alpha", &alpha)];
let got = shares(
&one,
std::slice::from_ref(&a),
"src/lib.rs\n 1-1 | x\nnotes.md:1-1 notes\n",
);
assert_eq!(got.len(), 1);
assert_eq!(got[0].paths.len(), 2, "{got:?}");
assert!(shares(&members, &roots, "No match for that query.").is_empty());
}
fn row<'a>(
query: &'a str,
qid: &'a str,
client: &'a str,
hits: i64,
) -> store::NewRetrieval<'a> {
store::NewRetrieval {
client,
client_version: "2.1.299",
session: "s1",
tool: "search",
query,
hits,
micros: 1000,
excerpt_tokens: 10 * hits,
whole_file_tokens: 100 * hits,
stale_hits: 0,
tokenizer: MODEL,
query_id: qid,
}
}
fn sql(db: &rusqlite::Connection, statement: &str) {
store::read_only(db, false).unwrap();
db.execute(statement, []).unwrap();
}
fn writable(s: &Semlith) -> &rusqlite::Connection {
store::read_only(s.db(), false).unwrap();
s.db()
}
#[test]
fn chain_v2_catches_a_query_id_rewrite_and_a_cut_tail() {
let tmp = tempfile::tempdir().unwrap();
let s = store_holding(&tmp.path().join("a"), &[]);
let db = writable(&s);
for (q, id) in [("a", "q1"), ("b", "q2"), ("c", "q3")] {
store::record_counted(db, &row(q, id, "Claude Code", 1), store::BASIS_TEXT).unwrap();
}
assert_eq!(store::ledger_break(db).unwrap(), None);
sql(db, "UPDATE retrievals SET query_id = 'same'");
assert!(
store::ledger_break(db).unwrap().is_some(),
"a query id rewrite verified"
);
sql(db, "UPDATE retrievals SET query_id = 'q' || id");
assert_eq!(store::ledger_break(db).unwrap(), None);
sql(
db,
"UPDATE retrievals SET whole_basis = 'bytes' WHERE id = 1",
);
assert_eq!(
store::ledger_break(db).unwrap(),
Some(1),
"the basis is covered"
);
sql(
db,
"UPDATE retrievals SET whole_basis = 'text' WHERE id = 1",
);
sql(db, "DELETE FROM retrievals WHERE id = 3");
assert_eq!(
store::ledger_break(db).unwrap(),
Some(3),
"deleting the newest row verified"
);
assert_eq!(store::ledger_repair(db).unwrap(), vec![3]);
assert_eq!(store::ledger_break(db).unwrap(), None);
}
#[test]
fn chain_v2_seals_usage_and_sync() {
let tmp = tempfile::tempdir().unwrap();
let s = store_holding(&tmp.path().join("a"), &[]);
let db = writable(&s);
store::record_counted(db, &row("a", "q1", "Claude Code", 1), store::BASIS_TEXT).unwrap();
let unfilled = store::unfilled(db, 0).unwrap();
store::fill_usage(
db,
&unfilled[0],
&store::Filled {
model: Some("claude-opus-5-5".into()),
cost_usd: Some(0.5),
usage_source: "test log".into(),
..Default::default()
},
)
.unwrap();
store::mark_synced(db, &[1], 99).unwrap();
assert_eq!(store::ledger_break(db).unwrap(), None);
sql(db, "UPDATE retrievals SET cost_usd = 0.01");
assert_eq!(
store::ledger_break(db).unwrap(),
Some(1),
"an edited cost verified"
);
}
#[test]
fn v1_rows_keep_verifying_beside_v2() {
let tmp = tempfile::tempdir().unwrap();
let s = store_holding(&tmp.path().join("a"), &[]);
let db = writable(&s);
store::record_legacy_retrieval(db, "cli", "old", 1, 10).unwrap();
store::record_v1_retrieval(db, &row("v1", "q1", "Zed", 1)).unwrap();
store::record_counted(db, &row("v2", "q2", "Zed", 1), store::BASIS_TEXT).unwrap();
store::record_v1_retrieval(db, &row("rolled back", "q3", "Zed", 1)).unwrap();
assert_eq!(store::ledger_break(db).unwrap(), None);
}
#[test]
fn figures_are_the_rows_recounted() {
let tmp = tempfile::tempdir().unwrap();
let a = store_holding(&tmp.path().join("a"), &[]);
let b = store_holding(&tmp.path().join("b"), &[]);
let (da, db) = (writable(&a), writable(&b));
let put = |db: &rusqlite::Connection, r: store::NewRetrieval<'_>| {
store::record_counted(db, &r, store::BASIS_TEXT).unwrap()
};
put(da, row("both", "q1", "Claude Code", 2));
put(db, row("both", "q1", "Claude Code", 1));
put(
da,
store::NewRetrieval {
session: "s2",
..row("none", "q2", "Cursor", 0)
},
);
put(
db,
store::NewRetrieval {
tool: "neighbors",
excerpt_tokens: 500,
..row("tiny", "q3", "Claude Code", 1)
},
);
put(
da,
store::NewRetrieval {
tool: RAW_READ,
..row("/x.rs", "q4", "Claude Code", 0)
},
);
put(
da,
store::NewRetrieval {
tool: RAW_READ,
..row("/y.rs", "q5", "Claude Code", 0)
},
);
put(
da,
store::NewRetrieval {
tool: "accept",
..row("{}", "q6", "portal", 0)
},
);
store::ledger_note(da, 1).unwrap();
put(db, row("by hand", "q7", "portal", 3));
let f = figures(&[("a", da), ("b", db)], Window::ALL, false).unwrap();
assert_eq!(f.queries, 3, "{f:#?}");
assert_eq!(f.credited, 2);
assert_eq!(f.zero_hit, 1);
assert_eq!(f.coverage, 67);
assert_eq!(f.refunds, 2);
assert_eq!(f.excerpt_tokens, 20 + 10 + 500);
assert_eq!(f.whole_file_tokens, 200 + 100 + 100);
assert_eq!(
f.net_tokens,
400 - 530,
"a negative net is summed, not clamped"
);
assert_eq!(f.clients, vec!["Claude Code", "Cursor"]);
assert_eq!(f.by_client["Claude Code"], 2, "retrievals, not rows");
assert_eq!(f.sessions, 2);
assert_eq!(f.human.queries, 1);
assert_eq!(f.human.net_tokens, 300 - 30);
let gaps: Vec<(&str, i64)> = f.gaps.iter().map(|g| (g.query.as_str(), g.asked)).collect();
assert_eq!(gaps, vec![("none", 1)]);
assert_eq!(f.p50_micros, 1000);
assert_eq!(f.tier.excerpt, MODEL);
assert_eq!(f.tier.whole, store::BASIS_TEXT);
let per: Vec<(String, i64)> = f
.stores
.iter()
.map(|s| (s.store.clone(), s.queries))
.collect();
assert_eq!(
per,
vec![("a".into(), 2), ("b".into(), 2)],
"per-store attribution kept"
);
let with = figures(&[("a", da), ("b", db)], Window::ALL, true).unwrap();
assert_eq!(with.queries, 4);
assert_eq!(percent(2, 3) + percent(1, 3), 100);
}
#[test]
fn a_fleet_reply_writes_a_row_per_store_that_answered() {
let tmp = tempfile::tempdir().unwrap();
let (a, b) = two_roots(tmp.path());
let (alpha_dir, beta_dir) = (tmp.path().join("alpha"), tmp.path().join("beta"));
drop(store_holding(&alpha_dir, &[a.join("notes.md")]));
drop(store_holding(&beta_dir, &[b.join("src").join("lib.rs")]));
let fleet = Fleet::open(&[alpha_dir.clone(), beta_dir.clone()]).unwrap();
assert_eq!(fleet.labels(), vec!["alpha", "beta"]);
let body = format!(
"alpha {}\n 1-1 | # notes\nbeta {}\n 1-1 x fn ยท stale | fn x() {{}}",
crate::plain(&a.join("notes.md").to_string_lossy()),
crate::plain(&b.join("src").join("lib.rs").to_string_lossy()),
);
let who = Who {
client: "harness",
session: "s",
version: "",
};
reply(
&fleet,
&who,
"search",
"notes",
&body,
&[],
std::time::Duration::from_millis(1),
);
drop(fleet);
let mut ids = Vec::new();
for dir in [&alpha_dir, &beta_dir] {
let s = Semlith::open(dir, None).unwrap();
let rows = store::retrievals(s.db(), 10).unwrap();
assert_eq!(rows.len(), 1, "{}: {rows:?}", dir.display());
assert_eq!(rows[0].hits, 1, "{}: {rows:?}", dir.display());
assert!(rows[0].whole_file_tokens > 0);
assert_eq!(store::ledger_break(s.db()).unwrap(), None);
ids.push(rows[0].query_id.clone());
}
assert_eq!(ids[0], ids[1], "one reply is one retrieval");
}
#[test]
fn a_zero_hit_reply_is_one_row_in_the_store_searched() {
let tmp = tempfile::tempdir().unwrap();
let (alpha_dir, beta_dir) = (tmp.path().join("alpha"), tmp.path().join("beta"));
drop(store_holding(&alpha_dir, &[]));
drop(store_holding(&beta_dir, &[]));
let fleet = Fleet::open(&[alpha_dir.clone(), beta_dir.clone()]).unwrap();
let who = Who {
client: "harness",
session: "s",
version: "",
};
let nothing = "No match for that query.";
let at = std::time::Duration::from_millis(1);
reply(
&fleet,
&who,
"search",
"q1",
nothing,
&["beta".to_string()],
at,
);
reply(&fleet, &who, "search", "q2", nothing, &[], at);
drop(fleet);
let queries = |dir: &Path| -> Vec<(String, i64)> {
let s = Semlith::open(dir, None).unwrap();
store::retrievals(s.db(), 10)
.unwrap()
.into_iter()
.map(|r| (r.query, r.hits))
.collect()
};
assert_eq!(queries(beta_dir.as_path()), vec![("q1".to_string(), 0)]);
assert_eq!(queries(alpha_dir.as_path()), vec![("q2".to_string(), 0)]);
}
}