herolib-ai 0.3.13

//! AI model definitions.
//!
//! This module defines the available AI models and their provider mappings.

use serde::{Deserialize, Serialize};

use crate::provider::Provider;

/// Available AI models with our own naming convention.
///
/// Each model maps to one or more providers, tried in order.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Model {
    // Fast, small models for quick tasks
    /// Llama 3.3 70B - Fast, capable model for general tasks.
    Llama3_3_70B,
    /// Llama 3.1 70B - Versatile model for various tasks.
    Llama3_1_70B,
    /// Llama 3.1 8B - Small, fast model for simple tasks.
    Llama3_1_8B,

    // Llama 4 models (latest generation)
    /// Llama 4 Maverick 17B - Latest Llama with 128 experts.
    Llama4Maverick17B,
    /// Llama 4 Scout 17B - Efficient Llama 4 with 16 experts.
    Llama4Scout17B,

    // Coding-focused models
    /// Qwen 2.5 Coder 32B - Specialized for code generation.
    Qwen2_5Coder32B,
    /// Qwen 3 32B - Latest Qwen model with strong reasoning.
    Qwen3_32B,
    /// DeepSeek Coder V2.5 - Advanced coding model.
    DeepSeekCoderV2_5,
    /// DeepSeek V3 - Latest DeepSeek model.
    DeepSeekV3,

    // Large reasoning models
    /// Llama 3.1 405B - Largest Llama model for complex tasks.
    Llama3_1_405B,
    /// Kimi K2 - Moonshot's 256K context model with agentic coding.
    KimiK2,

    // Mixture of Experts models
    /// Mixtral 8x7B - Efficient mixture of experts model.
    Mixtral8x7B,

    // Vision models
    /// Llama 3.2 90B Vision - Multimodal model with vision.
    Llama3_2_90BVision,
    /// Llama 3.2 11B Vision - Smaller vision model.
    Llama3_2_11BVision,

    // NVIDIA models
    /// Nemotron 3 Nano 30B - NVIDIA MoE model with reasoning support.
    NemotronNano30B,

    // OpenAI models
    /// GPT-OSS 120B - OpenAI's open-weight 120B MoE model with reasoning and tool use.
    GptOss120B,
    /// GPT-OSS 20B - Smaller, faster OpenAI open-weight model.
    GptOss20B,
}

/// Model information including provider mappings.
#[derive(Debug, Clone)]
pub struct ModelInfo {
    /// Our internal model name.
    pub model: Model,
    /// Human-readable description.
    pub description: &'static str,
    /// Context window size in tokens.
    pub context_window: usize,
    /// Provider mappings in order of preference.
    pub providers: Vec<ProviderMapping>,
}

/// Pricing information for a model (per million tokens).
#[derive(Debug, Clone, Copy)]
pub struct ModelPricing {
    /// Input price per million tokens in USD.
    pub input_per_million: f32,
    /// Output price per million tokens in USD.
    pub output_per_million: f32,
}

impl ModelPricing {
    /// Creates new pricing info.
    pub const fn new(input: f32, output: f32) -> Self {
        Self {
            input_per_million: input,
            output_per_million: output,
        }
    }

    /// Returns a formatted price string like "$0.05/$0.08"
    pub fn format(&self) -> String {
        format!("${:.2}/${:.2}", self.input_per_million, self.output_per_million)
    }
}

/// Mapping of a model to a specific provider.
#[derive(Debug, Clone)]
pub struct ProviderMapping {
    /// The provider.
    pub provider: Provider,
    /// The model name/ID used by this provider.
    pub model_id: &'static str,
    /// Pricing for this provider (if known).
    pub pricing: Option<ModelPricing>,
}

impl ProviderMapping {
    /// Creates a new provider mapping without pricing.
    pub const fn new(provider: Provider, model_id: &'static str) -> Self {
        Self { provider, model_id, pricing: None }
    }

    /// Creates a new provider mapping with pricing.
    pub const fn with_pricing(provider: Provider, model_id: &'static str, input: f32, output: f32) -> Self {
        Self {
            provider,
            model_id,
            pricing: Some(ModelPricing::new(input, output)),
        }
    }
}

impl Model {
    /// Returns the model information.
    pub fn info(&self) -> ModelInfo {
        match self {
            Model::Llama3_3_70B => ModelInfo {
                model: *self,
                description: "Llama 3.3 70B - Fast, capable model for general tasks",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "llama-3.3-70b-versatile", 0.59, 0.79),
                    ProviderMapping::new(Provider::SambaNova, "Meta-Llama-3.3-70B-Instruct"),
                    ProviderMapping::new(Provider::OpenRouter, "meta-llama/llama-3.3-70b-instruct"),
                ],
            },
            Model::Llama3_1_70B => ModelInfo {
                model: *self,
                description: "Llama 3.1 70B - Versatile model for various tasks",
                context_window: 131_072,
                providers: vec![
                    // Note: llama-3.1-70b-versatile deprecated on Groq as of 2025
                    ProviderMapping::new(Provider::SambaNova, "Meta-Llama-3.1-70B-Instruct"),
                    ProviderMapping::new(Provider::OpenRouter, "meta-llama/llama-3.1-70b-instruct"),
                ],
            },
            Model::Llama3_1_8B => ModelInfo {
                model: *self,
                description: "Llama 3.1 8B - Small, fast model for simple tasks",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "llama-3.1-8b-instant", 0.05, 0.08),
                    ProviderMapping::new(Provider::SambaNova, "Meta-Llama-3.1-8B-Instruct"),
                    ProviderMapping::new(Provider::OpenRouter, "meta-llama/llama-3.1-8b-instruct"),
                ],
            },
            Model::Llama4Maverick17B => ModelInfo {
                model: *self,
                description: "Llama 4 Maverick 17B - Latest Llama with 128 experts, strong reasoning",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "meta-llama/llama-4-maverick-17b-128e-instruct", 0.20, 0.60),
                ],
            },
            Model::Llama4Scout17B => ModelInfo {
                model: *self,
                description: "Llama 4 Scout 17B - Efficient Llama 4 with 16 experts, good balance",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "meta-llama/llama-4-scout-17b-16e-instruct", 0.11, 0.34),
                ],
            },
            Model::Qwen2_5Coder32B => ModelInfo {
                model: *self,
                description: "Qwen 2.5 Coder 32B - Specialized for code generation",
                context_window: 32_000,
                providers: vec![
                    // Note: qwen-2.5-coder-32b deprecated on Groq as of 2025
                    ProviderMapping::new(Provider::SambaNova, "Qwen2.5-Coder-32B-Instruct"),
                    ProviderMapping::new(Provider::OpenRouter, "qwen/qwen-2.5-coder-32b-instruct"),
                ],
            },
            Model::Qwen3_32B => ModelInfo {
                model: *self,
                description: "Qwen 3 32B - Latest Qwen model with strong reasoning",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "qwen/qwen3-32b", 0.29, 0.59),
                ],
            },
            Model::DeepSeekCoderV2_5 => ModelInfo {
                model: *self,
                description: "DeepSeek Coder V2.5 - Advanced coding model",
                context_window: 128_000,
                providers: vec![
                    ProviderMapping::new(Provider::OpenRouter, "deepseek/deepseek-coder"),
                    ProviderMapping::new(Provider::SambaNova, "DeepSeek-Coder-V2-Instruct"),
                ],
            },
            Model::DeepSeekV3 => ModelInfo {
                model: *self,
                description: "DeepSeek V3 - Latest DeepSeek model",
                context_window: 128_000,
                providers: vec![
                    ProviderMapping::new(Provider::OpenRouter, "deepseek/deepseek-chat"),
                    ProviderMapping::new(Provider::SambaNova, "DeepSeek-V3"),
                ],
            },
            Model::Llama3_1_405B => ModelInfo {
                model: *self,
                description: "Llama 3.1 405B - Largest Llama model for complex tasks",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::new(Provider::SambaNova, "Meta-Llama-3.1-405B-Instruct"),
                    ProviderMapping::new(
                        Provider::OpenRouter,
                        "meta-llama/llama-3.1-405b-instruct",
                    ),
                ],
            },
            Model::KimiK2 => ModelInfo {
                model: *self,
                description: "Kimi K2 - Moonshot's 256K context model with agentic coding",
                context_window: 262_144,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "moonshotai/kimi-k2-instruct-0905", 1.00, 3.00),
                ],
            },
            Model::Mixtral8x7B => ModelInfo {
                model: *self,
                description: "Mixtral 8x7B - Efficient mixture of experts model",
                context_window: 32_000,
                providers: vec![
                    // Note: mixtral-8x7b-32768 deprecated on Groq as of 2025
                    ProviderMapping::new(Provider::OpenRouter, "mistralai/mixtral-8x7b-instruct"),
                ],
            },
            Model::Llama3_2_90BVision => ModelInfo {
                model: *self,
                description: "Llama 3.2 90B Vision - Multimodal model with vision",
                context_window: 128_000,
                providers: vec![
                    // Note: llama-3.2-90b-vision-preview deprecated on Groq as of 2025
                    ProviderMapping::new(
                        Provider::OpenRouter,
                        "meta-llama/llama-3.2-90b-vision-instruct",
                    ),
                ],
            },
            Model::Llama3_2_11BVision => ModelInfo {
                model: *self,
                description: "Llama 3.2 11B Vision - Smaller vision model",
                context_window: 128_000,
                providers: vec![
                    // Note: llama-3.2-11b-vision-preview deprecated on Groq as of 2025
                    ProviderMapping::new(Provider::SambaNova, "Llama-3.2-11B-Vision-Instruct"),
                    ProviderMapping::new(
                        Provider::OpenRouter,
                        "meta-llama/llama-3.2-11b-vision-instruct",
                    ),
                ],
            },
            Model::NemotronNano30B => ModelInfo {
                model: *self,
                description: "Nemotron 3 Nano 30B - NVIDIA MoE model with reasoning support",
                context_window: 262_144,
                providers: vec![ProviderMapping::new(
                    Provider::OpenRouter,
                    "nvidia/nemotron-3-nano-30b-a3b",
                )],
            },
            Model::GptOss120B => ModelInfo {
                model: *self,
                description: "GPT-OSS 120B - OpenAI's open-weight 120B MoE model with reasoning and tool use",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "openai/gpt-oss-120b", 0.15, 0.60),
                    ProviderMapping::new(Provider::SambaNova, "gpt-oss-120b"),
                    ProviderMapping::new(Provider::OpenRouter, "openai/gpt-oss-120b"),
                ],
            },
            Model::GptOss20B => ModelInfo {
                model: *self,
                description: "GPT-OSS 20B - Smaller, faster OpenAI open-weight model",
                context_window: 131_072,
                providers: vec![
                    ProviderMapping::with_pricing(Provider::Groq, "openai/gpt-oss-20b", 0.075, 0.30),
                ],
            },
        }
    }

    /// Returns the human-readable name.
    pub fn name(&self) -> &'static str {
        match self {
            Model::Llama3_3_70B => "Llama 3.3 70B",
            Model::Llama3_1_70B => "Llama 3.1 70B",
            Model::Llama3_1_8B => "Llama 3.1 8B",
            Model::Llama4Maverick17B => "Llama 4 Maverick 17B",
            Model::Llama4Scout17B => "Llama 4 Scout 17B",
            Model::Qwen2_5Coder32B => "Qwen 2.5 Coder 32B",
            Model::Qwen3_32B => "Qwen 3 32B",
            Model::DeepSeekCoderV2_5 => "DeepSeek Coder V2.5",
            Model::DeepSeekV3 => "DeepSeek V3",
            Model::Llama3_1_405B => "Llama 3.1 405B",
            Model::KimiK2 => "Kimi K2",
            Model::Mixtral8x7B => "Mixtral 8x7B",
            Model::Llama3_2_90BVision => "Llama 3.2 90B Vision",
            Model::Llama3_2_11BVision => "Llama 3.2 11B Vision",
            Model::NemotronNano30B => "Nemotron 3 Nano 30B",
            Model::GptOss120B => "GPT-OSS 120B",
            Model::GptOss20B => "GPT-OSS 20B",
        }
    }

    /// Returns the default model for general tasks.
    pub fn default_general() -> Self {
        Model::Llama3_3_70B
    }

    /// Returns the default model for coding tasks.
    pub fn default_coding() -> Self {
        Model::Qwen2_5Coder32B
    }

    /// Returns all available models.
    pub fn all() -> &'static [Model] {
        &[
            Model::Llama3_3_70B,
            Model::Llama3_1_70B,
            Model::Llama3_1_8B,
            Model::Llama4Maverick17B,
            Model::Llama4Scout17B,
            Model::Qwen2_5Coder32B,
            Model::Qwen3_32B,
            Model::DeepSeekCoderV2_5,
            Model::DeepSeekV3,
            Model::Llama3_1_405B,
            Model::KimiK2,
            Model::Mixtral8x7B,
            Model::Llama3_2_90BVision,
            Model::Llama3_2_11BVision,
            Model::NemotronNano30B,
            Model::GptOss120B,
            Model::GptOss20B,
        ]
    }

    /// Returns models that are available for the given providers.
    ///
    /// A model is available if at least one of its provider mappings
    /// matches one of the configured providers.
    pub fn available_for_providers(providers: &[Provider]) -> Vec<Model> {
        Self::all()
            .iter()
            .filter(|model| {
                let info = model.info();
                info.providers.iter().any(|mapping| providers.contains(&mapping.provider))
            })
            .copied()
            .collect()
    }

    /// Returns the API-friendly ID for this model (lowercase, hyphenated).
    pub fn id(&self) -> &'static str {
        match self {
            Model::Llama3_3_70B => "llama3.3-70b",
            Model::Llama3_1_70B => "llama3.1-70b",
            Model::Llama3_1_8B => "llama3.1-8b",
            Model::Llama4Maverick17B => "llama4-maverick-17b",
            Model::Llama4Scout17B => "llama4-scout-17b",
            Model::Qwen2_5Coder32B => "qwen2.5-coder-32b",
            Model::Qwen3_32B => "qwen3-32b",
            Model::DeepSeekCoderV2_5 => "deepseek-coder-v2.5",
            Model::DeepSeekV3 => "deepseek-v3",
            Model::Llama3_1_405B => "llama3.1-405b",
            Model::KimiK2 => "kimi-k2",
            Model::Mixtral8x7B => "mixtral-8x7b",
            Model::Llama3_2_90BVision => "llama3.2-90b-vision",
            Model::Llama3_2_11BVision => "llama3.2-11b-vision",
            Model::NemotronNano30B => "nemotron-nano-30b",
            Model::GptOss120B => "gpt-oss-120b",
            Model::GptOss20B => "gpt-oss-20b",
        }
    }

    /// Parse a model from its API ID.
    pub fn from_id(id: &str) -> Option<Model> {
        match id.to_lowercase().as_str() {
            "llama3.3-70b" | "llama-3.3-70b" | "llama3_3_70b" => Some(Model::Llama3_3_70B),
            "llama3.1-70b" | "llama-3.1-70b" | "llama3_1_70b" => Some(Model::Llama3_1_70B),
            "llama3.1-8b" | "llama-3.1-8b" | "llama3_1_8b" => Some(Model::Llama3_1_8B),
            "llama4-maverick-17b" | "llama-4-maverick" => Some(Model::Llama4Maverick17B),
            "llama4-scout-17b" | "llama-4-scout" => Some(Model::Llama4Scout17B),
            "llama3.1-405b" | "llama-3.1-405b" | "llama3_1_405b" => Some(Model::Llama3_1_405B),
            "qwen2.5-coder-32b" | "qwen-2.5-coder-32b" | "qwen2_5coder32b" => Some(Model::Qwen2_5Coder32B),
            "qwen3-32b" | "qwen-3-32b" => Some(Model::Qwen3_32B),
            "deepseek-coder-v2.5" | "deepseek-coder" => Some(Model::DeepSeekCoderV2_5),
            "deepseek-v3" | "deepseek" => Some(Model::DeepSeekV3),
            "kimi-k2" | "kimi" => Some(Model::KimiK2),
            "mixtral-8x7b" | "mixtral8x7b" => Some(Model::Mixtral8x7B),
            "llama3.2-90b-vision" | "llama-3.2-90b-vision" => Some(Model::Llama3_2_90BVision),
            "llama3.2-11b-vision" | "llama-3.2-11b-vision" => Some(Model::Llama3_2_11BVision),
            "nemotron-nano-30b" | "nemotron-3-nano-30b" => Some(Model::NemotronNano30B),
            "gpt-oss-120b" | "gpt-oss" => Some(Model::GptOss120B),
            "gpt-oss-20b" => Some(Model::GptOss20B),
            _ => None,
        }
    }
}

impl ModelInfo {
    /// Returns the best (cheapest) pricing available for this model.
    pub fn best_pricing(&self) -> Option<ModelPricing> {
        self.providers
            .iter()
            .filter_map(|m| m.pricing)
            .min_by(|a, b| {
                let cost_a = a.input_per_million + a.output_per_million;
                let cost_b = b.input_per_million + b.output_per_million;
                cost_a.partial_cmp(&cost_b).unwrap_or(std::cmp::Ordering::Equal)
            })
    }

    /// Returns pricing for a specific provider, if available.
    pub fn pricing_for_provider(&self, provider: Provider) -> Option<ModelPricing> {
        self.providers
            .iter()
            .find(|m| m.provider == provider)
            .and_then(|m| m.pricing)
    }
}

impl std::fmt::Display for Model {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_model_info() {
        let info = Model::Llama3_3_70B.info();
        assert!(!info.providers.is_empty());
        assert!(info.context_window > 0);
    }

    #[test]
    fn test_all_models_have_providers() {
        for model in Model::all() {
            let info = model.info();
            assert!(
                !info.providers.is_empty(),
                "Model {} has no providers",
                model.name()
            );
        }
    }

    #[test]
    fn test_default_models() {
        assert_eq!(Model::default_general(), Model::Llama3_3_70B);
        assert_eq!(Model::default_coding(), Model::Qwen2_5Coder32B);
    }
}