aft/commands/semantic_search/
confidence.rs

1use std::cmp::Ordering;
2use std::fmt;
3use std::fs;
4use std::path::{Path, PathBuf};
5
6use serde::{Deserialize, Serialize};
7use serde_json::{Map, Value};
8
9use super::blocks::{BlockEntry, BlockReply, CanonicalList};
10use super::evidence_descriptor::{EvidenceDescriptor, EvidenceKind, EvidenceTier};
11
12pub const CONFIDENCE_THRESHOLD_RELATIVE_PATH: &str =
13    "benchmarks/aft-search/engine-fixtures/confidence-threshold.json";
14pub const FLAT_HEAD_LINE: &str = "flat head: ranks 1 and 2 are not meaningfully separated";
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17#[serde(rename_all = "snake_case")]
18pub enum Confidence {
19    High,
20    Low,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
24#[serde(rename_all = "snake_case")]
25pub enum ConfidenceBranch {
26    Empty,
27    SingletonExact,
28    SingletonNonExact,
29    ExactOverNonExact,
30    ExactEvidence,
31    NonExactDescriptor,
32    NonExactMargin,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
36pub struct ConfidenceDecision {
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub confidence: Option<Confidence>,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub flat_head_line: Option<&'static str>,
41    #[serde(skip)]
42    pub branch: ConfidenceBranch,
43}
44
45impl ConfidenceDecision {
46    fn empty() -> Self {
47        Self {
48            confidence: None,
49            flat_head_line: None,
50            branch: ConfidenceBranch::Empty,
51        }
52    }
53
54    fn high(branch: ConfidenceBranch) -> Self {
55        Self {
56            confidence: Some(Confidence::High),
57            flat_head_line: None,
58            branch,
59        }
60    }
61
62    fn singleton_low() -> Self {
63        Self {
64            confidence: Some(Confidence::Low),
65            flat_head_line: None,
66            branch: ConfidenceBranch::SingletonNonExact,
67        }
68    }
69
70    fn flat_head(branch: ConfidenceBranch) -> Self {
71        debug_assert!(matches!(
72            branch,
73            ConfidenceBranch::ExactEvidence | ConfidenceBranch::NonExactMargin
74        ));
75        Self {
76            confidence: Some(Confidence::Low),
77            flat_head_line: Some(FLAT_HEAD_LINE),
78            branch,
79        }
80    }
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
84#[serde(rename_all = "snake_case")]
85pub enum ConfidenceThresholdSource {
86    Provisional,
87    Calibrated,
88}
89
90#[derive(Debug, Clone, PartialEq, Serialize)]
91pub struct ConfidenceThreshold {
92    pub schema: u64,
93    pub source: ConfidenceThresholdSource,
94    pub margin: f32,
95    pub model_id: Option<String>,
96    pub calibrated_at: Option<String>,
97}
98
99impl ConfidenceThreshold {
100    pub fn load_at_startup(path: &Path) -> Result<Self, ConfidenceStartupError> {
101        let content = fs::read_to_string(path).map_err(|error| {
102            ConfidenceStartupError::new(path, "file", format!("failed to read file: {error}"))
103        })?;
104        Self::parse_at_startup(path, &content)
105    }
106
107    pub fn parse_at_startup(path: &Path, content: &str) -> Result<Self, ConfidenceStartupError> {
108        let value: Value = serde_json::from_str(content).map_err(|error| {
109            let field = if error.to_string().contains("number out of range")
110                && content.contains("\"margin\"")
111            {
112                "margin"
113            } else {
114                "json"
115            };
116            ConfidenceStartupError::new(path, field, format!("malformed JSON: {error}"))
117        })?;
118        let object = value.as_object().ok_or_else(|| {
119            ConfidenceStartupError::new(path, "json", "top-level value must be an object")
120        })?;
121
122        let schema = required_field(path, object, "schema")?
123            .as_u64()
124            .ok_or_else(|| ConfidenceStartupError::new(path, "schema", "must be the integer 1"))?;
125        if schema != 1 {
126            return Err(ConfidenceStartupError::new(
127                path,
128                "schema",
129                format!("unknown schema {schema}; expected 1"),
130            ));
131        }
132
133        let source = match required_field(path, object, "source")?.as_str() {
134            Some("provisional") => ConfidenceThresholdSource::Provisional,
135            Some("calibrated") => ConfidenceThresholdSource::Calibrated,
136            Some(other) => {
137                return Err(ConfidenceStartupError::new(
138                    path,
139                    "source",
140                    format!("unknown source {other:?}"),
141                ));
142            }
143            None => {
144                return Err(ConfidenceStartupError::new(
145                    path,
146                    "source",
147                    "must be \"provisional\" or \"calibrated\"",
148                ));
149            }
150        };
151
152        let margin_value = required_field(path, object, "margin")?;
153        let margin = match margin_value {
154            Value::Number(number) => number.as_f64(),
155            Value::String(text) if matches!(text.as_str(), "NaN" | "Infinity" | "-Infinity") => {
156                text.parse::<f64>().ok()
157            }
158            _ => None,
159        }
160        .ok_or_else(|| {
161            ConfidenceStartupError::new(path, "margin", "must be a finite non-negative number")
162        })?;
163        if !margin.is_finite() || margin < 0.0 || margin > f32::MAX as f64 {
164            return Err(ConfidenceStartupError::new(
165                path,
166                "margin",
167                "must be a finite non-negative number",
168            ));
169        }
170
171        let model_id = optional_string(path, object, "model_id")?;
172        let calibrated_at = optional_string(path, object, "calibrated_at")?;
173        if source == ConfidenceThresholdSource::Calibrated && model_id.is_none() {
174            return Err(ConfidenceStartupError::new(
175                path,
176                "model_id",
177                "must be non-null when source is \"calibrated\"",
178            ));
179        }
180
181        Ok(Self {
182            schema,
183            source,
184            margin: margin as f32,
185            model_id,
186            calibrated_at,
187        })
188    }
189}
190
191fn required_field<'a>(
192    path: &Path,
193    object: &'a Map<String, Value>,
194    field: &'static str,
195) -> Result<&'a Value, ConfidenceStartupError> {
196    object
197        .get(field)
198        .ok_or_else(|| ConfidenceStartupError::new(path, field, "required field is missing"))
199}
200
201fn optional_string(
202    path: &Path,
203    object: &Map<String, Value>,
204    field: &'static str,
205) -> Result<Option<String>, ConfidenceStartupError> {
206    match required_field(path, object, field)? {
207        Value::Null => Ok(None),
208        Value::String(value) => Ok(Some(value.clone())),
209        _ => Err(ConfidenceStartupError::new(
210            path,
211            field,
212            "must be a string or null",
213        )),
214    }
215}
216
217#[derive(Debug, Clone, PartialEq, Eq)]
218pub struct ConfidenceStartupError {
219    path: PathBuf,
220    field: &'static str,
221    message: String,
222}
223
224impl ConfidenceStartupError {
225    fn new(path: &Path, field: &'static str, message: impl Into<String>) -> Self {
226        Self {
227            path: path.to_path_buf(),
228            field,
229            message: message.into(),
230        }
231    }
232
233    pub fn path(&self) -> &Path {
234        &self.path
235    }
236
237    pub fn field(&self) -> &'static str {
238        self.field
239    }
240}
241
242impl fmt::Display for ConfidenceStartupError {
243    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
244        write!(
245            formatter,
246            "confidence threshold startup error in {} at field `{}`: {}",
247            self.path.display(),
248            self.field,
249            self.message
250        )
251    }
252}
253
254impl std::error::Error for ConfidenceStartupError {}
255
256#[derive(Debug, Clone, PartialEq, Eq)]
257pub enum ConfidenceError {
258    InvalidEvidenceDescriptor,
259    InvalidHeadOrder,
260    MissingFusionScore,
261    NonFiniteFusionScore,
262}
263
264impl fmt::Display for ConfidenceError {
265    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
266        match self {
267            Self::InvalidEvidenceDescriptor => {
268                formatter.write_str("confidence received an invalid evidence descriptor")
269            }
270            Self::InvalidHeadOrder => {
271                formatter.write_str("a non-exact rank 1 cannot precede an exact-tier rank 2")
272            }
273            Self::MissingFusionScore => {
274                formatter.write_str("non-exact confidence requires a frozen fusion score")
275            }
276            Self::NonFiniteFusionScore => {
277                formatter.write_str("non-exact confidence requires a finite fusion score")
278            }
279        }
280    }
281}
282
283impl std::error::Error for ConfidenceError {}
284
285/// The confidence view intentionally exposes evidence and frozen fusion only.
286/// Provenance is absent from this interface, so run-scoped lane positions cannot
287/// become a ranking input by accident.
288pub trait ConfidenceCandidate {
289    fn evidence_descriptor(&self) -> &EvidenceDescriptor;
290    fn frozen_fusion_score(&self) -> Option<f32>;
291}
292
293impl ConfidenceCandidate for BlockEntry {
294    fn evidence_descriptor(&self) -> &EvidenceDescriptor {
295        &self.result.evidence
296    }
297
298    fn frozen_fusion_score(&self) -> Option<f32> {
299        self.result.fusion_score
300    }
301}
302
303#[derive(Debug, Clone)]
304pub struct ConfidenceEngine {
305    threshold: ConfidenceThreshold,
306}
307
308/// The pinned threshold is a product input, so it lives beside the crate
309/// (`crates/aft/assets/`): a product build must not reach outside `crates/`
310/// (the release and Docker images copy only the crate tree). The benchmark
311/// keeps its own copy at `CONFIDENCE_THRESHOLD_RELATIVE_PATH`, and
312/// `engine_confidence_threshold_test` asserts the two are byte-identical so
313/// neither can drift from the other.
314pub const PINNED_CONFIDENCE_THRESHOLD_JSON: &str =
315    include_str!("../../../assets/search-confidence-threshold.pinned.json");
316
317impl ConfidenceEngine {
318    pub fn running() -> Self {
319        let path = Path::new(CONFIDENCE_THRESHOLD_RELATIVE_PATH);
320        let threshold =
321            ConfidenceThreshold::parse_at_startup(path, PINNED_CONFIDENCE_THRESHOLD_JSON)
322                .expect("embedded confidence threshold must satisfy the engine schema");
323        Self { threshold }
324    }
325
326    /// Loads the threshold artifact once and retains the validated value for all
327    /// queries served by this engine instance.
328    pub fn start(threshold_path: &Path) -> Result<Self, ConfidenceStartupError> {
329        Ok(Self {
330            threshold: ConfidenceThreshold::load_at_startup(threshold_path)?,
331        })
332    }
333
334    pub fn start_at_workspace_root(root: &Path) -> Result<Self, ConfidenceStartupError> {
335        Self::start(&root.join(CONFIDENCE_THRESHOLD_RELATIVE_PATH))
336    }
337
338    pub fn threshold(&self) -> &ConfidenceThreshold {
339        &self.threshold
340    }
341
342    /// Computes confidence from ranks 1 and 2 of the canonical list, never from
343    /// the returned page interval.
344    pub fn evaluate_reply(
345        &self,
346        reply: &BlockReply,
347    ) -> Result<ConfidenceDecision, ConfidenceError> {
348        self.evaluate_canonical_list(&reply.canonical_list)
349    }
350
351    pub fn evaluate_canonical_list(
352        &self,
353        list: &CanonicalList,
354    ) -> Result<ConfidenceDecision, ConfidenceError> {
355        let mut entries = list.entries();
356        self.evaluate_pair(
357            entries
358                .next()
359                .map(|entry| entry as &dyn ConfidenceCandidate),
360            entries
361                .next()
362                .map(|entry| entry as &dyn ConfidenceCandidate),
363        )
364    }
365
366    pub fn evaluate_candidates<C: ConfidenceCandidate>(
367        &self,
368        candidates: &[C],
369    ) -> Result<ConfidenceDecision, ConfidenceError> {
370        self.evaluate_pair(
371            candidates
372                .first()
373                .map(|candidate| candidate as &dyn ConfidenceCandidate),
374            candidates
375                .get(1)
376                .map(|candidate| candidate as &dyn ConfidenceCandidate),
377        )
378    }
379
380    fn evaluate_pair(
381        &self,
382        first: Option<&dyn ConfidenceCandidate>,
383        second: Option<&dyn ConfidenceCandidate>,
384    ) -> Result<ConfidenceDecision, ConfidenceError> {
385        let Some(first) = first else {
386            return Ok(ConfidenceDecision::empty());
387        };
388        let first_evidence = first.evidence_descriptor();
389        validate_evidence(first_evidence)?;
390
391        let Some(second) = second else {
392            return Ok(if first_evidence.tier == EvidenceTier::Exact {
393                ConfidenceDecision::high(ConfidenceBranch::SingletonExact)
394            } else {
395                ConfidenceDecision::singleton_low()
396            });
397        };
398        let second_evidence = second.evidence_descriptor();
399        validate_evidence(second_evidence)?;
400
401        match (first_evidence.tier, second_evidence.tier) {
402            (EvidenceTier::Exact, EvidenceTier::NonExact) => Ok(ConfidenceDecision::high(
403                ConfidenceBranch::ExactOverNonExact,
404            )),
405            (EvidenceTier::NonExact, EvidenceTier::Exact) => Err(ConfidenceError::InvalidHeadOrder),
406            (EvidenceTier::Exact, EvidenceTier::Exact) => {
407                if compare_exact_fields_1_to_4(first_evidence, second_evidence) == Ordering::Less {
408                    Ok(ConfidenceDecision::high(ConfidenceBranch::ExactEvidence))
409                } else {
410                    Ok(ConfidenceDecision::flat_head(
411                        ConfidenceBranch::ExactEvidence,
412                    ))
413                }
414            }
415            (EvidenceTier::NonExact, EvidenceTier::NonExact) => {
416                if first_evidence.exact_form != second_evidence.exact_form
417                    || first_evidence.generated != second_evidence.generated
418                {
419                    return Ok(ConfidenceDecision::high(
420                        ConfidenceBranch::NonExactDescriptor,
421                    ));
422                }
423
424                let first_score = frozen_score(first)?;
425                let second_score = frozen_score(second)?;
426                if first_score - second_score >= self.threshold.margin {
427                    Ok(ConfidenceDecision::high(ConfidenceBranch::NonExactMargin))
428                } else {
429                    Ok(ConfidenceDecision::flat_head(
430                        ConfidenceBranch::NonExactMargin,
431                    ))
432                }
433            }
434        }
435    }
436}
437
438fn validate_evidence(evidence: &EvidenceDescriptor) -> Result<(), ConfidenceError> {
439    if evidence.is_valid_shape() {
440        Ok(())
441    } else {
442        Err(ConfidenceError::InvalidEvidenceDescriptor)
443    }
444}
445
446fn frozen_score(candidate: &dyn ConfidenceCandidate) -> Result<f32, ConfidenceError> {
447    let score = candidate
448        .frozen_fusion_score()
449        .ok_or(ConfidenceError::MissingFusionScore)?;
450    if score.is_finite() {
451        Ok(score)
452    } else {
453        Err(ConfidenceError::NonFiniteFusionScore)
454    }
455}
456
457fn compare_exact_fields_1_to_4(
458    first: &EvidenceDescriptor,
459    second: &EvidenceDescriptor,
460) -> Ordering {
461    let kind_rank = |kind| match kind {
462        EvidenceKind::Definition => 0,
463        EvidenceKind::E1 => 1,
464        EvidenceKind::Anchored => 2,
465        EvidenceKind::E2 => 3,
466        EvidenceKind::None => usize::MAX,
467    };
468
469    match kind_rank(first.kind).cmp(&kind_rank(second.kind)) {
470        Ordering::Equal => {}
471        ordering => return ordering,
472    }
473
474    match first.kind {
475        EvidenceKind::Definition => Ordering::Equal,
476        EvidenceKind::E1 => second.occurrences.cmp(&first.occurrences),
477        EvidenceKind::Anchored => second
478            .matched_span
479            .cmp(&first.matched_span)
480            .then_with(|| first.gap_chars.cmp(&second.gap_chars)),
481        EvidenceKind::E2 => first.window_lines.cmp(&second.window_lines),
482        EvidenceKind::None => Ordering::Equal,
483    }
484}
aft/commands/semantic_search/confidence.rs

aft/commands/semantic_search/
confidence.rs
