Skip to main content

newt_core/
router.rs

1//! Tier-based router — NeMoCode inheritance.
2//!
3//! Heuristics are deliberately small in v0. The classifier returns a `Tier`;
4//! the inference layer then asks each registered backend `supports_tier(t)`
5//! and picks the first match in config order.
6
7use 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/// Result of [`Router::classify_detailed`]: the chosen tier, a confidence
19/// score in `[0.0, 1.0]`, and human-readable reasons explaining the decision.
20#[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    /// Create a router that always returns the given tier (confidence 1.0).
40    pub fn with_override(tier: Tier) -> Self {
41        Self {
42            tier_override: Some(tier),
43        }
44    }
45
46    /// Classify an incoming prompt. v0 heuristics: length + keyword triggers.
47    /// Refine with empirical signals before v1.
48    pub fn classify(&self, prompt: &str) -> Tier {
49        self.classify_detailed(prompt).tier
50    }
51
52    /// Like [`classify`](Self::classify), but returns confidence and reasons
53    /// alongside the tier.
54    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}