1use serde::{Deserialize, Serialize};
8
9mod classifier;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
12#[serde(rename_all = "UPPERCASE")]
13pub enum Tier {
14 Fast,
15 Standard,
16 Complex,
17 Review,
18}
19
20#[derive(Debug, Clone)]
23pub struct Classification {
24 pub tier: Tier,
25 pub confidence: f64,
26 pub reasons: Vec<String>,
27}
28
29#[derive(Debug, Default)]
30pub struct Router {
31 tier_override: Option<Tier>,
32}
33
34impl Router {
35 pub fn new() -> Self {
36 Self {
37 tier_override: None,
38 }
39 }
40
41 pub fn with_override(tier: Tier) -> Self {
43 Self {
44 tier_override: Some(tier),
45 }
46 }
47
48 pub fn classify(&self, prompt: &str) -> Tier {
51 self.classify_detailed(prompt).tier
52 }
53
54 pub fn classify_detailed(&self, prompt: &str) -> Classification {
57 classifier::classify_detailed(prompt, self.tier_override)
58 }
59}
60
61#[cfg(test)]
62mod tests {
63 use super::*;
64
65 #[test]
66 fn short_prompt_is_fast() {
67 assert_eq!(Router::new().classify("rename foo to bar"), Tier::Fast);
68 }
69
70 #[test]
71 fn review_keyword_routes_review() {
72 assert_eq!(Router::new().classify("review this diff"), Tier::Review);
73 }
74
75 #[test]
76 fn refactor_keyword_routes_complex() {
77 assert_eq!(
78 Router::new().classify("refactor the auth middleware to use traits"),
79 Tier::Complex
80 );
81 }
82
83 #[test]
84 fn classify_detailed_short_has_confidence() {
85 let c = Router::new().classify_detailed("fix typo");
86 assert_eq!(c.tier, Tier::Fast);
87 assert!(c.confidence >= 0.5, "expected confidence >= 0.5");
88 }
89
90 #[test]
91 fn classify_detailed_review_has_reasons() {
92 let c = Router::new().classify_detailed("review this PR");
93 assert!(!c.reasons.is_empty(), "reasons should be non-empty");
94 assert!(
95 c.reasons.iter().any(|r| r.contains("review")),
96 "reasons should mention 'review'"
97 );
98 }
99
100 #[test]
101 fn classify_detailed_confidence_bounded() {
102 for prompt in &["hi", "review me", "refactor everything", &"x".repeat(300)] {
103 let c = Router::new().classify_detailed(prompt);
104 assert!(
105 (0.0..=1.0).contains(&c.confidence),
106 "confidence {} out of [0,1] for {:?}",
107 c.confidence,
108 c.tier
109 );
110 }
111 }
112
113 #[test]
114 fn override_always_returns_its_tier() {
115 assert_eq!(
116 Router::with_override(Tier::Complex).classify("short"),
117 Tier::Complex
118 );
119 }
120
121 #[test]
122 fn override_has_full_confidence() {
123 let c = Router::with_override(Tier::Fast).classify_detailed("anything");
124 assert!(
125 (c.confidence - 1.0).abs() < f64::EPSILON,
126 "override confidence should be 1.0"
127 );
128 }
129
130 #[test]
131 fn override_reason_mentions_override() {
132 let c = Router::with_override(Tier::Review).classify_detailed("whatever");
133 assert!(
134 c.reasons.iter().any(|r| r.contains("override")),
135 "reasons should mention 'override'"
136 );
137 }
138}