brokecrew/src/router/mod.rs

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,
    router_model: std::sync::RwLock<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: std::sync::RwLock::new(router_model),
            complexity_threshold,
        }
    }

    pub fn set_router_model(&self, model: String) {
        if let Ok(mut write) = self.router_model.write() {
            *write = model;
        }
    }

    pub fn get_router_model(&self) -> String {
        self.router_model
            .read()
            .map(|m| m.clone())
            .unwrap_or_default()
    }

    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 router_model = self.get_router_model();
        let selected =
            LlmRouter::classify(client, &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);
    }
}