1use super::hybrid_rank::compare_fused;
4use super::{
5 WorkspaceHybridSearchHit, WorkspaceRerankAlgorithm, WorkspaceRerankFallbackReason,
6 WorkspaceRerankMode, WorkspaceRerankOptions, WorkspaceRerankStatus, WorkspaceRetrievalError,
7 WorkspaceRetrievalResult,
8};
9use std::cmp::Ordering;
10use std::collections::HashMap;
11use std::sync::Arc;
12
13mod features;
14use features::{candidate_features, scratch_bytes, CandidateFeatures};
15
16const MAX_CANDIDATES: usize = 100;
17const MAX_FEATURE_BYTES_PER_CANDIDATE: usize = 4 * 1024;
18const MAX_FINGERPRINTS_PER_CANDIDATE: usize = 128;
19const MAX_SCRATCH_BYTES: usize = 4 * 1024 * 1024;
20const MAX_RESULTS_PER_FILE: usize = 2;
21const NEAR_DUPLICATE_THRESHOLD: f64 = 0.85;
22const RELEVANCE_WEIGHT: f64 = 0.70;
23const CHANNEL_AGREEMENT_WEIGHT: f64 = 0.10;
24const DIVERSITY_WEIGHT: f64 = 0.20;
25const NEAR_DUPLICATE_PENALTY: f64 = 0.50;
26const MAX_CHANNELS: usize = 4;
27const FNV_OFFSET: u64 = 0xcbf29ce484222325;
28const FNV_PRIME: u64 = 0x100000001b3;
29
30pub(super) struct RerankOutcome {
31 pub hits: Vec<WorkspaceHybridSearchHit>,
32 pub status: WorkspaceRerankStatus,
33}
34
35pub(super) fn rerank_status_not_run(requested_mode: WorkspaceRerankMode) -> WorkspaceRerankStatus {
36 WorkspaceRerankStatus {
37 requested_mode,
38 applied_mode: WorkspaceRerankMode::RrfOnly,
39 algorithm: WorkspaceRerankAlgorithm::RrfK60,
40 input_candidates: 0,
41 evaluated_candidates: 0,
42 selected_candidates: 0,
43 near_duplicate_candidates: 0,
44 selected_near_duplicates: 0,
45 feature_bytes: 0,
46 accounted_scratch_bytes: 0,
47 candidate_truncated: false,
48 fallback: None,
49 }
50}
51
52impl WorkspaceRerankOptions {
53 pub fn validate(self) -> WorkspaceRetrievalResult<Self> {
58 validate_range(
59 self.max_candidates,
60 1,
61 MAX_CANDIDATES,
62 "rerank.max_candidates",
63 )?;
64 validate_range(
65 self.max_feature_bytes_per_candidate,
66 4,
67 MAX_FEATURE_BYTES_PER_CANDIDATE,
68 "rerank.max_feature_bytes_per_candidate",
69 )?;
70 validate_range(
71 self.max_fingerprints_per_candidate,
72 1,
73 MAX_FINGERPRINTS_PER_CANDIDATE,
74 "rerank.max_fingerprints_per_candidate",
75 )?;
76 validate_range(
77 self.max_scratch_bytes,
78 1,
79 MAX_SCRATCH_BYTES,
80 "rerank.max_scratch_bytes",
81 )?;
82 Ok(self)
83 }
84}
85
86fn validate_range(
87 value: usize,
88 minimum: usize,
89 maximum: usize,
90 field: &'static str,
91) -> WorkspaceRetrievalResult<()> {
92 if !(minimum..=maximum).contains(&value) {
93 return Err(WorkspaceRetrievalError::InvalidConfiguration {
94 field,
95 reason: "is outside the supported bounded range",
96 });
97 }
98 Ok(())
99}
100
101pub(super) fn rerank_fused_candidates(
102 mut candidates: Vec<WorkspaceHybridSearchHit>,
103 limit: usize,
104 options: WorkspaceRerankOptions,
105) -> RerankOutcome {
106 candidates.sort_by(compare_fused);
107 let input_candidates = candidates.len();
108 if options.mode == WorkspaceRerankMode::RrfOnly {
109 return rrf_only(
110 candidates,
111 limit,
112 input_candidates,
113 options.mode,
114 0,
115 false,
116 None,
117 );
118 }
119 if options.validate().is_err() {
120 return rrf_only(
121 candidates,
122 limit,
123 input_candidates,
124 options.mode,
125 0,
126 false,
127 Some(WorkspaceRerankFallbackReason::InvalidConfiguration),
128 );
129 }
130
131 let candidate_truncated = candidates.len() > options.max_candidates;
132 let evaluated_candidates = candidates.len().min(options.max_candidates);
133 let accounted_scratch_bytes = scratch_bytes(evaluated_candidates, options);
134 if accounted_scratch_bytes > options.max_scratch_bytes {
135 return rrf_only(
136 candidates,
137 limit,
138 input_candidates,
139 options.mode,
140 accounted_scratch_bytes,
141 candidate_truncated,
142 Some(WorkspaceRerankFallbackReason::ScratchBudgetExceeded),
143 );
144 }
145 candidates = bounded_candidate_pool(candidates, options.max_candidates);
146
147 let features = candidates
148 .iter()
149 .map(|candidate| candidate_features(&candidate.chunk.text, options))
150 .collect::<Vec<_>>();
151 let feature_bytes = features.iter().fold(0usize, |total, features| {
152 total.saturating_add(features.feature_bytes)
153 });
154 let near_duplicate_candidates = count_near_duplicates(&candidates, &features);
155 let (hits, selected_near_duplicates) = select_mmr(candidates, &features, limit);
156 let selected_candidates = hits.len();
157
158 RerankOutcome {
159 hits,
160 status: WorkspaceRerankStatus {
161 requested_mode: options.mode,
162 applied_mode: WorkspaceRerankMode::Deterministic,
163 algorithm: WorkspaceRerankAlgorithm::RrfK60DeterministicMmrV1,
164 input_candidates,
165 evaluated_candidates: features.len(),
166 selected_candidates,
167 near_duplicate_candidates,
168 selected_near_duplicates,
169 feature_bytes,
170 accounted_scratch_bytes,
171 candidate_truncated,
172 fallback: None,
173 },
174 }
175}
176
177fn bounded_candidate_pool(
178 candidates: Vec<WorkspaceHybridSearchHit>,
179 maximum: usize,
180) -> Vec<WorkspaceHybridSearchHit> {
181 if candidates.len() <= maximum {
182 return candidates;
183 }
184
185 let mut selected = Vec::<usize>::with_capacity(maximum);
186 for per_file_quota in 1..=MAX_RESULTS_PER_FILE {
187 for (index, candidate) in candidates.iter().enumerate() {
188 if selected.len() == maximum {
189 break;
190 }
191 if selected.contains(&index) {
192 continue;
193 }
194 let selected_for_file = selected
195 .iter()
196 .filter(|selected_index| {
197 candidates[**selected_index].chunk.path == candidate.chunk.path
198 })
199 .count();
200 if selected_for_file < per_file_quota {
201 selected.push(index);
202 }
203 }
204 }
205 for index in 0..candidates.len() {
206 if selected.len() == maximum {
207 break;
208 }
209 if !selected.contains(&index) {
210 selected.push(index);
211 }
212 }
213 selected.sort_unstable();
214
215 let mut selected = selected.into_iter().peekable();
216 candidates
217 .into_iter()
218 .enumerate()
219 .filter_map(|(index, candidate)| {
220 if selected.next_if_eq(&index).is_some() {
221 Some(candidate)
222 } else {
223 None
224 }
225 })
226 .collect()
227}
228
229fn rrf_only(
230 candidates: Vec<WorkspaceHybridSearchHit>,
231 limit: usize,
232 input_candidates: usize,
233 requested_mode: WorkspaceRerankMode,
234 accounted_scratch_bytes: usize,
235 candidate_truncated: bool,
236 fallback: Option<WorkspaceRerankFallbackReason>,
237) -> RerankOutcome {
238 let mut per_file = HashMap::<Arc<str>, usize>::new();
239 let hits = candidates
240 .into_iter()
241 .filter_map(|mut candidate| {
242 let count = per_file
243 .entry(Arc::clone(&candidate.chunk.path))
244 .or_default();
245 if *count >= MAX_RESULTS_PER_FILE {
246 return None;
247 }
248 *count += 1;
249 candidate.rerank_score = candidate.fused_score;
250 candidate.redundancy_score = 0.0;
251 Some(candidate)
252 })
253 .take(limit)
254 .collect::<Vec<_>>();
255 let selected_candidates = hits.len();
256 RerankOutcome {
257 hits,
258 status: WorkspaceRerankStatus {
259 requested_mode,
260 applied_mode: WorkspaceRerankMode::RrfOnly,
261 algorithm: WorkspaceRerankAlgorithm::RrfK60,
262 input_candidates,
263 evaluated_candidates: 0,
264 selected_candidates,
265 near_duplicate_candidates: 0,
266 selected_near_duplicates: 0,
267 feature_bytes: 0,
268 accounted_scratch_bytes,
269 candidate_truncated,
270 fallback,
271 },
272 }
273}
274
275fn select_mmr(
276 candidates: Vec<WorkspaceHybridSearchHit>,
277 features: &[CandidateFeatures],
278 limit: usize,
279) -> (Vec<WorkspaceHybridSearchHit>, usize) {
280 let exact_max = maximum_fused_score(&candidates, true);
281 let non_exact_max = maximum_fused_score(&candidates, false);
282 let mut selected = Vec::<usize>::with_capacity(limit.min(candidates.len()));
283 let mut selected_scores = Vec::<(f64, f64)>::with_capacity(selected.capacity());
284 let mut selected_flags = vec![false; candidates.len()];
285 let mut per_file = HashMap::<Arc<str>, usize>::new();
286 let mut selected_near_duplicates = 0usize;
287
288 while selected.len() < limit {
289 let select_exact = candidates.iter().enumerate().any(|(index, candidate)| {
290 !selected_flags[index]
291 && candidate.exact_identifier
292 && per_file.get(&candidate.chunk.path).copied().unwrap_or(0) < MAX_RESULTS_PER_FILE
293 });
294 let maximum = if select_exact {
295 exact_max
296 } else {
297 non_exact_max
298 };
299 let mut best = None::<(usize, f64, f64)>;
300 for (index, candidate) in candidates.iter().enumerate() {
301 if selected_flags[index]
302 || candidate.exact_identifier != select_exact
303 || per_file.get(&candidate.chunk.path).copied().unwrap_or(0) >= MAX_RESULTS_PER_FILE
304 {
305 continue;
306 }
307 let redundancy = selected.iter().fold(0.0_f64, |maximum, selected_index| {
308 maximum.max(candidate_similarity(
309 candidate,
310 &features[index],
311 &candidates[*selected_index],
312 &features[*selected_index],
313 ))
314 });
315 let score = selection_score(candidate, maximum, redundancy);
316 let replace = best.is_none_or(|(best_index, best_score, _)| {
317 score
318 .total_cmp(&best_score)
319 .then_with(|| best_index.cmp(&index))
320 == Ordering::Greater
321 });
322 if replace {
323 best = Some((index, score, redundancy));
324 }
325 }
326 let Some((index, score, redundancy)) = best else {
327 break;
328 };
329 selected_flags[index] = true;
330 *per_file
331 .entry(Arc::clone(&candidates[index].chunk.path))
332 .or_default() += 1;
333 if redundancy >= NEAR_DUPLICATE_THRESHOLD {
334 selected_near_duplicates = selected_near_duplicates.saturating_add(1);
335 }
336 selected.push(index);
337 selected_scores.push((score, redundancy));
338 }
339
340 let hits = selected
341 .into_iter()
342 .zip(selected_scores)
343 .map(|(index, (score, redundancy))| {
344 let mut candidate = candidates[index].clone();
345 candidate.rerank_score = score;
346 candidate.redundancy_score = redundancy;
347 candidate
348 })
349 .collect();
350 (hits, selected_near_duplicates)
351}
352
353fn maximum_fused_score(candidates: &[WorkspaceHybridSearchHit], exact: bool) -> f64 {
354 candidates
355 .iter()
356 .filter(|candidate| candidate.exact_identifier == exact)
357 .map(|candidate| candidate.fused_score)
358 .filter(|score| score.is_finite() && *score > 0.0)
359 .max_by(f64::total_cmp)
360 .unwrap_or(1.0)
361}
362
363fn selection_score(
364 candidate: &WorkspaceHybridSearchHit,
365 maximum_fused_score: f64,
366 redundancy: f64,
367) -> f64 {
368 let relevance = (candidate.fused_score / maximum_fused_score).clamp(0.0, 1.0);
369 let agreement = (candidate.channels.len() as f64 / MAX_CHANNELS as f64).clamp(0.0, 1.0);
370 let duplicate_penalty = if redundancy >= NEAR_DUPLICATE_THRESHOLD {
371 NEAR_DUPLICATE_PENALTY * redundancy
372 } else {
373 0.0
374 };
375 RELEVANCE_WEIGHT * relevance
376 + CHANNEL_AGREEMENT_WEIGHT * agreement
377 + DIVERSITY_WEIGHT * (1.0 - redundancy)
378 - duplicate_penalty
379}
380
381fn count_near_duplicates(
382 candidates: &[WorkspaceHybridSearchHit],
383 features: &[CandidateFeatures],
384) -> usize {
385 candidates
386 .iter()
387 .enumerate()
388 .filter(|(index, candidate)| {
389 (0..*index).any(|prior| {
390 candidate_similarity(
391 candidate,
392 &features[*index],
393 &candidates[prior],
394 &features[prior],
395 ) >= NEAR_DUPLICATE_THRESHOLD
396 })
397 })
398 .count()
399}
400
401fn candidate_similarity(
402 left: &WorkspaceHybridSearchHit,
403 left_features: &CandidateFeatures,
404 right: &WorkspaceHybridSearchHit,
405 right_features: &CandidateFeatures,
406) -> f64 {
407 interval_overlap(left, right).max(fingerprint_jaccard(
408 &left_features.fingerprints,
409 &right_features.fingerprints,
410 ))
411}
412
413fn interval_overlap(left: &WorkspaceHybridSearchHit, right: &WorkspaceHybridSearchHit) -> f64 {
414 if left.chunk.path != right.chunk.path {
415 return 0.0;
416 }
417 let intersection = left
418 .chunk
419 .end_byte
420 .min(right.chunk.end_byte)
421 .saturating_sub(left.chunk.start_byte.max(right.chunk.start_byte));
422 let denominator = left
423 .chunk
424 .end_byte
425 .saturating_sub(left.chunk.start_byte)
426 .min(right.chunk.end_byte.saturating_sub(right.chunk.start_byte));
427 if denominator == 0 {
428 0.0
429 } else {
430 intersection as f64 / denominator as f64
431 }
432}
433
434fn fingerprint_jaccard(left: &[u64], right: &[u64]) -> f64 {
435 if left.is_empty() || right.is_empty() {
436 return 0.0;
437 }
438 let mut left_index = 0usize;
439 let mut right_index = 0usize;
440 let mut intersection = 0usize;
441 while left_index < left.len() && right_index < right.len() {
442 match left[left_index].cmp(&right[right_index]) {
443 Ordering::Less => left_index += 1,
444 Ordering::Greater => right_index += 1,
445 Ordering::Equal => {
446 intersection += 1;
447 left_index += 1;
448 right_index += 1;
449 }
450 }
451 }
452 let union = left
453 .len()
454 .saturating_add(right.len())
455 .saturating_sub(intersection);
456 intersection as f64 / union as f64
457}