use serde::{Deserialize, Serialize}; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum Tier { Free, Paid, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct ModelCandidate { pub id: String, pub tier: Tier, #[serde(default)] pub is_stealth: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Persona { pub name: String, pub system_prompt: String, #[serde(default)] pub trigger_words: Vec, pub candidates: Vec, } #[derive(Debug, Clone)] pub struct ModelAttempt { pub model_id: String, pub error: Option, } #[derive(Debug, Clone)] pub struct DispatchResult { pub persona: String, pub model_used: String, pub tier_used: Tier, pub reply: String, pub fell_back: bool, pub attempts: Vec, } #[derive(Debug, Clone, Default)] pub struct PersonaRegistry { personas: Vec, } impl PersonaRegistry { pub fn new(personas: Vec) -> Self { Self { personas } } } impl From> for PersonaRegistry { fn from(personas: Vec) -> Self { Self::new(personas) } } impl PersonaRegistry { pub fn get(&self, name: &str) -> Option<&Persona> { self.personas.iter().find(|p| p.name == name) } pub fn list(&self) -> &[Persona] { &self.personas } /// Returns the ordered list of eligible model candidates for a persona. /// /// - If `allow_stealth` is false, excludes candidates where `is_stealth == true`. /// - If `require_paid` is true (e.g. complexity escalation), filters for `Tier::Paid` candidates. pub fn get_eligible_candidates( &self, persona_name: &str, allow_stealth: bool, require_paid: bool, ) -> Result, String> { let persona = self .get(persona_name) .ok_or_else(|| format!("Persona '{persona_name}' not found"))?; let eligible: Vec<&ModelCandidate> = persona .candidates .iter() .filter(|c| { if !allow_stealth && c.is_stealth { return false; } if require_paid && c.tier != Tier::Paid { return false; } true }) .collect(); if eligible.is_empty() { Err(format!( "No eligible model candidates for persona '{persona_name}' (stealth: {allow_stealth}, paid: {require_paid})" )) } else { Ok(eligible) } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_candidate_filtering_and_fallback_order() { let candidates = vec![ ModelCandidate { id: "stealth-model:free".into(), tier: Tier::Free, is_stealth: true, }, ModelCandidate { id: "qwen-free".into(), tier: Tier::Free, is_stealth: false, }, ModelCandidate { id: "gemini-paid".into(), tier: Tier::Paid, is_stealth: false, }, ]; let persona = Persona { name: "coder".into(), system_prompt: "You are a coder".into(), trigger_words: vec!["code".into()], candidates, }; let registry = PersonaRegistry::new(vec![persona]); // When allow_stealth is false and require_paid is false let free_no_stealth = registry .get_eligible_candidates("coder", false, false) .expect("should return candidates"); assert_eq!(free_no_stealth.len(), 2); assert_eq!(free_no_stealth[0].id, "qwen-free"); assert_eq!(free_no_stealth[1].id, "gemini-paid"); // When allow_stealth is true let with_stealth = registry .get_eligible_candidates("coder", true, false) .expect("should return candidates"); assert_eq!(with_stealth.len(), 3); assert_eq!(with_stealth[0].id, "stealth-model:free"); // When require_paid is true (complexity escalation) let paid_only = registry .get_eligible_candidates("coder", true, true) .expect("should return candidates"); assert_eq!(paid_only.len(), 1); assert_eq!(paid_only[0].id, "gemini-paid"); // Non-existent persona returns error assert!(registry .get_eligible_candidates("unknown", false, false) .is_err()); } }