use std::collections::HashSet;
use anyhow::Result;
use once_cell::sync::Lazy;
use regex::Regex;
use serde_json::Value;
use unicode_segmentation::UnicodeSegmentation;
use crate::llmtrim::gate::{GateKind, PlanEntry, Transform};
use crate::llmtrim::ir::Request;
use crate::llmtrim::provider::Provider;
use crate::llmtrim::select::{self, Item, Weights};
use crate::llmtrim::stages::sizing::optimal_keep;
use crate::llmtrim::stages::tools::{detect_lang, lex_words, stopword_set};
pub struct RetrieveStage {
pub keep_ratio: f64,
pub min_segment_chars: usize,
pub reorder: bool,
pub mmr: bool,
pub mmr_lambda: f64,
pub sentence: bool,
}
impl Transform for RetrieveStage {
fn name(&self) -> &str {
"retrieve"
}
fn gate_kind(&self) -> GateKind {
GateKind::InputTokens
}
fn apply(
&self,
req: &mut Request,
provider: &dyn Provider,
_plan: &mut Vec<PlanEntry>,
) -> Result<()> {
let pointers = crate::llmtrim::cache_zone::compressible_pointers(req, provider);
let min = self.min_segment_chars;
use crate::llmtrim::provider::Role;
let segs: Vec<Seg> = pointers
.iter()
.map(|p| Seg {
idx: crate::llmtrim::provider::turn_index(p),
role: provider.role_at(req, p),
len: req.get_str(p).map(|s| s.chars().count()).unwrap_or(0),
is_tool_result: crate::llmtrim::provider::is_tool_result_ptr(req, p),
ptr: p.clone(),
})
.collect();
let is_system = |s: &Seg| s.role.is_none() || s.role == Some(Role::System);
let last_user = segs
.iter()
.filter(|s| s.role == Some(Role::User) && !s.is_tool_result)
.filter_map(|s| s.idx)
.max();
let is_last_user = |s: &Seg| s.idx.is_some() && s.idx == last_user && !s.is_tool_result;
let other_long_context = segs
.iter()
.any(|s| s.len >= min && !is_system(s) && !is_last_user(s));
let pinned: Vec<bool> = segs
.iter()
.map(|s| is_system(s) || (is_last_user(s) && other_long_context))
.collect();
let query_text: String = segs
.iter()
.filter(|s| is_last_user(s) || s.len < min)
.filter_map(|s| req.get_str(&s.ptr))
.collect::<Vec<_>>()
.join(" ");
let query = lex_words(&query_text);
for (k, s) in segs.iter().enumerate() {
if pinned[k] || s.len < min {
continue; }
let Some(text) = req.get_str(&s.ptr).map(str::to_string) else {
continue;
};
if serde_json::from_str::<Value>(text.trim())
.is_ok_and(|v| v.is_array() || v.is_object())
{
continue;
}
if let Some(rebuilt) = prune_protecting_directives(&text, &query, self) {
req.set(&s.ptr, Value::String(rebuilt));
}
}
Ok(())
}
}
static DIRECTIVE_BLOCK: Lazy<Regex> = Lazy::new(|| {
Regex::new(
r"(?is)<(?:system-reminder|system|instructions|important[_-]?instructions|directive)\b[^>]*>.*?</(?:system-reminder|system|instructions|important[_-]?instructions|directive)>",
)
.unwrap()
});
fn directive_spans(text: &str) -> Vec<(usize, usize)> {
DIRECTIVE_BLOCK
.find_iter(text)
.map(|m| (m.start(), m.end()))
.collect()
}
fn prune_protecting_directives(
text: &str,
query: &[String],
cfg: &RetrieveStage,
) -> Option<String> {
let spans = directive_spans(text);
if spans.is_empty() {
return prune_one(text, query, cfg);
}
let mut out = String::new();
let mut pos = 0usize;
let mut changed = false;
for (start, end) in spans {
changed |= prune_gap(&text[pos..start], query, cfg, &mut out);
out.push_str(&text[start..end]); pos = end;
}
changed |= prune_gap(&text[pos..], query, cfg, &mut out);
changed.then_some(out)
}
fn prune_gap(gap: &str, query: &[String], cfg: &RetrieveStage, out: &mut String) -> bool {
if gap.len() >= cfg.min_segment_chars
&& let Some(pruned) = prune_one(gap, query, cfg)
{
out.push_str(&pruned);
true
} else {
out.push_str(gap);
false
}
}
fn prune_one(text: &str, query: &[String], cfg: &RetrieveStage) -> Option<String> {
if cfg.sentence {
rebuild_sentence(text, query, cfg.keep_ratio)
} else {
rebuild_chunked(text, query, cfg)
}
}
fn rebuild_sentence(text: &str, query: &[String], keep_ratio: f64) -> Option<String> {
if query.is_empty() {
return None;
}
let chunks = sentence_chunks(text);
if chunks.len() < 3 {
return None;
}
let stops = stopword_set(text);
let kept = prune_sentences(&chunks, query, stops, keep_ratio);
if kept.len() >= chunks.len() {
return None; }
Some(rebuild(&chunks, &kept, PARA_SEP))
}
fn rebuild_chunked(text: &str, query: &[String], cfg: &RetrieveStage) -> Option<String> {
let (chunks, sep) = chunk_with_sep(text);
if chunks.len() < 2 {
return None;
}
let keep = ((chunks.len() as f64) * cfg.keep_ratio).ceil().max(1.0) as usize;
if keep >= chunks.len() {
return None; }
let protected = failure_protected(&chunks);
let with_protected = |mut idx: Vec<usize>| -> Vec<usize> {
idx.extend_from_slice(&protected);
idx.sort_unstable();
idx.dedup();
idx
};
let ranked = if query.is_empty() {
if chunks.len() > TEXTRANK_MAX_CHUNKS {
return Some(rebuild(
&chunks,
&with_protected(head_tail_keep(keep, chunks.len())),
sep,
));
}
textrank_rank(&chunks)
} else {
bm25_rank(&chunks, query)
};
if !cfg.reorder && !cfg.mmr {
return Some(rebuild(
&chunks,
&with_protected(budgeted_select(&chunks, &ranked, keep, cfg)),
sep,
));
}
let mut chosen = if cfg.mmr {
mmr_order(&ranked, &chunks, keep, cfg.mmr_lambda)
} else {
ranked.iter().copied().take(keep).collect::<Vec<usize>>()
};
chosen.extend_from_slice(&protected);
pin_boundaries(&mut chosen, chunks.len());
if cfg.reorder {
Some(rebuild_ordered(&chunks, &u_shape(&chosen), sep))
} else {
chosen.sort_unstable();
chosen.dedup();
Some(rebuild(&chunks, &chosen, sep))
}
}
fn failure_protected(chunks: &[String]) -> Vec<usize> {
use crate::llmtrim::stages::toolout::signals::STRONG;
static TAP_FAIL: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?i)^not ok\b").unwrap());
static TAP_RECORD: Lazy<Regex> =
Lazy::new(|| Regex::new(r"(?i)^(not ok\b|ok\s+\d|\d+\.\.\d+|#)").unwrap());
let mut out = Vec::new();
let mut indent_block = false; let mut tap_block = false; for (i, c) in chunks.iter().enumerate() {
let t = c.trim_start();
if tap_block && !TAP_RECORD.is_match(t) {
out.push(i); continue;
}
tap_block = false;
if STRONG.is_match(c) {
out.push(i);
indent_block = true;
tap_block = TAP_FAIL.is_match(t);
} else if indent_block && (c.starts_with(' ') || c.starts_with('\t')) {
out.push(i); } else {
indent_block = false;
}
}
out
}
const TEXTRANK_MAX_CHUNKS: usize = 2000;
fn head_tail_keep(keep: usize, n: usize) -> Vec<usize> {
let keep = keep.min(n);
let head = (keep * 2 / 3).max(1);
let tail = keep.saturating_sub(head);
let mut idx: Vec<usize> = (0..head.min(n)).collect();
idx.extend((n.saturating_sub(tail))..n);
pin_boundaries(&mut idx, n);
idx.sort_unstable();
idx.dedup();
idx
}
fn pin_boundaries(chosen: &mut Vec<usize>, n: usize) {
for b in [0, n - 1] {
if !chosen.contains(&b) {
chosen.push(b);
}
}
}
struct Seg {
ptr: String,
idx: Option<usize>,
role: Option<crate::llmtrim::provider::Role>,
len: usize,
is_tool_result: bool,
}
fn sentence_chunks(text: &str) -> Vec<String> {
let sentences: Vec<String> = text
.split_sentence_bounds()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
if sentences.len() > 1 {
sentences
} else {
chunk(text)
}
}
fn prune_sentences(
chunks: &[String],
query: &[String],
stops: &HashSet<&str>,
keep_ratio: f64,
) -> Vec<usize> {
let n = chunks.len();
let qset: HashSet<&str> = query
.iter()
.map(String::as_str)
.filter(|w| w.len() >= 2 && !stops.contains(w))
.collect();
if qset.is_empty() {
return (0..n).collect();
}
let mut buf = String::new();
let mut hits: HashSet<&str> = HashSet::new();
let score: Vec<usize> = chunks
.iter()
.map(|c| {
hits.clear();
for w in c.unicode_words() {
buf.clear();
buf.extend(w.chars().flat_map(char::to_lowercase));
if let Some(&q) = qset.get(buf.as_str()) {
hits.insert(q);
}
}
hits.len()
})
.collect();
if score.iter().all(|&s| s == 0) {
return (0..n).collect();
}
let relevant: Vec<bool> = score.iter().map(|&s| s > 0).collect();
let mut protected = vec![false; n];
for i in failure_protected(chunks) {
protected[i] = true;
}
let mut keep = vec![false; n];
for (i, k) in keep.iter_mut().enumerate() {
let neighbour = (i > 0 && relevant[i - 1]) || (i + 1 < n && relevant[i + 1]);
if relevant[i] || neighbour || protected[i] || i == 0 || i == n - 1 {
*k = true;
}
}
let target = ((n as f64) * keep_ratio).ceil().max(1.0) as usize;
let kept_count = keep.iter().filter(|&&k| k).count();
if kept_count > target {
let best = (0..n).max_by_key(|&i| score[i]).unwrap_or(0);
let mut droppable: Vec<usize> = (0..n)
.filter(|&i| keep[i] && !protected[i] && i != 0 && i != n - 1 && i != best)
.collect();
droppable.sort_by_key(|&i| (score[i], i)); let mut excess = kept_count - target;
for i in droppable {
if excess == 0 {
break;
}
keep[i] = false;
excess -= 1;
}
}
(0..n).filter(|&i| keep[i]).collect()
}
const PARA_SEP: &str = "\n\n";
const LINE_SEP: &str = "\n";
const TILE_SEP: &str = " ";
fn chunk_with_sep(text: &str) -> (Vec<String>, &'static str) {
let paras: Vec<String> = text
.split("\n\n")
.map(|p| p.trim().to_string())
.filter(|p| !p.is_empty())
.collect();
if paras.len() > 1 {
return (paras, PARA_SEP);
}
if let Some(tiles) = texttile(text) {
return (tiles, TILE_SEP);
}
let lines: Vec<String> = text
.lines()
.map(|l| l.trim().to_string())
.filter(|l| !l.is_empty())
.collect();
if lines.len() > 1 {
(lines, LINE_SEP)
} else {
(vec![text.trim().to_string()], PARA_SEP)
}
}
fn chunk(text: &str) -> Vec<String> {
chunk_with_sep(text).0
}
const TILE_BLOCK_SENTENCES: usize = 2;
const TILE_MIN_SENTENCES: usize = 6;
const TILE_MIN_TILE_SENTENCES: usize = 2;
const TILE_MAX_TILE_CHARS: usize = 1200;
const TILE_DEPTH_CUTOFF_SIGMA: f64 = 0.5;
const TILE_VALLEY_SIM_RATIO: f64 = 0.6;
fn texttile(text: &str) -> Option<Vec<String>> {
let sentences: Vec<&str> = text
.split_sentence_bounds()
.map(str::trim)
.filter(|s| !s.is_empty())
.collect();
if sentences.len() < TILE_MIN_SENTENCES {
return None;
}
let stops = stopword_set(text);
let bags: Vec<Vec<String>> = sentences
.iter()
.map(|s| {
lex_words(s)
.into_iter()
.filter(|w| w.chars().count() >= 2 && !stops.contains(w.as_str()))
.collect()
})
.collect();
let n = sentences.len();
let sims: Vec<f64> = (0..n - 1)
.map(|g| {
let left = block_counts(&bags, g + 1 - TILE_BLOCK_SENTENCES.min(g + 1), g + 1);
let right = block_counts(&bags, g + 1, (g + 1 + TILE_BLOCK_SENTENCES).min(n));
cosine(&left, &right)
})
.collect();
let depths = gap_depths(&sims);
let lens: Vec<usize> = sentences.iter().map(|s| s.chars().count()).collect();
let boundaries = pick_boundaries(&depths, &sims, &lens);
if boundaries.is_empty() {
return None; }
let mut tiles: Vec<String> = Vec::new();
let mut start = 0usize;
for &g in &boundaries {
tiles.push(sentences[start..=g].join(" "));
start = g + 1;
}
tiles.push(sentences[start..].join(" "));
Some(tiles)
}
fn block_counts(
bags: &[Vec<String>],
lo: usize,
hi: usize,
) -> std::collections::HashMap<&str, f64> {
let mut counts: std::collections::HashMap<&str, f64> = std::collections::HashMap::new();
for bag in &bags[lo..hi] {
for tok in bag {
*counts.entry(tok.as_str()).or_insert(0.0) += 1.0;
}
}
counts
}
fn cosine(
a: &std::collections::HashMap<&str, f64>,
b: &std::collections::HashMap<&str, f64>,
) -> f64 {
if a.is_empty() || b.is_empty() {
return 0.0;
}
let dot: f64 = a
.iter()
.map(|(t, &av)| av * b.get(t).copied().unwrap_or(0.0))
.sum();
let na: f64 = a.values().map(|v| v * v).sum::<f64>().sqrt();
let nb: f64 = b.values().map(|v| v * v).sum::<f64>().sqrt();
if na == 0.0 || nb == 0.0 {
0.0
} else {
dot / (na * nb)
}
}
fn gap_depths(sims: &[f64]) -> Vec<f64> {
let n = sims.len();
(0..n)
.map(|i| {
let mut left_peak = sims[i];
let mut j = i;
while j > 0 && sims[j - 1] >= sims[j] {
j -= 1;
left_peak = left_peak.max(sims[j]);
}
let mut right_peak = sims[i];
let mut k = i;
while k + 1 < n && sims[k + 1] >= sims[k] {
k += 1;
right_peak = right_peak.max(sims[k]);
}
(left_peak - sims[i]) + (right_peak - sims[i])
})
.collect()
}
fn pick_boundaries(depths: &[f64], sims: &[f64], lens: &[usize]) -> Vec<usize> {
let positive: Vec<f64> = depths.iter().copied().filter(|&d| d > 0.0).collect();
if positive.is_empty() {
return Vec::new();
}
let mean_depth = positive.iter().sum::<f64>() / positive.len() as f64;
let var = positive
.iter()
.map(|d| (d - mean_depth).powi(2))
.sum::<f64>()
/ positive.len() as f64;
let cutoff = mean_depth + TILE_DEPTH_CUTOFF_SIGMA * var.sqrt();
let mean_sim = sims.iter().sum::<f64>() / sims.len() as f64;
let valley_ceiling = TILE_VALLEY_SIM_RATIO * mean_sim;
let n = depths.len();
let mut chosen: Vec<usize> = Vec::new();
let mut last_cut: isize = -1; for g in 0..n {
let is_local_max = (g == 0 || depths[g] >= depths[g - 1])
&& (g + 1 == n || depths[g] >= depths[g + 1])
&& depths[g] > 0.0;
if is_local_max
&& depths[g] >= cutoff
&& sims[g] <= valley_ceiling
&& fits_min_tile(g, last_cut, lens.len())
{
chosen.push(g);
last_cut = g as isize;
}
}
split_oversized_tiles(&mut chosen, lens);
chosen
}
fn fits_min_tile(g: usize, last_cut: isize, n_sentences: usize) -> bool {
(g as isize - last_cut) >= TILE_MIN_TILE_SENTENCES as isize
&& (n_sentences - 1 - g) >= TILE_MIN_TILE_SENTENCES
}
fn split_oversized_tiles(chosen: &mut Vec<usize>, lens: &[usize]) {
if chosen.is_empty() {
return; }
let total: usize = lens.iter().sum();
if total <= TILE_MAX_TILE_CHARS {
return; }
let mut prefix = vec![0usize; lens.len() + 1];
for (i, &l) in lens.iter().enumerate() {
prefix[i + 1] = prefix[i] + l;
}
let span_chars = |start: usize, end: usize| prefix[end + 1] - prefix[start];
loop {
let starts: Vec<usize> = std::iter::once(0)
.chain(chosen.iter().map(|&g| g + 1))
.collect();
let ends: Vec<usize> = chosen
.iter()
.copied()
.chain(std::iter::once(lens.len() - 1))
.collect();
let mut added = false;
for (&start, &end) in starts.iter().zip(&ends) {
if span_chars(start, end) <= TILE_MAX_TILE_CHARS {
continue;
}
let mid = prefix[start] + span_chars(start, end) / 2;
let best = (start..end)
.filter(|&c| {
(c - start + 1) >= TILE_MIN_TILE_SENTENCES
&& (end - c) >= TILE_MIN_TILE_SENTENCES
})
.min_by_key(|&c| (prefix[c + 1].abs_diff(mid), c));
if let Some(c) = best {
chosen.push(c);
added = true;
}
}
if !added {
break; }
chosen.sort_unstable();
chosen.dedup();
}
}
fn bm25_language(sample: &str) -> bm25::Language {
use bm25::Language as B;
use whatlang::Lang;
match detect_lang(sample) {
Some(Lang::Ara) => B::Arabic,
Some(Lang::Dan) => B::Danish,
Some(Lang::Nld) => B::Dutch,
Some(Lang::Fra) => B::French,
Some(Lang::Deu) => B::German,
Some(Lang::Ell) => B::Greek,
Some(Lang::Hun) => B::Hungarian,
Some(Lang::Ita) => B::Italian,
Some(Lang::Nob) => B::Norwegian,
Some(Lang::Por) => B::Portuguese,
Some(Lang::Ron) => B::Romanian,
Some(Lang::Rus) => B::Russian,
Some(Lang::Spa) => B::Spanish,
Some(Lang::Swe) => B::Swedish,
Some(Lang::Tam) => B::Tamil,
Some(Lang::Tur) => B::Turkish,
_ => B::English,
}
}
const BM25_PLUS_DELTA: f64 = 1.0;
struct Bm25Index {
tokenizer: bm25::DefaultTokenizer,
embedder: bm25::Embedder<u32, bm25::DefaultTokenizer>,
scorer: bm25::Scorer<usize>,
chunk_tokens: Vec<Vec<String>>,
n: usize,
}
impl Bm25Index {
fn build(chunks: &[String]) -> Bm25Index {
use bm25::{DefaultTokenizer, EmbedderBuilder, Scorer, Tokenizer};
let refs: Vec<&str> = chunks.iter().map(String::as_str).collect();
let tokenizer = DefaultTokenizer::builder()
.language_mode(bm25_language(&refs.join(" ")))
.normalization(false)
.build();
let embedder =
EmbedderBuilder::<u32>::with_tokenizer_and_fit_to_corpus(tokenizer, &refs).build();
let tokenizer = DefaultTokenizer::builder()
.language_mode(bm25_language(&refs.join(" ")))
.normalization(false)
.build();
let chunk_tokens: Vec<Vec<String>> = refs.iter().map(|c| tokenizer.tokenize(c)).collect();
let mut scorer = Scorer::<usize>::new();
for (i, c) in refs.iter().enumerate() {
scorer.upsert(&i, embedder.embed(c));
}
Bm25Index {
tokenizer,
embedder,
scorer,
chunk_tokens,
n: chunks.len(),
}
}
fn query_terms(&self, query: &[String]) -> Vec<String> {
use bm25::Tokenizer;
let mut terms = self.tokenizer.tokenize(&query.join(" "));
terms.sort_unstable();
terms.dedup();
terms
}
fn idf(&self, term: &str) -> f64 {
let df = self
.chunk_tokens
.iter()
.filter(|toks| toks.iter().any(|t| t == term))
.count();
let n = self.n as f64;
(1.0 + (n - df as f64 + 0.5) / (df as f64 + 0.5)).ln()
}
fn scores(&self, terms: &[String], delta: f64) -> Vec<f64> {
let q = self.embedder.embed(&terms.join(" "));
let mut score = vec![0.0f64; self.n];
for m in self.scorer.matches(&q) {
if m.id < self.n {
score[m.id] = m.score as f64;
}
}
if delta != 0.0 {
let idf: Vec<f64> = terms.iter().map(|t| self.idf(t)).collect();
for (i, toks) in self.chunk_tokens.iter().enumerate() {
let present: HashSet<&str> = toks.iter().map(String::as_str).collect();
for (t, &w) in terms.iter().zip(&idf) {
if present.contains(t.as_str()) {
score[i] += w * delta;
}
}
}
}
score
}
fn weighted_scores(&self, weighted: &[(String, f64)], delta: f64) -> Vec<f64> {
let mut score = vec![0.0f64; self.n];
let one = [String::new()];
for (term, weight) in weighted {
if *weight == 0.0 {
continue;
}
let mut slot = one.clone();
slot[0] = term.clone();
for (i, s) in self.scores(&slot, delta).into_iter().enumerate() {
score[i] += weight * s;
}
}
score
}
}
#[cfg(test)]
fn bm25_scores(chunks: &[String], query: &[String], delta: f64) -> Vec<f64> {
let index = Bm25Index::build(chunks);
let terms = index.query_terms(query);
index.scores(&terms, delta)
}
fn bm25_rank(chunks: &[String], query: &[String]) -> Vec<usize> {
let index = Bm25Index::build(chunks);
let terms = index.query_terms(query);
let mut scores = index.scores(&terms, BM25_PLUS_DELTA);
rm3_rescore(&index, &terms, &mut scores);
argsort_desc(&scores)
}
const RM3_SPARSE_QUERY_TERMS: usize = 4;
const RM3_FEEDBACK_CHUNKS: usize = 4;
const RM3_EXPANSION_TERMS: usize = 10;
const RM3_LAMBDA: f64 = 0.5;
const RM3_FLATNESS_RATIO: f64 = 1.15;
fn rm3_rescore(index: &Bm25Index, query_terms: &[String], scores: &mut [f64]) {
let n = index.n;
if n < 3 || query_terms.is_empty() {
return; }
let order = argsort_desc(scores);
let m = RM3_FEEDBACK_CHUNKS.min(n);
if !rm3_should_fire(query_terms.len(), scores, &order, m) {
return;
}
let feedback = &order[..m];
let expansion = rm3_expansion_terms(index, query_terms, feedback, scores);
if expansion.is_empty() {
return; }
let mut weighted: Vec<(String, f64)> = query_terms
.iter()
.map(|t| (t.clone(), RM3_LAMBDA))
.collect();
let total: f64 = expansion.iter().map(|(_, w)| *w).sum();
if total > 0.0 {
for (term, w) in &expansion {
weighted.push((term.clone(), (1.0 - RM3_LAMBDA) * (w / total)));
}
}
let rescored = index.weighted_scores(&weighted, BM25_PLUS_DELTA);
scores.copy_from_slice(&rescored);
}
fn rm3_should_fire(n_query_terms: usize, scores: &[f64], order: &[usize], m: usize) -> bool {
if n_query_terms <= RM3_SPARSE_QUERY_TERMS {
return true;
}
let top = scores[order[0]];
let frontier = scores[order[m - 1]];
top > 0.0 && frontier > 0.0 && top <= frontier * RM3_FLATNESS_RATIO
}
fn rm3_expansion_terms(
index: &Bm25Index,
query_terms: &[String],
feedback: &[usize],
scores: &[f64],
) -> Vec<(String, f64)> {
let sample: String = feedback
.iter()
.filter_map(|&i| index.chunk_tokens.get(i))
.flat_map(|t| t.iter())
.map(String::as_str)
.collect::<Vec<_>>()
.join(" ");
let stops = stopword_set(&sample);
let query_set: HashSet<&str> = query_terms.iter().map(String::as_str).collect();
let mut weight: std::collections::HashMap<&str, f64> = std::collections::HashMap::new();
for &c in feedback {
let Some(toks) = index.chunk_tokens.get(c) else {
continue;
};
let len = toks.len() as f64;
if len == 0.0 {
continue;
}
let sc = scores.get(c).copied().unwrap_or(0.0);
if sc <= 0.0 {
continue; }
let mut tf: std::collections::HashMap<&str, f64> = std::collections::HashMap::new();
for t in toks {
*tf.entry(t.as_str()).or_insert(0.0) += 1.0;
}
for (term, count) in tf {
if term.chars().count() < 2 || stops.contains(term) || query_set.contains(term) {
continue;
}
*weight.entry(term).or_insert(0.0) += (count / len) * sc;
}
}
let mut ranked: Vec<(String, f64)> = weight
.into_iter()
.map(|(t, w)| (t.to_string(), w))
.collect();
ranked.sort_by(|a, b| {
b.1.partial_cmp(&a.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.0.cmp(&b.0))
});
ranked.truncate(RM3_EXPANSION_TERMS);
ranked
}
fn textrank_rank(chunks: &[String]) -> Vec<usize> {
let n = chunks.len();
let toks: Vec<HashSet<String>> = chunks
.iter()
.map(|c| lex_words(c).into_iter().collect())
.collect();
let ln_len: Vec<f64> = toks.iter().map(|t| ((t.len() + 1) as f64).ln()).collect();
let mut w = vec![vec![0.0f64; n]; n];
for i in 0..n {
for j in (i + 1)..n {
let inter = toks[i].intersection(&toks[j]).count() as f64;
let denom = ln_len[i] + ln_len[j];
let weight = if denom > 0.0 { inter / denom } else { 0.0 };
w[i][j] = weight;
w[j][i] = weight;
}
}
let damping = 0.85;
let mut score = vec![1.0 / n as f64; n];
let out_sum: Vec<f64> = w.iter().map(|row| row.iter().sum()).collect();
for _ in 0..30 {
let mut next = vec![(1.0 - damping) / n as f64; n];
for i in 0..n {
for j in 0..n {
if out_sum[j] > 0.0 {
next[i] += damping * (w[j][i] / out_sum[j]) * score[j];
}
}
}
score = next;
}
argsort_desc(&score)
}
fn argsort_desc(scores: &[f64]) -> Vec<usize> {
let mut idx: Vec<usize> = (0..scores.len()).collect();
idx.sort_by(|&a, &b| {
scores[b]
.partial_cmp(&scores[a])
.unwrap_or(std::cmp::Ordering::Equal)
.then(a.cmp(&b))
});
idx
}
#[cfg(test)]
fn top_k(ranked: &[usize], keep: usize) -> Vec<usize> {
let mut kept: Vec<usize> = ranked.iter().copied().take(keep).collect();
kept.sort_unstable();
kept
}
fn chunk_cost(chunk: &str) -> usize {
lex_words(chunk).len().max(1)
}
const RETRIEVE_WEIGHTS: Weights = Weights {
lambda: 0.5,
saturation: 0.3,
};
fn budgeted_select(
chunks: &[String],
ranked: &[usize],
keep: usize,
cfg: &RetrieveStage,
) -> Vec<usize> {
let n = chunks.len();
let costs: Vec<usize> = chunks.iter().map(|c| chunk_cost(c)).collect();
let total_tokens: usize = costs.iter().sum();
let min_cost = *costs.iter().min().unwrap_or(&1);
let budget = if cfg.keep_ratio > 0.0 && cfg.keep_ratio < 1.0 {
((total_tokens as f64) * cfg.keep_ratio).ceil() as usize
} else {
let k_sat = optimal_keep(
&chunks.iter().map(String::as_str).collect::<Vec<_>>(),
1,
keep,
);
ranked.iter().take(k_sat).map(|&i| costs[i]).sum()
}
.max(min_cost);
let mut rel = vec![0.0f64; n];
for (pos, &i) in ranked.iter().enumerate() {
if i < n {
rel[i] = (n - pos) as f64 / n as f64;
}
}
let items: Vec<Item> = (0..n)
.map(|i| Item::from_text(&chunks[i], costs[i], rel[i]))
.collect();
let mut chosen = select::select(&items, budget, &RETRIEVE_WEIGHTS);
pin_boundaries(&mut chosen, n);
chosen.sort_unstable();
chosen.dedup();
chosen
}
fn mmr_order(ranked: &[usize], chunks: &[String], keep: usize, lambda: f64) -> Vec<usize> {
let toks: Vec<HashSet<String>> = chunks
.iter()
.map(|c| lex_words(c).into_iter().collect())
.collect();
let total = ranked.len();
let denom = total.max(1) as f64;
let mut rank = vec![total; chunks.len()];
for (pos, &i) in ranked.iter().enumerate() {
if i < rank.len() {
rank[i] = pos;
}
}
let rel = |idx: usize| -> f64 {
let pos = rank.get(idx).copied().unwrap_or(total);
(total - pos) as f64 / denom
};
let sim = |a: usize, b: usize| -> f64 {
let (ta, tb) = (&toks[a], &toks[b]);
if ta.is_empty() || tb.is_empty() {
return 0.0;
}
let inter = ta.intersection(tb).count() as f64;
let uni = ta.union(tb).count() as f64;
if uni > 0.0 { inter / uni } else { 0.0 }
};
let mut selected: Vec<usize> = Vec::new();
let mut pool: Vec<usize> = ranked.to_vec();
while selected.len() < keep && !pool.is_empty() {
let best = pool
.iter()
.copied()
.max_by(|&a, &b| {
let score = |i: usize| {
let red = selected.iter().map(|&s| sim(i, s)).fold(0.0_f64, f64::max);
lambda * rel(i) - (1.0 - lambda) * red
};
score(a)
.partial_cmp(&score(b))
.unwrap_or(std::cmp::Ordering::Equal)
})
.unwrap();
selected.push(best);
pool.retain(|&x| x != best);
}
selected
}
fn u_shape(ranked: &[usize]) -> Vec<usize> {
let mut seen = HashSet::new();
let (mut front, mut back) = (Vec::new(), Vec::new());
for (i, &x) in ranked.iter().filter(|&&x| seen.insert(x)).enumerate() {
if i % 2 == 0 {
front.push(x)
} else {
back.push(x)
}
}
back.reverse();
front.extend(back);
front
}
fn rebuild_ordered(chunks: &[String], order: &[usize], sep: &str) -> String {
let dropped = chunks.len() - order.len();
let mut parts: Vec<String> = order.iter().map(|&i| chunks[i].clone()).collect();
if dropped > 0 {
parts.push(format!(
"[… {dropped} chunk(s) omitted, kept chunks reordered by relevance …]"
));
}
parts.join(sep)
}
fn rebuild(chunks: &[String], keep_idx: &[usize], sep: &str) -> String {
let keep_set: HashSet<usize> = keep_idx.iter().copied().collect();
let mut parts: Vec<String> = Vec::new();
let mut dropped = 0usize;
for (i, c) in chunks.iter().enumerate() {
if keep_set.contains(&i) {
if dropped > 0 {
parts.push(format!("[… {dropped} chunk(s) omitted …]"));
dropped = 0;
}
parts.push(c.clone());
} else {
dropped += 1;
}
}
if dropped > 0 {
parts.push(format!("[… {dropped} chunk(s) omitted …]"));
}
parts.join(sep)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::llmtrim::ir::ProviderKind;
use crate::llmtrim::pipeline;
use crate::llmtrim::provider::OpenAiProvider;
use crate::llmtrim::tokenizer::counter_for;
use serde_json::json;
fn doc() -> String {
[
"The cafeteria serves lunch from noon until two in the afternoon.",
"Parking is available in the north lot for all visitors and staff.",
"The quarterly revenue figure for the logistics division was 4.2 million.",
"Recycling bins are located on every floor near the elevators.",
"Office hours run from nine to five on weekdays only.",
]
.join("\n\n")
}
#[test]
fn bm25_recall_at_k_finds_the_answer_chunk() {
let chunks = chunk(&doc());
let query = lex_words("what was the quarterly revenue for logistics");
let ranked = bm25_rank(&chunks, &query);
assert_eq!(ranked[0], 2, "BM25 should rank the answer chunk top");
assert!(top_k(&ranked, 2).contains(&2));
}
#[test]
fn bm25_ranks_answer_in_non_latin_script() {
let chunks = vec![
"食堂在中午到下午两点之间供应午餐".to_string(),
"北方停车场为所有访客和员工提供停车位".to_string(),
"物流部门本季度的收入为四百二十万".to_string(),
"回收箱位于每层楼电梯附近".to_string(),
"工作时间为工作日上午九点到下午五点".to_string(),
];
let query = lex_words("物流部门本季度的收入是多少");
let ranked = bm25_rank(&chunks, &query);
assert_eq!(ranked[0], 2, "BM25 ranks the revenue chunk top in Chinese");
}
#[test]
fn textrank_ranks_all_chunks() {
let chunks = chunk(&doc());
let ranked = textrank_rank(&chunks);
assert_eq!(ranked.len(), chunks.len());
let unique: HashSet<usize> = ranked.iter().copied().collect();
assert_eq!(unique.len(), chunks.len(), "ranking is a permutation");
}
#[test]
fn rebuild_collapses_dropped_runs() {
let chunks: Vec<String> = (0..5).map(|i| format!("chunk{i}")).collect();
let out = rebuild(&chunks, &[0, 3], PARA_SEP);
assert_eq!(
out,
"chunk0\n\n[… 2 chunk(s) omitted …]\n\nchunk3\n\n[… 1 chunk(s) omitted …]"
);
}
#[test]
fn rebuild_rejoins_line_chunks_at_line_grain() {
let log = (0..6)
.map(|i| format!("line{i}"))
.collect::<Vec<_>>()
.join("\n");
let (chunks, sep) = chunk_with_sep(&log);
assert_eq!(sep, "\n", "line-chunked text remembers the line grain");
let out = rebuild(&chunks, &[0, 1, 5], sep);
assert!(
!out.contains("\n\n"),
"no paragraph gaps in a line-chunked rejoin: {out:?}"
);
assert!(out.contains("line0\nline1"), "kept lines stay line-joined");
}
#[test]
fn head_tail_keep_is_boundary_safe_and_bounded() {
let n = 10_000;
let keep = head_tail_keep(30, n);
assert!(keep.contains(&0) && keep.contains(&(n - 1)), "edges pinned");
assert!(
keep.len() <= 32,
"bounded to the budget, got {}",
keep.len()
);
assert!(keep.windows(2).all(|w| w[0] < w[1]), "sorted, deduped");
}
#[test]
fn sentence_chunks_split_on_terminal_punctuation() {
let c = sentence_chunks("First fact here. Second fact there! Third one? Done.");
assert_eq!(c.len(), 4, "split into four sentences");
assert!(c[0].contains("First fact"));
}
#[test]
fn prune_sentences_keeps_relevant_and_neighbours_drops_filler() {
let chunks: Vec<String> = [
"The vault access code is 7741.", "Lunch is served at noon each day.", "Parking validation is at the desk.", "Weather today is sunny and warm.", "Use the vault code to open the safe.", ]
.iter()
.map(|s| s.to_string())
.collect();
let query = vec!["vault".to_string(), "code".to_string()];
let stops = stopword_set("the vault is here and the code is there for a test");
let kept = prune_sentences(&chunks, &query, stops, 0.9);
assert!(
kept.contains(&0) && kept.contains(&4),
"relevant sentences kept"
);
assert!(!kept.contains(&2), "zero-overlap middle filler dropped");
let mut sorted = kept.clone();
sorted.sort_unstable();
assert_eq!(kept, sorted, "kept in original order");
let tight = prune_sentences(&chunks, &query, stops, 0.4);
assert!(
tight.contains(&0) && tight.contains(&4),
"answer + boundaries protected under aggressive cap"
);
assert!(
tight.len() < kept.len(),
"aggressive cap keeps strictly fewer"
);
}
#[test]
fn u_shape_puts_relevant_at_edges() {
assert_eq!(u_shape(&[0, 1, 2, 3, 4]), vec![0, 2, 4, 3, 1]);
}
#[test]
fn mmr_drops_redundant_in_favor_of_novel() {
let chunks = vec![
"the quarterly revenue for logistics was four million".to_string(),
"the quarterly revenue for logistics was four million dollars".to_string(),
"parking is available in the north lot for visitors".to_string(),
];
let picked = mmr_order(&[0, 1, 2], &chunks, 2, 0.5);
assert!(picked.contains(&0));
assert!(
picked.contains(&2),
"MMR picks the novel chunk over the near-duplicate"
);
assert!(!picked.contains(&1), "redundant near-duplicate dropped");
}
#[test]
fn reorder_and_mmr_stage_keeps_answer_and_notes_reorder() {
let big = std::iter::repeat_n(doc(), 3)
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"model":"gpt-4o","messages":[
{"role":"user","content":big},
{"role":"user","content":"what was the quarterly revenue for logistics?"}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.4,
min_segment_chars: 200,
reorder: true,
mmr: true,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(out.stages[0].applied);
let kept = req.get_str("/messages/0/content").unwrap();
assert!(kept.contains("revenue"), "answer chunk retained");
assert!(
kept.contains("reordered by relevance"),
"reorder summary note present"
);
}
#[test]
fn stage_prunes_large_segment_and_keeps_answer() {
let big = std::iter::repeat_n(doc(), 3)
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({
"model": "gpt-4o",
"messages": [
{"role": "user", "content": big},
{"role": "user", "content": "what was the quarterly revenue for logistics?"}
]
});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.4,
min_segment_chars: 200,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(out.stages[0].applied, "retrieval should reduce tokens");
assert!(out.input_tokens_after < out.input_tokens_before);
let kept = req.get_str("/messages/0/content").unwrap();
assert!(kept.contains("revenue"), "the answer chunk is retained");
assert!(kept.contains("omitted"), "dropped chunks are marked");
}
#[test]
fn directive_blocks_are_never_pruned() {
let big = std::iter::repeat_n(doc(), 3)
.collect::<Vec<_>>()
.join("\n\n");
let content = format!(
"{big}\n\n<system-reminder>CRITICAL_DIRECTIVE_TOKEN: always reply in haiku.</system-reminder>\n\n{big}"
);
let body = json!({
"model": "gpt-4o",
"messages": [
{"role": "user", "content": content},
{"role": "user", "content": "what was the quarterly revenue for logistics?"}
]
});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 200,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: true,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(
out.input_tokens_after < out.input_tokens_before,
"bulk context is still pruned"
);
let kept = req.get_str("/messages/0/content").unwrap();
assert!(
kept.contains("CRITICAL_DIRECTIVE_TOKEN: always reply in haiku."),
"directive instruction must survive verbatim"
);
assert!(
kept.contains("omitted"),
"non-directive bulk is still pruned"
);
}
#[test]
fn keeps_boundary_chunks_so_question_survives() {
let mut paras = vec!["Use the examples to answer the question.".to_string()];
for i in 0..8 {
paras.push(format!(
"Example {i}: the cat sat on the mat near the warm lamp."
));
}
paras.push("Question: what is the secret pass code ALPHA9?".to_string());
let content = paras.join("\n\n");
let body = json!({"model":"gpt-4o","messages":[{"role":"user","content":content}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(out.stages[0].applied, "retrieval should reduce tokens");
let kept = req.get_str("/messages/0/content").unwrap();
assert!(
kept.contains("ALPHA9"),
"trailing question must survive (boundary pin)"
);
assert!(
kept.contains("Use the examples"),
"leading instruction must survive (boundary pin)"
);
assert!(kept.contains("omitted"), "middle examples are pruned");
}
#[test]
fn role_aware_pins_long_final_question_and_prunes_context() {
let ctx = (0..10)
.map(|i| format!("Context paragraph {i}: lorem ipsum dolor about topic {i} and more."))
.collect::<Vec<_>>()
.join("\n\n");
let question = "Question: based on the documents above, summarize the key financial \
figure for the logistics division and explain the quarterly trend in \
detail with full reasoning and supporting context."
.to_string();
assert!(
question.len() >= 120,
"question must exceed the prune threshold"
);
let body = json!({"model":"gpt-4o","messages":[
{"role":"user","content":ctx},
{"role":"user","content":question}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(out.stages[0].applied, "context pruned -> tokens cut");
assert_eq!(
req.get_str("/messages/1/content").unwrap(),
question,
"final user turn pinned verbatim"
);
assert!(
req.get_str("/messages/0/content")
.unwrap()
.contains("omitted"),
"bulk context is pruned"
);
}
#[test]
fn tool_result_on_final_turn_is_not_mistaken_for_the_live_question() {
use crate::llmtrim::provider::AnthropicProvider;
let earlier_question =
"what does the claude-api skill say about streaming responses?".to_string();
let skill_dump = (0..20)
.map(|i| {
format!(
"Section {i}: unrelated reference material about topic {i} \
with lots of filler prose that has nothing to do with streaming."
)
})
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"model":"claude-3-5-sonnet-20241022","messages":[
{"role":"user","content":earlier_question},
{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Skill","input":{}}]},
{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":skill_dump}]}]});
let mut req = Request::from_value(ProviderKind::Anthropic, body);
let counter = counter_for(ProviderKind::Anthropic, None).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &AnthropicProvider, counter.as_ref(), &stages);
assert!(
out.stages[0].applied,
"tool_result dump pruned -> tokens cut"
);
assert_eq!(
req.get_str("/messages/0/content").unwrap(),
earlier_question,
"earlier user turn untouched (too short to prune)"
);
assert!(
req.get_str("/messages/2/content/0/content")
.unwrap()
.contains("omitted"),
"tool_result on the newest turn is pruned like any other bulk context, not pinned"
);
}
#[test]
fn mixed_message_pins_only_the_text_block_not_the_tool_result_riding_with_it() {
use crate::llmtrim::provider::AnthropicProvider;
let follow_up_question = "Question: based on the tool output above, summarize the key \
financial figure for the logistics division and explain the \
quarterly trend in detail with full reasoning and context."
.to_string();
assert!(
follow_up_question.len() >= 120,
"question must exceed the prune threshold"
);
let skill_dump = (0..20)
.map(|i| {
format!(
"Section {i}: unrelated reference material about topic {i} \
with lots of filler prose that has nothing to do with the question."
)
})
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"model":"claude-3-5-sonnet-20241022","messages":[
{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Skill","input":{}}]},
{"role":"user","content":[
{"type":"tool_result","tool_use_id":"t1","content":skill_dump},
{"type":"text","text":follow_up_question}]}]});
let mut req = Request::from_value(ProviderKind::Anthropic, body);
let counter = counter_for(ProviderKind::Anthropic, None).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &AnthropicProvider, counter.as_ref(), &stages);
assert!(
out.stages[0].applied,
"tool_result riding with the question is pruned -> tokens cut"
);
assert_eq!(
req.get_str("/messages/1/content/1/text").unwrap(),
follow_up_question,
"trailing question text pinned verbatim even though it shares a turn with the tool_result"
);
assert!(
req.get_str("/messages/1/content/0/content")
.unwrap()
.contains("omitted"),
"tool_result sharing the final turn with the question is still pruned"
);
}
#[test]
fn tool_result_mid_conversation_is_unaffected_by_the_pin_fix() {
use crate::llmtrim::provider::AnthropicProvider;
let final_question = "Question: based on everything above, summarize the key financial \
figure for the logistics division and explain the quarterly \
trend in detail with full reasoning and supporting context."
.to_string();
assert!(
final_question.len() >= 120,
"question must exceed the prune threshold"
);
let skill_dump = (0..20)
.map(|i| {
format!(
"Section {i}: unrelated reference material about topic {i} \
with lots of filler prose that has nothing to do with the question."
)
})
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"model":"claude-3-5-sonnet-20241022","messages":[
{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Skill","input":{}}]},
{"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":skill_dump}]},
{"role":"assistant","content":"Here's a summary of the skill docs."},
{"role":"user","content":final_question}]});
let mut req = Request::from_value(ProviderKind::Anthropic, body);
let counter = counter_for(ProviderKind::Anthropic, None).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &AnthropicProvider, counter.as_ref(), &stages);
assert!(
out.stages[0].applied,
"mid-conversation tool_result pruned -> tokens cut"
);
assert_eq!(
req.get_str("/messages/3/content").unwrap(),
final_question,
"the genuinely final user turn is still pinned verbatim"
);
assert!(
req.get_str("/messages/1/content/0/content")
.unwrap()
.contains("omitted"),
"mid-conversation tool_result is pruned like any other bulk context"
);
}
#[test]
fn gemini_function_response_on_final_turn_is_not_mistaken_for_the_live_question() {
use crate::llmtrim::provider::GoogleProvider;
let earlier_question =
"what does the claude-api skill say about streaming responses?".to_string();
let skill_dump = (0..20)
.map(|i| {
format!(
"Section {i}: unrelated reference material about topic {i} \
with lots of filler prose that has nothing to do with streaming."
)
})
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"contents":[
{"role":"user","parts":[{"text":earlier_question}]},
{"role":"model","parts":[{"functionCall":{"name":"Skill","args":{}}}]},
{"role":"user","parts":[{"functionResponse":{"name":"Skill","response":{"result":skill_dump}}}]}]});
let mut req = Request::from_value(ProviderKind::Google, body);
let counter = counter_for(ProviderKind::Google, Some("gemini-1.5-pro")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.3,
min_segment_chars: 120,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &GoogleProvider, counter.as_ref(), &stages);
assert!(
out.stages[0].applied,
"functionResponse dump pruned -> tokens cut"
);
assert_eq!(
req.get_str("/contents/0/parts/0/text").unwrap(),
earlier_question,
"earlier user turn untouched (too short to prune)"
);
assert!(
req.get_str("/contents/2/parts/0/functionResponse/response/result")
.unwrap()
.contains("omitted"),
"functionResponse on the newest turn is pruned like any other bulk context, not pinned"
);
}
#[test]
fn role_aware_works_on_google_contents_shape() {
use crate::llmtrim::provider::GoogleProvider;
let big = std::iter::repeat_n(doc(), 3)
.collect::<Vec<_>>()
.join("\n\n");
let body = json!({"contents":[
{"role":"user","parts":[{"text":big}]},
{"role":"user","parts":[{"text":"what was the quarterly revenue for logistics?"}]}],
"generationConfig":{"maxOutputTokens":64}});
let mut req = Request::from_value(ProviderKind::Google, body);
let counter = counter_for(ProviderKind::Google, Some("gemini-1.5-pro")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.4,
min_segment_chars: 200,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &GoogleProvider, counter.as_ref(), &stages);
assert!(
out.stages[0].applied,
"retrieval must run on the Google shape, not no-op"
);
let kept = req.get_str("/contents/0/parts/0/text").unwrap();
assert!(kept.contains("revenue"), "answer chunk retained");
assert!(kept.contains("omitted"), "bulk context pruned");
assert_eq!(
req.get_str("/contents/1/parts/0/text").unwrap(),
"what was the quarterly revenue for logistics?",
"final user question pinned verbatim"
);
}
#[test]
fn short_segments_are_left_alone() {
let body =
json!({"model":"gpt-4o","messages":[{"role":"user","content":"short question?"}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(RetrieveStage {
keep_ratio: 0.4,
min_segment_chars: 200,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
})];
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(
!out.stages[0].applied,
"short content is the query, not pruned"
);
}
fn retrieve_cfg(keep_ratio: f64) -> RetrieveStage {
RetrieveStage {
keep_ratio,
min_segment_chars: 200,
reorder: false,
mmr: false,
mmr_lambda: 0.5,
sentence: false,
}
}
fn tap_log() -> String {
let mut lines = vec!["TAP version 13".to_string()];
for i in 1..=52 {
if i == 19 {
lines.extend(
[
"not ok 19 - normalize: Windows backslash path matches POSIX entry",
" ---",
" failureType: 'testCodeFailure'",
" error: |-",
" Expected values to be strictly deep-equal:",
" + Symbol(FAIL_OPEN)",
" - true",
" ...",
]
.map(String::from),
);
} else {
lines.push(format!(
"ok {i} - isInList: case {i} returns expected result"
));
lines.push(format!(" duration_ms: {}.123456", i * 3));
}
}
lines.push("1..52".to_string());
lines.push("# pass 51".to_string());
lines.push("# fail 1".to_string());
lines.join("\n")
}
#[test]
fn tap_failure_block_survives_chunk_pruning() {
let query = lex_words("run the tests");
let out =
rebuild_chunked(&tap_log(), &query, &retrieve_cfg(0.1)).expect("large log prunes");
assert!(out.contains("not ok 19"), "failure line survives: {out}");
assert!(
out.contains("testCodeFailure") && out.contains("Symbol(FAIL_OPEN)"),
"indented diagnostic block survives with its failure: {out}"
);
assert!(out.contains("# fail 1"), "summary fail count survives");
assert!(out.contains("omitted"), "passing noise still pruned");
}
#[test]
fn tap_diagnostic_survives_without_indentation() {
let chunks: Vec<String> = tap_log().lines().map(|l| l.trim().to_string()).collect();
let kept = failure_protected(&chunks);
let joined: String = kept.iter().map(|&i| chunks[i].as_str()).collect();
assert!(joined.contains("not ok 19"));
assert!(
joined.contains("Symbol(FAIL_OPEN)"),
"trimmed diagnostic protected"
);
assert!(
!joined.contains("ok 20"),
"block closes at the next test point"
);
}
#[test]
fn tap_failure_block_survives_sentence_pruning() {
let chunks: Vec<String> = tap_log().lines().map(String::from).collect();
let query = lex_words("run the tests");
let stops = stopword_set(&tap_log());
let kept = prune_sentences(&chunks, &query, stops, 0.1);
let joined: String = kept.iter().map(|&i| chunks[i].as_str()).collect();
assert!(
joined.contains("not ok 19"),
"failure line survives the cap"
);
assert!(
joined.contains("Symbol(FAIL_OPEN)"),
"diagnostic survives the cap"
);
}
#[test]
fn budgeted_select_prunes_redundant_near_duplicates_topk_kept() {
let chunks: Vec<String> = vec![
"alpha preamble about the office building main entrance and lobby desk".to_string(),
"the quarterly revenue for the logistics division was four point two million"
.to_string(),
"the quarterly revenue for the logistics division was four point two million"
.to_string(),
"the quarterly revenue for the logistics division was four point two million"
.to_string(),
"recycling bins sit on every floor right beside the north stairwell doors".to_string(),
"visitor parking validation happens at the front security reception window".to_string(),
"omega closing remarks about the parking garage exit gate and ramp barrier".to_string(),
];
let ranked = vec![1, 2, 3, 4, 5, 0, 6];
let keep = ((chunks.len() as f64) * 0.5).ceil() as usize;
let old = top_k(&ranked, keep + 2); assert!(
old.contains(&1) && old.contains(&2) && old.contains(&3),
"baseline top-k kept all three redundant copies: {old:?}"
);
let chosen = budgeted_select(&chunks, &ranked, keep, &retrieve_cfg(0.5));
let dup_kept = [1usize, 2, 3]
.iter()
.filter(|&&i| chosen.contains(&i))
.count();
assert!(
dup_kept <= 1,
"at most one revenue copy is kept (redundancy pruned): {chosen:?}"
);
assert!(
chosen.contains(&4) || chosen.contains(&5),
"a novel chunk is preferred over a redundant duplicate: {chosen:?}"
);
assert!(
chosen.contains(&0) && chosen.contains(&6),
"distinct boundary chunks retained: {chosen:?}"
);
}
#[test]
fn bm25_plus_long_match_beats_absence() {
let filler = "and the report also covered many other routine operational matters \
across the wider organisation in considerable repetitive detail";
let long_with_term = format!("the logistics division summary {filler} {filler} {filler}");
let short_without = "weather today is mild".to_string();
let chunks = vec![long_with_term, short_without];
let query = lex_words("logistics");
let plain = bm25_scores(&chunks, &query, 0.0);
let plus = bm25_scores(&chunks, &query, BM25_PLUS_DELTA);
assert_eq!(plain[1], 0.0, "absence scores zero under plain BM25");
assert_eq!(plus[1], 0.0, "δ adds nothing to a chunk lacking the term");
assert!(plus[0] > 0.0, "the matching chunk scores positive");
assert!(
plus[0] > plain[0],
"δ raises the matched long chunk's floor"
);
assert!(plus[0] > plus[1], "occurrence beats absence under BM25+");
}
#[test]
fn bm25_plus_keeps_equal_density_ranking() {
let unit = "the logistics division handled freight";
let short = unit.to_string();
let long = format!("{unit} {unit} {unit}");
let chunks = vec![short, long];
let query = lex_words("logistics");
let plain = bm25_scores(&chunks, &query, 0.0);
let plus = bm25_scores(&chunks, &query, BM25_PLUS_DELTA);
assert_eq!(
argsort_desc(&plain),
argsort_desc(&plus),
"δ floor preserves the equal-density ranking"
);
let bonus_short = plus[0] - plain[0];
let bonus_long = plus[1] - plain[1];
assert!(
(bonus_short - bonus_long).abs() < 1e-9,
"equal idf·δ bonus applied to both present chunks"
);
}
fn rm3_corpus() -> Vec<String> {
vec![
"the telescope observed the distant galaxy nebula and faint orbiting comet".to_string(),
"lunch in the cafeteria is served from noon until two each weekday".to_string(),
"the galaxy and the nebula drift past a comet far beyond the orbiting moon".to_string(),
"parking validation stamps are available at the front reception desk".to_string(),
"recycling bins sit on every floor right next to the elevator bank".to_string(),
]
}
#[test]
fn rm3_lifts_answer_sharing_vocabulary_with_top_chunk() {
let chunks = rm3_corpus();
let query = lex_words("telescope");
let baseline = argsort_desc(&bm25_scores(&chunks, &query, BM25_PLUS_DELTA));
let answer_baseline = baseline.iter().position(|&i| i == 2).unwrap();
let with_rm3 = bm25_rank(&chunks, &query);
let answer_rm3 = with_rm3.iter().position(|&i| i == 2).unwrap();
assert!(
answer_rm3 < answer_baseline,
"RM3 lifts the vocabulary-sharing answer chunk (was #{answer_baseline}, now #{answer_rm3})"
);
}
#[test]
fn rm3_is_a_noop_for_a_rich_decisive_query() {
let chunks = vec![
"the quarterly revenue figure for the logistics division reached four million"
.to_string(),
"lunch in the cafeteria is served from noon until two each weekday".to_string(),
"parking validation stamps are available at the front reception desk".to_string(),
"recycling bins sit on every floor right next to the elevator bank".to_string(),
"office working hours run from nine until five on weekdays only".to_string(),
];
let query = lex_words("quarterly revenue logistics division four million figure");
assert!(
{
let idx = Bm25Index::build(&chunks);
idx.query_terms(&query).len() > RM3_SPARSE_QUERY_TERMS
},
"fixture must exercise the rich-query branch"
);
let baseline = argsort_desc(&bm25_scores(&chunks, &query, BM25_PLUS_DELTA));
let with_rm3 = bm25_rank(&chunks, &query);
assert_eq!(
baseline, with_rm3,
"rich decisive query: RM3 leaves the ranking untouched"
);
assert_eq!(with_rm3[0], 0, "the revenue chunk still ranks first");
}
fn two_topic_prose() -> String {
"The garden soil was rich so the garden plants grew fast. \
Garden plants need water and the soil keeps the plants healthy. \
A healthy garden has deep soil and the plants love the garden. \
The garden plants and garden soil made the garden thrive. \
The telescope showed the stars and the night sky was clear. \
Stars in the sky shone bright through the telescope each night. \
The telescope tracked the stars while the sky stayed dark. \
A clear sky let the telescope find faint stars in the night."
.to_string()
}
#[test]
fn texttiling_splits_between_two_topics_not_mid_topic() {
let tiles = texttile(&two_topic_prose()).expect("two-topic prose should tile");
assert!(tiles.len() >= 2, "at least one boundary placed");
assert!(
tiles[0].contains("garden"),
"first tile holds the garden topic"
);
assert!(
!tiles[0].contains("telescope") && !tiles[0].contains("stars"),
"the boundary is at the topic shift, not mid-astronomy"
);
assert!(
tiles
.iter()
.any(|t| t.contains("telescope") || t.contains("stars")),
"a later tile holds the telescope topic"
);
}
#[test]
fn budgeted_stage_drops_duplicate_paragraphs_and_cuts_tokens() {
fn count_tokens(text: &str) -> usize {
text.split_whitespace().count()
}
let mut paras: Vec<String> = Vec::new();
for _ in 0..10 {
paras.push(
"The quarterly revenue figure for the logistics division was 4.2 million."
.to_string(),
);
}
paras.push("Parking is available in the north lot for all visitors.".to_string());
paras.push("Recycling bins are located on every floor near the elevators.".to_string());
paras.push("Office hours run from nine to five on weekdays only.".to_string());
let big = paras.join("\n\n");
let body = json!({"model":"gpt-4o","messages":[
{"role":"user","content":big},
{"role":"user","content":"what was the quarterly revenue for logistics?"}]});
let mut req = Request::from_value(ProviderKind::OpenAi, body);
let counter = counter_for(ProviderKind::OpenAi, Some("gpt-4o")).unwrap();
let stages: Vec<Box<dyn Transform>> = vec![Box::new(retrieve_cfg(0.5))];
let before = req
.get_str("/messages/0/content")
.map(count_tokens)
.unwrap();
let out = pipeline::run(&mut req, &OpenAiProvider, counter.as_ref(), &stages);
assert!(out.stages[0].applied, "duplicate-heavy context is pruned");
let kept = req.get_str("/messages/0/content").unwrap();
assert!(kept.contains("revenue"), "the answer survives");
assert!(kept.contains("omitted"), "redundant copies are elided");
assert!(
count_tokens(kept) < before,
"tokens reduced ({} -> {})",
before,
count_tokens(kept)
);
}
#[test]
fn texttiling_falls_back_on_single_topic() {
let single = "The logistics division shipped freight to every regional warehouse. \
Freight volumes rose as the division added new regional routes. \
Each warehouse tracked its freight against the division forecast. \
The division reviewed warehouse freight costs every quarter. \
Regional freight delays pushed the division to add more warehouses. \
Freight forecasts guided how the division staffed each warehouse."
.to_string();
assert_eq!(texttile(&single), None, "single-topic prose falls back");
}
#[test]
fn texttiling_is_deterministic() {
let text = two_topic_prose();
assert_eq!(texttile(&text), texttile(&text), "tiling is deterministic");
}
#[test]
fn texttiling_subdivides_oversized_tile_after_a_real_boundary() {
let mut s = String::from(
"The garden soil was rich so the garden plants grew fast. \
Garden plants need water and the garden soil keeps the plants healthy. \
A healthy garden has deep soil and the plants love the garden. ",
);
for _ in 0..20 {
s.push_str(
"The telescope tracked the bright stars across the clear night sky tonight. ",
);
}
let tiles = texttile(&s).expect("a real boundary plus an oversized tile should tile");
assert!(tiles.len() >= 3, "oversized telescope tile is subdivided");
assert!(
tiles[0].contains("garden") && !tiles[0].contains("telescope"),
"the first (real) boundary still separates the topics"
);
assert!(
tiles
.iter()
.all(|t| t.chars().count() <= TILE_MAX_TILE_CHARS),
"no tile exceeds the max-char cap after subdivision"
);
}
}