aft/commands/semantic_search/
scoring.rs

1use std::collections::BTreeMap;
2use std::fmt;
3
4use serde::{Deserialize, Serialize};
5
6use super::plan_table::{PlanTable, SearchLaneKind, SearchShape};
7
8/// One depth-limited lane's canonical contribution to a candidate.
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct LaneContribution {
11    pub lane: SearchLaneKind,
12    pub position: usize,
13    pub raw_score: f32,
14}
15
16/// A contribution admitted at the candidate's own tier depth.
17#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
18pub struct AdmittedContribution {
19    pub lane: SearchLaneKind,
20    pub position: usize,
21    pub raw_score: f32,
22}
23
24/// The fixed scoring inputs for one depth-limited lane.
25#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
26pub struct LaneScoringRule {
27    pub weight: f32,
28    pub rrf_constant: f32,
29    pub plan_order_index: usize,
30}
31
32/// Request-independent scoring rules copied from the pinned plan table.
33#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
34pub struct ScoringPolicy {
35    rules: BTreeMap<SearchLaneKind, LaneScoringRule>,
36}
37
38impl ScoringPolicy {
39    pub fn new(lexical: LaneScoringRule, semantic: LaneScoringRule) -> Result<Self, ScoringError> {
40        let mut rules = BTreeMap::new();
41        rules.insert(SearchLaneKind::Lexical, lexical);
42        rules.insert(SearchLaneKind::Semantic, semantic);
43        let policy = Self { rules };
44        policy.validate()?;
45        Ok(policy)
46    }
47
48    /// An empty policy is useful for exact-only block construction. If a
49    /// non-exact candidate reaches scoring, the missing rule is a hard error.
50    pub fn empty() -> Self {
51        Self {
52            rules: BTreeMap::new(),
53        }
54    }
55
56    pub fn from_plan_table(table: &PlanTable, shape: SearchShape) -> Result<Self, ScoringError> {
57        let shape_name = shape.as_str();
58        let lane_entries = table
59            .entries
60            .get(shape_name)
61            .ok_or_else(|| ScoringError::MissingPlanShape(shape_name.to_string()))?;
62
63        let rule = |lane: SearchLaneKind| -> Result<LaneScoringRule, ScoringError> {
64            let entry = lane_entries
65                .get(lane.as_str())
66                .ok_or(ScoringError::MissingLaneRule(lane))?;
67            let weight = entry.weight.ok_or(ScoringError::MissingWeight(lane))?;
68            let rrf_constant = entry
69                .rrf_constant
70                .ok_or(ScoringError::MissingRrfConstant(lane))?;
71            Ok(LaneScoringRule {
72                weight,
73                rrf_constant,
74                plan_order_index: entry.plan_order_index,
75            })
76        };
77
78        Self::new(
79            rule(SearchLaneKind::Lexical)?,
80            rule(SearchLaneKind::Semantic)?,
81        )
82    }
83
84    pub fn rule(&self, lane: SearchLaneKind) -> Option<LaneScoringRule> {
85        self.rules.get(&lane).copied()
86    }
87
88    fn validate(&self) -> Result<(), ScoringError> {
89        for (lane, rule) in &self.rules {
90            if *lane == SearchLaneKind::Exact {
91                return Err(ScoringError::ExactLaneHasScoringRule);
92            }
93            if !rule.weight.is_finite() || rule.weight < 0.0 {
94                return Err(ScoringError::InvalidWeight(*lane));
95            }
96            if !rule.rrf_constant.is_finite() || rule.rrf_constant < 0.0 {
97                return Err(ScoringError::InvalidRrfConstant(*lane));
98            }
99        }
100        Ok(())
101    }
102}
103
104/// Scores frozen into a non-exact block entry.
105#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
106pub struct FrozenScores {
107    pub fusion_score: f32,
108    pub lane_score: f32,
109    pub best_lane: SearchLaneKind,
110    pub admitted: Vec<AdmittedContribution>,
111}
112
113#[derive(Debug, Clone, PartialEq, Eq)]
114pub enum ScoringError {
115    EmptyAdmittedSet,
116    ExactContribution,
117    DuplicateLane(SearchLaneKind),
118    NonFiniteRawScore(SearchLaneKind),
119    MissingLaneRule(SearchLaneKind),
120    MissingPlanShape(String),
121    MissingWeight(SearchLaneKind),
122    MissingRrfConstant(SearchLaneKind),
123    InvalidWeight(SearchLaneKind),
124    InvalidRrfConstant(SearchLaneKind),
125    ExactLaneHasScoringRule,
126}
127
128impl fmt::Display for ScoringError {
129    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
130        match self {
131            Self::EmptyAdmittedSet => {
132                formatter.write_str("non-exact candidate has an empty admitted contribution set")
133            }
134            Self::ExactContribution => {
135                formatter.write_str("the exact lane is depth-exempt and cannot be scored")
136            }
137            Self::DuplicateLane(lane) => write!(
138                formatter,
139                "candidate has more than one canonical contribution from the {lane} lane"
140            ),
141            Self::NonFiniteRawScore(lane) => {
142                write!(formatter, "{lane} contribution has a non-finite raw score")
143            }
144            Self::MissingLaneRule(lane) => {
145                write!(formatter, "scoring policy has no rule for the {lane} lane")
146            }
147            Self::MissingPlanShape(shape) => {
148                write!(formatter, "plan table has no entry for shape {shape}")
149            }
150            Self::MissingWeight(lane) => write!(formatter, "{lane} lane has no scoring weight"),
151            Self::MissingRrfConstant(lane) => {
152                write!(formatter, "{lane} lane has no RRF constant")
153            }
154            Self::InvalidWeight(lane) => write!(formatter, "{lane} lane has an invalid weight"),
155            Self::InvalidRrfConstant(lane) => {
156                write!(formatter, "{lane} lane has an invalid RRF constant")
157            }
158            Self::ExactLaneHasScoringRule => {
159                formatter.write_str("the exact lane must remain score-free")
160            }
161        }
162    }
163}
164
165impl std::error::Error for ScoringError {}
166
167/// Return only contributions available before the candidate's own tier cutoff.
168///
169/// The caller may pass everything observed by a deeper request. Contributions
170/// at or beyond `tier_depth` are deliberately ignored, so they cannot change a
171/// block that was already defined at the shallower depth.
172pub fn admitted_contributions(
173    contributions: &[LaneContribution],
174    tier_depth: usize,
175) -> Result<Vec<AdmittedContribution>, ScoringError> {
176    let mut admitted = Vec::new();
177    let mut seen = BTreeMap::new();
178
179    for contribution in contributions {
180        if contribution.lane == SearchLaneKind::Exact {
181            return Err(ScoringError::ExactContribution);
182        }
183        if !contribution.raw_score.is_finite() {
184            return Err(ScoringError::NonFiniteRawScore(contribution.lane));
185        }
186        if seen.insert(contribution.lane, ()).is_some() {
187            return Err(ScoringError::DuplicateLane(contribution.lane));
188        }
189        if contribution.position < tier_depth {
190            admitted.push(AdmittedContribution {
191                lane: contribution.lane,
192                position: contribution.position,
193                raw_score: contribution.raw_score,
194            });
195        }
196    }
197
198    admitted.sort_by_key(|contribution| contribution.lane.default_plan_order_index());
199    Ok(admitted)
200}
201
202/// Compute one non-exact candidate's block-frozen fusion and lane scores.
203///
204/// Fusion uses weighted reciprocal rank from canonical zero-based positions,
205/// without normalization by retrieved-set or block size. `lane_score` is the
206/// maximum raw lane score; fixed plan-table order breaks equal-score lane ties.
207pub fn freeze_non_exact_scores(
208    contributions: &[LaneContribution],
209    tier_depth: usize,
210    policy: &ScoringPolicy,
211) -> Result<FrozenScores, ScoringError> {
212    let admitted = admitted_contributions(contributions, tier_depth)?;
213    if admitted.is_empty() {
214        return Err(ScoringError::EmptyAdmittedSet);
215    }
216
217    let mut fusion_score = 0.0f32;
218    let mut best: Option<(f32, usize, SearchLaneKind)> = None;
219
220    for contribution in &admitted {
221        let rule = policy
222            .rule(contribution.lane)
223            .ok_or(ScoringError::MissingLaneRule(contribution.lane))?;
224        fusion_score += rule.weight / (rule.rrf_constant + contribution.position as f32 + 1.0);
225
226        let proposed = (
227            contribution.raw_score,
228            rule.plan_order_index,
229            contribution.lane,
230        );
231        let replace = best.is_none_or(|current| {
232            proposed.0 > current.0 || (proposed.0 == current.0 && proposed.1 < current.1)
233        });
234        if replace {
235            best = Some(proposed);
236        }
237    }
238
239    let (lane_score, _, best_lane) = best.expect("non-empty admitted set has a best lane");
240    Ok(FrozenScores {
241        fusion_score,
242        lane_score,
243        best_lane,
244        admitted,
245    })
246}
aft/commands/semantic_search/scoring.rs

aft/commands/semantic_search/
scoring.rs
