coderlib 0.1.0

//! Model definitions and utilities
//!
//! This module contains predefined model configurations for various LLM providers.

use crate::llm::{Model, ModelCapabilities};

/// Get predefined models for OpenAI
pub fn openai_models() -> Vec<Model> {
    vec![
        Model {
            id: "gpt-4".to_string(),
            name: "GPT-4".to_string(),
            provider: "openai".to_string(),
            context_length: 8192,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00003,
            cost_per_output_token: 0.00006,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "gpt-4-turbo".to_string(),
            name: "GPT-4 Turbo".to_string(),
            provider: "openai".to_string(),
            context_length: 128000,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.00001,
            cost_per_output_token: 0.00003,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "gpt-3.5-turbo".to_string(),
            name: "GPT-3.5 Turbo".to_string(),
            provider: "openai".to_string(),
            context_length: 16385,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.0000005,
            cost_per_output_token: 0.0000015,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get predefined models for Anthropic
pub fn anthropic_models() -> Vec<Model> {
    vec![
        Model {
            id: "claude-3-5-sonnet-20241022".to_string(),
            name: "Claude 3.5 Sonnet".to_string(),
            provider: "anthropic".to_string(),
            context_length: 200000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.000003,
            cost_per_output_token: 0.000015,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "claude-3-haiku-20240307".to_string(),
            name: "Claude 3 Haiku".to_string(),
            provider: "anthropic".to_string(),
            context_length: 200000,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.00000025,
            cost_per_output_token: 0.00000125,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get predefined models for Google Gemini
pub fn gemini_models() -> Vec<Model> {
    vec![
        Model {
            id: "gemini-1.5-pro".to_string(),
            name: "Gemini 1.5 Pro".to_string(),
            provider: "gemini".to_string(),
            context_length: 2000000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.00000125,
            cost_per_output_token: 0.000005,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "gemini-1.5-flash".to_string(),
            name: "Gemini 1.5 Flash".to_string(),
            provider: "gemini".to_string(),
            context_length: 1000000,
            max_output_tokens: 8192,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: true,
            cost_per_input_token: 0.000000075,
            cost_per_output_token: 0.0000003,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get Groq models
pub fn groq_models() -> Vec<Model> {
    vec![
        Model {
            id: "llama-3.1-70b-versatile".to_string(),
            name: "Llama 3.1 70B".to_string(),
            provider: "groq".to_string(),
            context_length: 131072,
            max_output_tokens: 8000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00000059,
            cost_per_output_token: 0.00000079,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "llama-3.1-8b-instant".to_string(),
            name: "Llama 3.1 8B Instant".to_string(),
            provider: "groq".to_string(),
            context_length: 131072,
            max_output_tokens: 8000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00000005,
            cost_per_output_token: 0.00000008,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "mixtral-8x7b-32768".to_string(),
            name: "Mixtral 8x7B".to_string(),
            provider: "groq".to_string(),
            context_length: 32768,
            max_output_tokens: 8000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00000024,
            cost_per_output_token: 0.00000024,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get Together AI models
pub fn together_models() -> Vec<Model> {
    vec![
        Model {
            id: "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo".to_string(),
            name: "Llama 3.1 70B Instruct Turbo".to_string(),
            provider: "together".to_string(),
            context_length: 131072,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00000088,
            cost_per_output_token: 0.00000088,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo".to_string(),
            name: "Llama 3.1 8B Instruct Turbo".to_string(),
            provider: "together".to_string(),
            context_length: 131072,
            max_output_tokens: 4096,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00000018,
            cost_per_output_token: 0.00000018,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get Cohere models
pub fn cohere_models() -> Vec<Model> {
    vec![
        Model {
            id: "command-r-plus".to_string(),
            name: "Command R+".to_string(),
            provider: "cohere".to_string(),
            context_length: 128000,
            max_output_tokens: 4000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.000003,
            cost_per_output_token: 0.000015,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "command-r".to_string(),
            name: "Command R".to_string(),
            provider: "cohere".to_string(),
            context_length: 128000,
            max_output_tokens: 4000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.0000005,
            cost_per_output_token: 0.0000015,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get SambaNova models
pub fn sambanova_models() -> Vec<Model> {
    vec![
        Model {
            id: "Meta-Llama-3.1-70B-Instruct".to_string(),
            name: "Llama 3.1 70B Instruct".to_string(),
            provider: "sambanova".to_string(),
            context_length: 131072,
            max_output_tokens: 4096,
            supports_tools: false,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.0000005,
            cost_per_output_token: 0.0000005,
            capabilities: ModelCapabilities::default(),
        },
        Model {
            id: "Meta-Llama-3.1-8B-Instruct".to_string(),
            name: "Llama 3.1 8B Instruct".to_string(),
            provider: "sambanova".to_string(),
            context_length: 131072,
            max_output_tokens: 4096,
            supports_tools: false,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.0000001,
            cost_per_output_token: 0.0000001,
            capabilities: ModelCapabilities::default(),
        },
    ]
}

/// Get all predefined models
pub fn all_models() -> Vec<Model> {
    let mut models = Vec::new();

    // Always include these providers (no feature flags needed)
    models.extend(openai_models());
    models.extend(anthropic_models());
    models.extend(groq_models());
    models.extend(together_models());
    models.extend(cohere_models());
    models.extend(sambanova_models());
    models.extend(gemini_models());

    models
}

/// Find a model by ID
pub fn find_model_by_id(id: &str) -> Option<Model> {
    all_models().into_iter().find(|model| model.id == id)
}

/// Find models by provider
pub fn find_models_by_provider(provider: &str) -> Vec<Model> {
    all_models()
        .into_iter()
        .filter(|model| model.provider == provider)
        .collect()
}

/// Get the default model for a provider
pub fn default_model_for_provider(provider: &str) -> Option<Model> {
    match provider {
        "openai" => find_model_by_id("gpt-4"),
        "anthropic" => find_model_by_id("claude-3-5-sonnet-20241022"),
        "groq" => find_model_by_id("llama-3.1-70b-versatile"),
        "together" => find_model_by_id("meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"),
        "cohere" => find_model_by_id("command-r-plus"),
        "sambanova" => find_model_by_id("Meta-Llama-3.1-70B-Instruct"),
        "perplexity" => find_model_by_id("llama-3.1-70b-versatile"), // Use Groq model as fallback
        "gemini" => find_model_by_id("gemini-1.5-pro"),
        _ => None,
    }
}

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

    #[test]
    fn test_find_model_by_id() {
        let model = find_model_by_id("gpt-4");
        assert!(model.is_some());
        assert_eq!(model.unwrap().name, "GPT-4");
        
        let missing = find_model_by_id("nonexistent");
        assert!(missing.is_none());
    }

    #[test]
    fn test_find_models_by_provider() {
        let openai_models = find_models_by_provider("openai");
        assert!(!openai_models.is_empty());
        assert!(openai_models.iter().all(|m| m.provider == "openai"));
    }

    #[test]
    fn test_default_model_for_provider() {
        let default = default_model_for_provider("openai");
        assert!(default.is_some());
        assert_eq!(default.unwrap().id, "gpt-4");
        
        let missing = default_model_for_provider("unknown");
        assert!(missing.is_none());
    }

    #[test]
    fn test_all_models() {
        let models = all_models();
        assert!(!models.is_empty());
        
        // Check that we have models from different providers
        let providers: std::collections::HashSet<_> = models.iter().map(|m| &m.provider).collect();
        assert!(providers.len() > 0);
    }
}