Skip to content

Commit f73a66f

Browse files
committed
feat(pilot): add binary pruning and pre-filtering for wide nodes
Add binary pruning functionality to quickly filter relevant candidates before full scoring. Introduce pre-filtering using NodeScorer to reduce LLM token costs when nodes have many children. - Add Prune InterventionPoint for binary relevance filtering - Implement PrefilterConfig and PruneConfig with configurable thresholds - Add pre-filtering logic in score_candidates_detailed to narrow candidate sets using NodeScorer before LLM calls - Implement binary_prune method in LLM pilot for quick yes/no relevance decisions - Update metrics collection to track pruning interventions - Add comprehensive tests for new configuration options
1 parent 4a392bb commit f73a66f

12 files changed

Lines changed: 380 additions & 8 deletions

File tree

rust/src/metrics/pilot.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@ pub enum InterventionPoint {
1818
Backtrack,
1919
/// Evaluating content sufficiency.
2020
Evaluate,
21+
/// Binary pruning for wide nodes.
22+
Prune,
2123
}
2224

2325
/// Helper to store f64 as u64 bits for atomic operations.
@@ -87,7 +89,7 @@ impl PilotMetrics {
8789
InterventionPoint::Start => {
8890
self.start_guidance_calls.fetch_add(1, Ordering::Relaxed);
8991
}
90-
InterventionPoint::Fork => {
92+
InterventionPoint::Fork | InterventionPoint::Prune => {
9193
self.fork_decisions.fetch_add(1, Ordering::Relaxed);
9294
}
9395
InterventionPoint::Backtrack => {

rust/src/retrieval/pilot/config.rs

Lines changed: 177 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@ pub struct PilotConfig {
2727
pub guide_at_backtrack: bool,
2828
/// Optional path to custom prompt templates.
2929
pub prompt_template_path: Option<String>,
30+
/// Pre-filtering configuration for reducing candidates before Pilot.
31+
pub prefilter: PrefilterConfig,
32+
/// Binary pruning configuration for quick relevance filtering.
33+
pub prune: PruneConfig,
3034
}
3135

3236
impl Default for PilotConfig {
@@ -38,6 +42,8 @@ impl Default for PilotConfig {
3842
guide_at_start: true,
3943
guide_at_backtrack: true,
4044
prompt_template_path: None,
45+
prefilter: PrefilterConfig::default(),
46+
prune: PruneConfig::default(),
4147
}
4248
}
4349
}
@@ -51,7 +57,7 @@ impl PilotConfig {
5157
}
5258
}
5359

54-
/// Create a high-quality config (more LLM calls).
60+
/// Create a high-quality config (more LLM calls, generous pre-filter).
5561
pub fn high_quality() -> Self {
5662
Self {
5763
mode: PilotMode::Aggressive,
@@ -71,10 +77,20 @@ impl PilotConfig {
7177
guide_at_start: true,
7278
guide_at_backtrack: true,
7379
prompt_template_path: None,
80+
prefilter: PrefilterConfig {
81+
threshold: 20,
82+
max_to_pilot: 20,
83+
enabled: true,
84+
},
85+
prune: PruneConfig {
86+
enabled: true,
87+
threshold: 25,
88+
min_keep: 5,
89+
},
7490
}
7591
}
7692

77-
/// Create a low-cost config (fewer LLM calls).
93+
/// Create a low-cost config (fewer LLM calls, aggressive pre-filter).
7894
pub fn low_cost() -> Self {
7995
Self {
8096
mode: PilotMode::Conservative,
@@ -94,13 +110,33 @@ impl PilotConfig {
94110
guide_at_start: false,
95111
guide_at_backtrack: true,
96112
prompt_template_path: None,
113+
prefilter: PrefilterConfig {
114+
threshold: 8,
115+
max_to_pilot: 8,
116+
enabled: true,
117+
},
118+
prune: PruneConfig {
119+
enabled: true,
120+
threshold: 12,
121+
min_keep: 2,
122+
},
97123
}
98124
}
99125

100126
/// Create a pure algorithm config (no LLM calls).
101127
pub fn algorithm_only() -> Self {
102128
Self {
103129
mode: PilotMode::AlgorithmOnly,
130+
prefilter: PrefilterConfig {
131+
threshold: 15,
132+
max_to_pilot: 15,
133+
enabled: false,
134+
},
135+
prune: PruneConfig {
136+
enabled: false,
137+
threshold: 20,
138+
min_keep: 3,
139+
},
104140
..Default::default()
105141
}
106142
}
@@ -228,6 +264,88 @@ impl InterventionConfig {
228264
}
229265
}
230266

267+
/// Configuration for NodeScorer-based pre-filtering before Pilot scoring.
268+
///
269+
/// When a node has many children, sending all to the LLM is wasteful.
270+
/// Pre-filtering uses cheap NodeScorer (keyword/BM25) to narrow the
271+
/// candidate set before expensive Pilot (LLM) scoring.
272+
#[derive(Debug, Clone, Serialize, Deserialize)]
273+
pub struct PrefilterConfig {
274+
/// Minimum number of candidates to trigger pre-filtering.
275+
///
276+
/// When `candidates.len()` exceeds this threshold, NodeScorer
277+
/// pre-filters before sending to Pilot.
278+
/// Default: 15.
279+
pub threshold: usize,
280+
281+
/// Maximum number of candidates passed to Pilot after pre-filtering.
282+
///
283+
/// NodeScorer's top-N are kept; the rest get NodeScorer-only scores.
284+
/// Default: 15.
285+
pub max_to_pilot: usize,
286+
287+
/// Whether pre-filtering is enabled.
288+
/// Default: true.
289+
pub enabled: bool,
290+
}
291+
292+
impl Default for PrefilterConfig {
293+
fn default() -> Self {
294+
Self {
295+
threshold: 15,
296+
max_to_pilot: 15,
297+
enabled: true,
298+
}
299+
}
300+
}
301+
302+
impl PrefilterConfig {
303+
/// Check if pre-filtering should be applied given the candidate count.
304+
pub fn should_prefilter(&self, candidate_count: usize) -> bool {
305+
self.enabled && candidate_count > self.threshold
306+
}
307+
}
308+
309+
/// Configuration for binary pruning before full Pilot scoring.
310+
///
311+
/// After P2 pre-filtering, if candidates still exceed this threshold,
312+
/// a lightweight LLM call asks "which are relevant?" before the full
313+
/// scoring call. This reduces the number of candidates that receive
314+
/// expensive detailed scoring.
315+
#[derive(Debug, Clone, Serialize, Deserialize)]
316+
pub struct PruneConfig {
317+
/// Whether binary pruning is enabled.
318+
/// Default: true.
319+
pub enabled: bool,
320+
321+
/// Trigger threshold — binary prune activates when the candidate
322+
/// count (after P2 pre-filtering) exceeds this value.
323+
/// Default: 20.
324+
pub threshold: usize,
325+
326+
/// Minimum candidates to keep after pruning, even if LLM says
327+
/// fewer are relevant. Prevents over-aggressive pruning.
328+
/// Default: 3.
329+
pub min_keep: usize,
330+
}
331+
332+
impl Default for PruneConfig {
333+
fn default() -> Self {
334+
Self {
335+
enabled: true,
336+
threshold: 20,
337+
min_keep: 3,
338+
}
339+
}
340+
}
341+
342+
impl PruneConfig {
343+
/// Check if binary pruning should be applied given the candidate count.
344+
pub fn should_prune(&self, candidate_count: usize) -> bool {
345+
self.enabled && candidate_count > self.threshold
346+
}
347+
}
348+
231349
#[cfg(test)]
232350
mod tests {
233351
use super::*;
@@ -269,11 +387,68 @@ mod tests {
269387
fn test_pilot_config_presets() {
270388
let high = PilotConfig::high_quality();
271389
assert_eq!(high.mode, PilotMode::Aggressive);
390+
assert!(high.prefilter.enabled);
391+
assert_eq!(high.prefilter.threshold, 20);
272392

273393
let low = PilotConfig::low_cost();
274394
assert_eq!(low.mode, PilotMode::Conservative);
395+
assert!(low.prefilter.enabled);
396+
assert_eq!(low.prefilter.threshold, 8);
275397

276398
let algo = PilotConfig::algorithm_only();
277399
assert_eq!(algo.mode, PilotMode::AlgorithmOnly);
400+
assert!(!algo.prefilter.enabled);
401+
}
402+
403+
#[test]
404+
fn test_prefilter_config_default() {
405+
let cfg = PrefilterConfig::default();
406+
assert!(cfg.enabled);
407+
assert_eq!(cfg.threshold, 15);
408+
assert_eq!(cfg.max_to_pilot, 15);
409+
}
410+
411+
#[test]
412+
fn test_prefilter_should_prefilter() {
413+
let cfg = PrefilterConfig::default();
414+
assert!(!cfg.should_prefilter(15)); // at threshold
415+
assert!(!cfg.should_prefilter(10)); // below
416+
assert!(cfg.should_prefilter(16)); // above
417+
418+
let disabled = PrefilterConfig { enabled: false, ..Default::default() };
419+
assert!(!disabled.should_prefilter(100));
420+
}
421+
422+
#[test]
423+
fn test_prune_config_default() {
424+
let cfg = PruneConfig::default();
425+
assert!(cfg.enabled);
426+
assert_eq!(cfg.threshold, 20);
427+
assert_eq!(cfg.min_keep, 3);
428+
}
429+
430+
#[test]
431+
fn test_prune_should_prune() {
432+
let cfg = PruneConfig::default();
433+
assert!(!cfg.should_prune(20)); // at threshold
434+
assert!(!cfg.should_prune(15)); // below
435+
assert!(cfg.should_prune(21)); // above
436+
437+
let disabled = PruneConfig { enabled: false, ..Default::default() };
438+
assert!(!disabled.should_prune(100));
439+
}
440+
441+
#[test]
442+
fn test_pilot_config_presets_prune() {
443+
let high = PilotConfig::high_quality();
444+
assert!(high.prune.enabled);
445+
assert_eq!(high.prune.threshold, 25);
446+
447+
let low = PilotConfig::low_cost();
448+
assert!(low.prune.enabled);
449+
assert_eq!(low.prune.threshold, 12);
450+
451+
let algo = PilotConfig::algorithm_only();
452+
assert!(!algo.prune.enabled);
278453
}
279454
}

rust/src/retrieval/pilot/decision.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -220,6 +220,8 @@ pub enum InterventionPoint {
220220
Backtrack,
221221
/// Evaluating a specific node for relevance.
222222
Evaluate,
223+
/// Binary pruning — quick yes/no relevance filter for wide nodes.
224+
Prune,
223225
}
224226

225227
impl InterventionPoint {
@@ -230,6 +232,7 @@ impl InterventionPoint {
230232
Self::Fork => "fork",
231233
Self::Backtrack => "backtrack",
232234
Self::Evaluate => "evaluate",
235+
Self::Prune => "prune",
233236
}
234237
}
235238
}

rust/src/retrieval/pilot/decision_scorer.rs

Lines changed: 79 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,14 @@ pub struct ScoredCandidate {
112112
/// from the Pilot. Use this when the search algorithm needs to
113113
/// record why each path step was taken (e.g., for beam search
114114
/// reasoning history).
115+
///
116+
/// # Pre-filtering
117+
///
118+
/// When a node has many children (exceeding `prefilter.threshold`),
119+
/// NodeScorer pre-filters candidates before sending to Pilot. This
120+
/// reduces LLM token cost and latency. Candidates filtered out still
121+
/// receive NodeScorer-only scores in the final merge, so no results
122+
/// are lost.
115123
pub async fn score_candidates_detailed(
116124
tree: &DocumentTree,
117125
candidates: &[NodeId],
@@ -139,20 +147,88 @@ pub async fn score_candidates_detailed(
139147
// Determine parent node (last in path) for cache key
140148
let parent = path.last().copied().unwrap_or(tree.root());
141149

150+
// === PRE-FILTERING ===
151+
// When candidates exceed the threshold, use NodeScorer to narrow
152+
// the set before sending to Pilot (LLM). Filtered-out candidates
153+
// still get NodeScorer-only scores in the final merge below.
154+
let prefilter_cfg = &p.config().prefilter;
155+
let pilot_candidates: Vec<NodeId> = if prefilter_cfg.should_prefilter(candidates.len()) {
156+
let scorer = NodeScorer::new(ScoringContext::new(query));
157+
let mut sorted = scorer.score_and_sort(tree, candidates);
158+
let pilot_max = prefilter_cfg.max_to_pilot.min(candidates.len());
159+
sorted.truncate(pilot_max);
160+
let ids: Vec<NodeId> = sorted.into_iter().map(|(id, _)| id).collect();
161+
tracing::debug!(
162+
"Pre-filtered: {} candidates -> {} to Pilot (threshold={})",
163+
candidates.len(),
164+
ids.len(),
165+
prefilter_cfg.threshold,
166+
);
167+
ids
168+
} else {
169+
candidates.to_vec()
170+
};
171+
172+
// === BINARY PRUNING ===
173+
// After P2 pre-filtering, if candidates still exceed the prune
174+
// threshold, ask Pilot for a quick yes/no filter before the
175+
// expensive full-scoring call.
176+
let prune_cfg = &p.config().prune;
177+
let pilot_candidates = if prune_cfg.should_prune(pilot_candidates.len()) {
178+
let mut prune_state = SearchState::new(
179+
tree, query, path, &pilot_candidates, visited,
180+
);
181+
prune_state.step_reasons = step_reasons;
182+
183+
if let Some(relevant_ids) = p.binary_prune(&prune_state).await {
184+
let relevant_set: HashSet<NodeId> = relevant_ids.iter().copied().collect();
185+
let mut pruned: Vec<NodeId> = pilot_candidates
186+
.iter()
187+
.filter(|id| relevant_set.contains(id))
188+
.copied()
189+
.collect();
190+
191+
// Enforce min_keep to prevent over-aggressive pruning
192+
if pruned.len() < prune_cfg.min_keep {
193+
// Fill from the top of pilot_candidates that weren't pruned
194+
for id in &pilot_candidates {
195+
if pruned.len() >= prune_cfg.min_keep {
196+
break;
197+
}
198+
if !relevant_set.contains(id) {
199+
pruned.push(*id);
200+
}
201+
}
202+
}
203+
204+
tracing::debug!(
205+
"Binary prune: {} candidates -> {} relevant (min_keep={})",
206+
pilot_candidates.len(),
207+
pruned.len(),
208+
prune_cfg.min_keep,
209+
);
210+
pruned
211+
} else {
212+
pilot_candidates
213+
}
214+
} else {
215+
pilot_candidates
216+
};
217+
142218
// Check cache first
143219
let decision = if let Some(c) = cache {
144220
if let Some(cached) = c.get(query, parent).await {
145221
tracing::trace!("Pilot cache hit for parent={:?}", parent);
146222
cached
147223
} else {
148-
let mut state = SearchState::new(tree, query, path, candidates, visited);
224+
let mut state = SearchState::new(tree, query, path, &pilot_candidates, visited);
149225
state.step_reasons = step_reasons;
150226
let d = p.decide(&state).await;
151227
c.put(query, parent, &d).await;
152228
d
153229
}
154230
} else {
155-
let mut state = SearchState::new(tree, query, path, candidates, visited);
231+
let mut state = SearchState::new(tree, query, path, &pilot_candidates, visited);
156232
state.step_reasons = step_reasons;
157233
p.decide(&state).await
158234
};
@@ -163,7 +239,7 @@ pub async fn score_candidates_detailed(
163239
pilot_data.insert(ranked.node_id, (ranked.score, ranked.reason.clone()));
164240
}
165241

166-
// Compute NodeScorer fallback scores
242+
// Compute NodeScorer fallback scores for ALL original candidates
167243
let scorer_weight = 1.0 - pilot_weight;
168244
let confidence = decision.confidence;
169245
let effective_pilot = pilot_weight * confidence;

0 commit comments

Comments
 (0)