feat: implement two-stage router with keyword filter and complexity escalation

src/lib.rs

@@ -4,8 +4,10 @@ pub mod budget;
 pub mod client;
 pub mod persona;
 pub mod registry;
+pub mod router;
 
 pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
 pub use client::{ClientError, ModelClient, OpenRouterClient};
 pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};
 pub use registry::{ModelInfo, ModelPricing, ModelRegistry};
+pub use router::{KeywordRouter, LlmRouter, RouteDecision, RouteStage, Router};

src/router/keyword.rs

@@ -0,0 +1,110 @@
+use crate::persona::Persona;
+
+pub struct KeywordRouter;
+
+impl KeywordRouter {
+    /// Stage A: Free, instant pre-filtering against persona trigger words.
+    /// Uses word boundary scanning to avoid false substring positives.
+    pub fn classify<'a>(personas: &'a [Persona], prompt: &str) -> Option<&'a str> {
+        let lower = prompt.to_lowercase();
+        // Tokenize prompt into words for fast lookup
+        let words: Vec<&str> = lower
+            .split(|c: char| !c.is_alphanumeric() && c != '_' && c != '-')
+            .filter(|w| !w.is_empty())
+            .collect();
+
+        if words.is_empty() {
+            return None;
+        }
+
+        let mut best_match: Option<(&'a str, usize)> = None;
+        let mut tie = false;
+
+        for persona in personas {
+            let mut score = 0;
+            for trigger in &persona.trigger_words {
+                let trig_lower = trigger.to_lowercase();
+                // Check exact word token match or multi-word phrase match
+                if trig_lower.contains(' ') {
+                    if lower.contains(&trig_lower) {
+                        score += 3;
+                    }
+                } else if words.iter().any(|&w| w == trig_lower) {
+                    score += 1;
+                }
+            }
+
+            if score > 0 {
+                match best_match {
+                    None => {
+                        best_match = Some((persona.name.as_str(), score));
+                        tie = false;
+                    }
+                    Some((_, best_score)) if score > best_score => {
+                        best_match = Some((persona.name.as_str(), score));
+                        tie = false;
+                    }
+                    Some((_, best_score)) if score == best_score => {
+                        tie = true;
+                    }
+                    _ => {}
+                }
+            }
+        }
+
+        // Return persona only if unambiguous winner found
+        if tie {
+            None
+        } else {
+            best_match.map(|(name, _)| name)
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_test_personas() -> Vec<Persona> {
+        vec![
+            Persona {
+                name: "coder".into(),
+                system_prompt: "coder prompt".into(),
+                trigger_words: vec!["build".into(), "code".into(), "script".into(), "website".into(), "api".into()],
+                candidates: vec![],
+            },
+            Persona {
+                name: "reviewer".into(),
+                system_prompt: "reviewer prompt".into(),
+                trigger_words: vec!["review".into(), "audit".into(), "critique".into(), "lint".into()],
+                candidates: vec![],
+            },
+            Persona {
+                name: "architect".into(),
+                system_prompt: "architect prompt".into(),
+                trigger_words: vec!["architecture".into(), "design".into(), "scale".into()],
+                candidates: vec![],
+            },
+        ]
+    }
+
+    #[test]
+    fn test_stage_a_keyword_classification() {
+        let personas = make_test_personas();
+
+        let cases = vec![
+            ("build me a proxy checker site", Some("coder")),
+            ("write a python script to parse logs", Some("coder")),
+            ("please review this pull request for bugs", Some("reviewer")),
+            ("audit our authentication system", Some("reviewer")),
+            ("design a high scale distributed cache architecture", Some("architect")),
+            ("hello, what is the weather today?", None),
+            ("how are you doing?", None),
+        ];
+
+        for (prompt, expected) in cases {
+            let matched = KeywordRouter::classify(&personas, prompt);
+            assert_eq!(matched, expected, "Prompt failed: '{}'", prompt);
+        }
+    }
+}

src/router/llm.rs

@@ -0,0 +1,152 @@
+use crate::client::{ClientError, ModelClient};
+
+pub struct LlmRouter;
+
+impl LlmRouter {
+    pub async fn classify(
+        client: &dyn ModelClient,
+        router_model: &str,
+        known_personas: &[String],
+        default_persona: &str,
+        prompt: &str,
+    ) -> Result<String, ClientError> {
+        let system_prompt = format!(
+            "You are a strict single-token persona router. Output ONLY the persona name from this list: [{}]. Do not output explanations, markdown, or punctuation.",
+            known_personas.join(", ")
+        );
+
+        let user_prompt = format!("Classify this request to the most appropriate persona:\n\"{}\"", prompt);
+
+        let raw_output = match client.complete(router_model, &system_prompt, &user_prompt).await {
+            Ok(out) => out,
+            Err(e) => {
+                tracing::warn!("Router LLM call failed, falling back to default persona: {e}");
+                return Ok(default_persona.to_string());
+            }
+        };
+
+        let cleaned = Self::sanitize_output(&raw_output);
+        if known_personas.iter().any(|p| p.eq_ignore_ascii_case(&cleaned)) {
+            // Find matched casing
+            if let Some(matched) = known_personas.iter().find(|p| p.eq_ignore_ascii_case(&cleaned)) {
+                return Ok(matched.clone());
+            }
+        }
+
+        tracing::info!(
+            "Router model returned unrecognized persona '{}', defaulting to '{}'",
+            raw_output.trim(),
+            default_persona
+        );
+        Ok(default_persona.to_string())
+    }
+
+    pub fn sanitize_output(output: &str) -> String {
+        // Strip markdown, punctuation, and extra whitespace
+        output
+            .trim()
+            .trim_matches(|c: char| !c.is_alphanumeric() && c != '_' && c != '-')
+            .to_lowercase()
+    }
+
+    pub fn estimate_complexity(prompt: &str) -> u8 {
+        let mut score = 1u8;
+        if prompt.len() > 800 {
+            score += 3;
+        } else if prompt.len() > 300 {
+            score += 1;
+        }
+
+        if prompt.contains("```") {
+            score += 2;
+        }
+
+        let heavy_keywords = [
+            "multi-file",
+            "entire project",
+            "refactor all",
+            "microservice",
+            "database migration",
+            "distributed",
+            "architecture",
+            "production ready",
+        ];
+        let lower = prompt.to_lowercase();
+        for kw in heavy_keywords {
+            if lower.contains(kw) {
+                score += 2;
+            }
+        }
+
+        score.min(10)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use async_trait::async_trait;
+    use tokio::sync::mpsc;
+
+    struct MockRouterClient {
+        reply: String,
+    }
+
+    #[async_trait]
+    impl ModelClient for MockRouterClient {
+        async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result<String, ClientError> {
+            Ok(self.reply.clone())
+        }
+
+        async fn complete_stream(
+            &self,
+            _model: &str,
+            _system: &str,
+            _user: &str,
+            _chunk_sender: mpsc::Sender<String>,
+        ) -> Result<String, ClientError> {
+            Ok(self.reply.clone())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_llm_router_valid_and_fallback_parsing() {
+        let known = vec!["coder".into(), "reviewer".into(), "general".into()];
+
+        // 1. Direct valid persona
+        let client_valid = MockRouterClient {
+            reply: "coder".into(),
+        };
+        let res = LlmRouter::classify(&client_valid, "router-model", &known, "general", "build something")
+            .await
+            .unwrap();
+        assert_eq!(res, "coder");
+
+        // 2. Extra whitespace & quotes
+        let client_messy = MockRouterClient {
+            reply: "  \"REVIEWER\" \n".into(),
+        };
+        let res = LlmRouter::classify(&client_messy, "router-model", &known, "general", "audit this")
+            .await
+            .unwrap();
+        assert_eq!(res, "reviewer");
+
+        // 3. Hallucinated / unrecognized persona defaults to default_persona
+        let client_hallucinated = MockRouterClient {
+            reply: "I believe a poet would be best for this query".into(),
+        };
+        let res = LlmRouter::classify(&client_hallucinated, "router-model", &known, "general", "write poem")
+            .await
+            .unwrap();
+        assert_eq!(res, "general");
+    }
+
+    #[test]
+    fn test_complexity_estimation() {
+        let simple_prompt = "hello what is 2 + 2?";
+        assert!(LlmRouter::estimate_complexity(simple_prompt) < 7);
+
+        let complex_prompt = "Refactor all services into a distributed microservice architecture with database migration and tests:\n```rust\nfn main() {}\n```";
+        assert!(LlmRouter::estimate_complexity(complex_prompt) >= 7);
+    }
+}

src/router/mod.rs

@@ -0,0 +1,135 @@
+pub mod keyword;
+pub mod llm;
+
+pub use keyword::KeywordRouter;
+pub use llm::LlmRouter;
+
+use crate::client::{ClientError, ModelClient};
+use crate::persona::PersonaRegistry;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RouteStage {
+    StageAKeyword,
+    StageBLlm,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct RouteDecision {
+    pub persona: String,
+    pub stage_used: RouteStage,
+    pub escalate_to_paid: bool,
+    pub complexity_score: u8,
+}
+
+pub struct Router {
+    pub default_persona: String,
+    pub router_model: String,
+    pub complexity_threshold: u8,
+}
+
+impl Router {
+    pub fn new(default_persona: String, router_model: String, complexity_threshold: u8) -> Self {
+        Self {
+            default_persona,
+            router_model,
+            complexity_threshold,
+        }
+    }
+
+    pub async fn route(
+        &self,
+        registry: &PersonaRegistry,
+        client: &dyn ModelClient,
+        prompt: &str,
+        force_escalate: bool,
+    ) -> Result<RouteDecision, ClientError> {
+        let complexity_score = LlmRouter::estimate_complexity(prompt);
+        let escalate_to_paid = force_escalate || complexity_score >= self.complexity_threshold;
+
+        // Stage A: Free, instant keyword matching
+        if let Some(matched) = KeywordRouter::classify(registry.list(), prompt) {
+            return Ok(RouteDecision {
+                persona: matched.to_string(),
+                stage_used: RouteStage::StageAKeyword,
+                escalate_to_paid,
+                complexity_score,
+            });
+        }
+
+        // Stage B: LLM classification fallback
+        let known: Vec<String> = registry.list().iter().map(|p| p.name.clone()).collect();
+        let selected = LlmRouter::classify(
+            client,
+            &self.router_model,
+            &known,
+            &self.default_persona,
+            prompt,
+        )
+        .await?;
+
+        Ok(RouteDecision {
+            persona: selected,
+            stage_used: RouteStage::StageBLlm,
+            escalate_to_paid,
+            complexity_score,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::persona::Persona;
+    use async_trait::async_trait;
+    use tokio::sync::mpsc;
+
+    struct DummyClient;
+
+    #[async_trait]
+    impl ModelClient for DummyClient {
+        async fn complete(&self, _model: &str, _system: &str, _user: &str) -> Result<String, ClientError> {
+            Ok("general".into())
+        }
+        async fn complete_stream(
+            &self,
+            _model: &str,
+            _system: &str,
+            _user: &str,
+            _chunk_sender: mpsc::Sender<String>,
+        ) -> Result<String, ClientError> {
+            Ok("general".into())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_combined_router_stage_a_and_escalate() {
+        let personas = vec![
+            Persona {
+                name: "coder".into(),
+                system_prompt: "coder".into(),
+                trigger_words: vec!["code".into(), "build".into()],
+                candidates: vec![],
+            },
+            Persona {
+                name: "general".into(),
+                system_prompt: "general".into(),
+                trigger_words: vec!["hello".into()],
+                candidates: vec![],
+            },
+        ];
+        let registry = PersonaRegistry::new(personas);
+        let router = Router::new("general".into(), "meta-llama/llama-3.2-3b-instruct:free".into(), 7);
+        let client = DummyClient;
+
+        // Matches Stage A (build) without escalation
+        let dec1 = router.route(&registry, &client, "build a script", false).await.unwrap();
+        assert_eq!(dec1.persona, "coder");
+        assert_eq!(dec1.stage_used, RouteStage::StageAKeyword);
+        assert!(!dec1.escalate_to_paid);
+
+        // Matches Stage A with forced escalation
+        let dec2 = router.route(&registry, &client, "build a script", true).await.unwrap();
+        assert_eq!(dec2.persona, "coder");
+        assert!(dec2.escalate_to_paid);
+    }
+}