Skip to main content

_diffctx/
mode.rs

1use crate::config::limits::PPR;
2use crate::config::mode::mode as mode_config;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum ScoringMode {
6    Ppr,
7    Ego,
8    Bm25,
9    Rrf,
10    Pit,
11}
12
13impl ScoringMode {
14    pub fn from_str(s: &str) -> Result<Self, String> {
15        match s.to_lowercase().as_str() {
16            "ppr" => Ok(Self::Ppr),
17            "ego" => Ok(Self::Ego),
18            "bm25" => Ok(Self::Bm25),
19            "rrf" => Ok(Self::Rrf),
20            "pit" => Ok(Self::Pit),
21            other => Err(format!(
22                "unknown scoring_mode '{other}': expected one of {}",
23                SCORING_MODE_NAMES.join("|")
24            )),
25        }
26    }
27}
28
29/// The single source of truth for what `--scoring` accepts.
30///
31/// Both CLIs enumerate the accepted values for their own argument parsers, and
32/// both silently kept their own copy: `pit` was reachable through the engine and
33/// the eval harness while `diffctx --scoring pit` rejected it as invalid. A test
34/// pins this array against `from_str` in both directions.
35pub const SCORING_MODE_NAMES: &[&str] = &["ppr", "ego", "bm25", "rrf", "pit"];
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum DiscoveryKind {
39    Default,
40    Ensemble,
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum ScoringKind {
45    Ppr,
46    Ego,
47    Bm25,
48    Rrf,
49    Pit,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum ObjectiveMode {
54    Submodular,
55    BoltzmannModular,
56}
57
58impl ObjectiveMode {
59    pub fn from_str(s: &str) -> Self {
60        match s.to_lowercase().as_str() {
61            "boltzmann" | "boltzmann_modular" | "modular_boltzmann" => Self::BoltzmannModular,
62            _ => Self::Submodular,
63        }
64    }
65}
66
67#[derive(Debug, Clone)]
68pub struct PipelineConfig {
69    pub discovery: DiscoveryKind,
70    pub scoring: ScoringKind,
71    pub objective: ObjectiveMode,
72    pub bm25_top_k: usize,
73    pub ego_depth: usize,
74    pub ppr_alpha: f64,
75}
76
77impl PipelineConfig {
78    pub fn from_mode(mode: ScoringMode) -> Self {
79        let m = mode_config();
80        match mode {
81            ScoringMode::Ppr => Self {
82                discovery: DiscoveryKind::Ensemble,
83                scoring: ScoringKind::Ppr,
84                bm25_top_k: m.bm25_top_k_primary,
85                ego_depth: m.ego_depth_default,
86                ppr_alpha: PPR.alpha,
87                objective: ObjectiveMode::Submodular,
88            },
89            ScoringMode::Ego => Self {
90                discovery: DiscoveryKind::Ensemble,
91                scoring: ScoringKind::Ego,
92                bm25_top_k: m.bm25_top_k_primary,
93                ego_depth: m.ego_depth_extended,
94                ppr_alpha: PPR.alpha,
95                objective: ObjectiveMode::Submodular,
96            },
97            ScoringMode::Bm25 => Self {
98                discovery: DiscoveryKind::Ensemble,
99                scoring: ScoringKind::Bm25,
100                bm25_top_k: m.bm25_top_k_off,
101                ego_depth: m.ego_depth_default,
102                ppr_alpha: PPR.alpha,
103                objective: ObjectiveMode::Submodular,
104            },
105            // Discovery is deliberately identical to EGO's: the fusion gain
106            // must come from ranking the same universe better, not from a
107            // wider universe that would confound it with a discovery change.
108            ScoringMode::Rrf => Self {
109                discovery: DiscoveryKind::Ensemble,
110                scoring: ScoringKind::Rrf,
111                bm25_top_k: m.bm25_top_k_primary,
112                ego_depth: m.ego_depth_extended,
113                ppr_alpha: PPR.alpha,
114                objective: ObjectiveMode::Submodular,
115            },
116            // Same universe and the same depth as Rrf, so the two fusion modes
117            // differ only in how the two signals are combined. Comparing PIT
118            // against RRF is then a comparison of blends, not of candidate
119            // supply.
120            ScoringMode::Pit => Self {
121                discovery: DiscoveryKind::Ensemble,
122                scoring: ScoringKind::Pit,
123                bm25_top_k: m.bm25_top_k_primary,
124                ego_depth: m.ego_depth_extended,
125                ppr_alpha: PPR.alpha,
126                objective: ObjectiveMode::Submodular,
127            },
128        }
129    }
130}
131
132#[cfg(test)]
133mod scoring_mode_name_tests {
134    use super::{SCORING_MODE_NAMES, ScoringMode};
135
136    #[test]
137    fn every_advertised_name_parses() {
138        for name in SCORING_MODE_NAMES {
139            assert!(
140                ScoringMode::from_str(name).is_ok(),
141                "advertised scoring mode does not parse: {name}"
142            );
143        }
144    }
145
146    /// The other direction, which is the one that actually broke: a mode added
147    /// to `from_str` but not to the advertised list is invisible to `--scoring`.
148    #[test]
149    fn every_parsable_mode_is_advertised() {
150        for candidate in ["ppr", "ego", "bm25", "rrf", "pit"] {
151            if ScoringMode::from_str(candidate).is_ok() {
152                assert!(
153                    SCORING_MODE_NAMES.contains(&candidate),
154                    "{candidate} parses but is not advertised to the CLIs"
155                );
156            }
157        }
158    }
159}