From d9fb777867e8499bc7b9580a38c0654a598cad3c Mon Sep 17 00:00:00 2001 From: oiupoyt Date: Sat, 26 Sep 2026 13:49:16 +0530 Subject: [PATCH] feat: implement two-stage router with keyword filter and complexity escalation --- src/lib.rs | 2 + src/router/keyword.rs | 110 ++++++++++++++++++++++++++++++ src/router/llm.rs | 152 ++++++++++++++++++++++++++++++++++++++++++ src/router/mod.rs | 135 +++++++++++++++++++++++++++++++++++++ 4 files changed, 399 insertions(+) create mode 100644 src/router/keyword.rs create mode 100644 src/router/llm.rs create mode 100644 src/router/mod.rs diff --git a/src/lib.rs b/src/lib.rs index 6f63ef9..6728615 100644 --- a/src/lib.rs +++ b/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}; diff --git a/src/router/keyword.rs b/src/router/keyword.rs new file mode 100644 index 0000000..9548c46 --- /dev/null +++ b/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 { + 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); + } + } +} diff --git a/src/router/llm.rs b/src/router/llm.rs new file mode 100644 index 0000000..585b81b --- /dev/null +++ b/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 { + 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 { + Ok(self.reply.clone()) + } + + async fn complete_stream( + &self, + _model: &str, + _system: &str, + _user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + 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); + } +} diff --git a/src/router/mod.rs b/src/router/mod.rs new file mode 100644 index 0000000..df8cf5a --- /dev/null +++ b/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 { + 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 = 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 { + Ok("general".into()) + } + async fn complete_stream( + &self, + _model: &str, + _system: &str, + _user: &str, + _chunk_sender: mpsc::Sender, + ) -> Result { + 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(®istry, &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(®istry, &client, "build a script", true).await.unwrap(); + assert_eq!(dec2.persona, "coder"); + assert!(dec2.escalate_to_paid); + } +}