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
29pub 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 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 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 #[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}