1use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10#[serde(rename_all = "UPPERCASE")]
11pub enum Tier {
12 Fast,
13 Standard,
14 Complex,
15 Review,
16}
17
18#[derive(Debug, Clone)]
21pub struct Classification {
22 pub tier: Tier,
23 pub confidence: f64,
24 pub reasons: Vec<String>,
25}
26
27#[derive(Debug, Default)]
28pub struct Router {
29 tier_override: Option<Tier>,
30}
31
32impl Router {
33 pub fn new() -> Self {
34 Self {
35 tier_override: None,
36 }
37 }
38
39 pub fn with_override(tier: Tier) -> Self {
41 Self {
42 tier_override: Some(tier),
43 }
44 }
45
46 pub fn classify(&self, prompt: &str) -> Tier {
49 self.classify_detailed(prompt).tier
50 }
51
52 pub fn classify_detailed(&self, prompt: &str) -> Classification {
55 if let Some(tier) = self.tier_override {
56 return Classification {
57 tier,
58 confidence: 1.0,
59 reasons: vec!["tier override active".to_string()],
60 };
61 }
62
63 let len = prompt.len();
64 let lower = prompt.to_ascii_lowercase();
65
66 if lower.contains("review") || lower.contains("grade") || lower.contains("critique") {
67 return Classification {
68 tier: Tier::Review,
69 confidence: 0.85,
70 reasons: vec!["review keyword detected".to_string()],
71 };
72 }
73 if lower.contains("refactor") || lower.contains("redesign") || lower.contains("architect") {
74 return Classification {
75 tier: Tier::Complex,
76 confidence: 0.80,
77 reasons: vec!["complex keyword detected".to_string()],
78 };
79 }
80 if len < 200 {
81 return Classification {
82 tier: Tier::Fast,
83 confidence: 0.70,
84 reasons: vec![format!("short prompt ({len} chars < 200)")],
85 };
86 }
87 Classification {
88 tier: Tier::Standard,
89 confidence: 0.5,
90 reasons: vec!["no keyword match, length >= 200".to_string()],
91 }
92 }
93}
94
95#[cfg(test)]
96mod tests {
97 use super::*;
98
99 #[test]
100 fn short_prompt_is_fast() {
101 assert_eq!(Router::new().classify("rename foo to bar"), Tier::Fast);
102 }
103
104 #[test]
105 fn review_keyword_routes_review() {
106 assert_eq!(Router::new().classify("review this diff"), Tier::Review);
107 }
108
109 #[test]
110 fn refactor_keyword_routes_complex() {
111 assert_eq!(
112 Router::new().classify("refactor the auth middleware to use traits"),
113 Tier::Complex
114 );
115 }
116
117 #[test]
118 fn classify_detailed_short_has_confidence() {
119 let c = Router::new().classify_detailed("fix typo");
120 assert_eq!(c.tier, Tier::Fast);
121 assert!(c.confidence >= 0.5, "expected confidence >= 0.5");
122 }
123
124 #[test]
125 fn classify_detailed_review_has_reasons() {
126 let c = Router::new().classify_detailed("review this PR");
127 assert!(!c.reasons.is_empty(), "reasons should be non-empty");
128 assert!(
129 c.reasons.iter().any(|r| r.contains("review")),
130 "reasons should mention 'review'"
131 );
132 }
133
134 #[test]
135 fn classify_detailed_confidence_bounded() {
136 for prompt in &["hi", "review me", "refactor everything", &"x".repeat(300)] {
137 let c = Router::new().classify_detailed(prompt);
138 assert!(
139 (0.0..=1.0).contains(&c.confidence),
140 "confidence {} out of [0,1] for {:?}",
141 c.confidence,
142 c.tier
143 );
144 }
145 }
146
147 #[test]
148 fn override_always_returns_its_tier() {
149 assert_eq!(
150 Router::with_override(Tier::Complex).classify("short"),
151 Tier::Complex
152 );
153 }
154
155 #[test]
156 fn override_has_full_confidence() {
157 let c = Router::with_override(Tier::Fast).classify_detailed("anything");
158 assert!(
159 (c.confidence - 1.0).abs() < f64::EPSILON,
160 "override confidence should be 1.0"
161 );
162 }
163
164 #[test]
165 fn override_reason_mentions_override() {
166 let c = Router::with_override(Tier::Review).classify_detailed("whatever");
167 assert!(
168 c.reasons.iter().any(|r| r.contains("override")),
169 "reasons should mention 'override'"
170 );
171 }
172}