use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct ModelRecommendation {
recommendations: HashMap<String, String>,
}
impl ModelRecommendation {
pub fn with_defaults() -> Self {
let mut map = HashMap::new();
map.insert("simple".to_string(), "claude-haiku-4-5-20251001".to_string());
map.insert("fast".to_string(), "claude-haiku-4-5-20251001".to_string());
map.insert("cheap".to_string(), "claude-haiku-4-5-20251001".to_string());
map.insert("quick".to_string(), "claude-haiku-4-5-20251001".to_string());
map.insert("balanced".to_string(), "claude-sonnet-4-5-20250929".to_string());
map.insert("general".to_string(), "claude-sonnet-4-5-20250929".to_string());
map.insert("normal".to_string(), "claude-sonnet-4-5-20250929".to_string());
map.insert("standard".to_string(), "claude-sonnet-4-5-20250929".to_string());
map.insert("latest".to_string(), "claude-sonnet-4-5-20250929".to_string());
map.insert("complex".to_string(), "claude-opus-4-6".to_string());
map.insert("best".to_string(), "claude-opus-4-6".to_string());
map.insert("quality".to_string(), "claude-opus-4-6".to_string());
map.insert("critical".to_string(), "claude-opus-4-6".to_string());
map.insert("advanced".to_string(), "claude-opus-4-6".to_string());
Self { recommendations: map }
}
pub fn custom(recommendations: HashMap<String, String>) -> Self {
Self { recommendations }
}
pub fn suggest(&self, task_type: &str) -> Option<&str> {
self.recommendations.get(task_type).map(|s| s.as_str())
}
pub fn add(&mut self, task_type: impl Into<String>, model: impl Into<String>) {
self.recommendations.insert(task_type.into(), model.into());
}
pub fn remove(&mut self, task_type: &str) -> Option<String> {
self.recommendations.remove(task_type)
}
pub fn task_types(&self) -> Vec<&str> {
self.recommendations.keys().map(|s| s.as_str()).collect()
}
pub fn all_recommendations(&self) -> &HashMap<String, String> {
&self.recommendations
}
}
impl Default for ModelRecommendation {
fn default() -> Self {
ModelRecommendation::with_defaults()
}
}
pub fn cheapest_model() -> &'static str {
"claude-haiku-4-5-20251001"
}
pub fn balanced_model() -> &'static str {
"claude-sonnet-4-5-20250929"
}
pub fn latest_sonnet() -> &'static str {
"claude-sonnet-4-5-20250929"
}
pub fn best_model() -> &'static str {
"claude-opus-4-6"
}
pub fn estimate_cost_multiplier(model: &str) -> f64 {
match model {
"haiku"
| "claude-haiku-4-5-20251001" | "claude-3-5-haiku-20241022" => 1.0,
"sonnet"
| "claude-sonnet-4-5-20250929" | "claude-sonnet-4-20250514" | "claude-3-5-sonnet-20241022" => 5.0,
"opus"
| "claude-opus-4-6" | "claude-opus-4-1-20250805" | "claude-opus-4-20250514" => 15.0,
_ => 5.0,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_recommendations() {
let recommender = ModelRecommendation::default();
assert_eq!(recommender.suggest("simple"), Some("claude-haiku-4-5-20251001"));
assert_eq!(recommender.suggest("fast"), Some("claude-haiku-4-5-20251001"));
assert_eq!(recommender.suggest("balanced"), Some("claude-sonnet-4-5-20250929"));
assert_eq!(recommender.suggest("latest"), Some("claude-sonnet-4-5-20250929"));
assert_eq!(recommender.suggest("complex"), Some("claude-opus-4-6"));
assert_eq!(recommender.suggest("unknown"), None);
}
#[test]
fn test_custom_recommendations() {
let mut map = HashMap::new();
map.insert("code_review".to_string(), "sonnet".to_string());
let recommender = ModelRecommendation::custom(map);
assert_eq!(recommender.suggest("code_review"), Some("sonnet"));
}
#[test]
fn test_add_remove() {
let mut recommender = ModelRecommendation::default();
recommender.add("my_task", "sonnet");
assert_eq!(recommender.suggest("my_task"), Some("sonnet"));
recommender.remove("my_task");
assert_eq!(recommender.suggest("my_task"), None);
}
#[test]
fn test_cost_multipliers() {
assert_eq!(estimate_cost_multiplier("haiku"), 1.0);
assert_eq!(estimate_cost_multiplier("sonnet"), 5.0);
assert_eq!(estimate_cost_multiplier("opus"), 15.0);
assert_eq!(estimate_cost_multiplier("claude-haiku-4-5-20251001"), 1.0);
assert_eq!(estimate_cost_multiplier("claude-sonnet-4-5-20250929"), 5.0);
assert_eq!(estimate_cost_multiplier("claude-opus-4-6"), 15.0);
}
#[test]
fn test_quick_helpers() {
assert_eq!(cheapest_model(), "claude-haiku-4-5-20251001");
assert_eq!(balanced_model(), "claude-sonnet-4-5-20250929");
assert_eq!(latest_sonnet(), "claude-sonnet-4-5-20250929");
assert_eq!(best_model(), "claude-opus-4-6");
}
}