diff --git a/src/lib.rs b/src/lib.rs index 8700f32..f161cf8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,5 @@ //! BrokeCrew: A Rust-Native Multi-Model AI Team Library -pub const VERSION: &str = env!("CARGO_PKG_VERSION"); +pub mod persona; + +pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier}; diff --git a/src/persona.rs b/src/persona.rs new file mode 100644 index 0000000..9390016 --- /dev/null +++ b/src/persona.rs @@ -0,0 +1,157 @@ +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 } + } + + 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()); + } +}