Skip to main content

newt_core/router/
mod.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
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/// Result of [`Router::classify_detailed`]: the chosen tier, a confidence
21/// score in `[0.0, 1.0]`, and human-readable reasons explaining the decision.
22#[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    /// Create a router that always returns the given tier (confidence 1.0).
42    pub fn with_override(tier: Tier) -> Self {
43        Self {
44            tier_override: Some(tier),
45        }
46    }
47
48    /// Classify an incoming prompt. v0 heuristics: length + keyword triggers.
49    /// Refine with empirical signals before v1.
50    pub fn classify(&self, prompt: &str) -> Tier {
51        self.classify_detailed(prompt).tier
52    }
53
54    /// Like [`classify`](Self::classify), but returns confidence and reasons
55    /// alongside the tier.
56    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}