coderlib 0.1.0

//! LLM integration module for CoderLib
//!
//! This module provides abstractions and implementations for various LLM providers,
//! including OpenAI, Anthropic, Google Gemini, and others.

pub mod providers;
pub mod models;

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::sync::mpsc;

use crate::core::error::ProviderError;
use crate::core::TokenUsage;
use crate::storage::Message;
use crate::tools::Tool;

/// Trait for LLM providers
#[async_trait]
pub trait Provider: Send + Sync {
    /// Send messages to the LLM and get a complete response
    async fn send_messages(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<ProviderResponse, ProviderError>;
    
    /// Stream response from the LLM
    async fn stream_response(
        &self,
        messages: Vec<Message>,
        tools: Vec<Box<dyn Tool>>,
    ) -> Result<mpsc::Receiver<Result<ProviderEvent, ProviderError>>, ProviderError>;
    
    /// Get the model information for this provider
    fn model(&self) -> &Model;
    
    /// Get the provider name
    fn name(&self) -> &str;
    
    /// Check if the provider is available/configured
    async fn is_available(&self) -> bool;
}

/// Response from an LLM provider
#[derive(Debug, Clone)]
pub struct ProviderResponse {
    /// The generated content
    pub content: String,
    
    /// Any tool calls made by the LLM
    pub tool_calls: Vec<ProviderToolCall>,
    
    /// Token usage information
    pub token_usage: Option<TokenUsage>,
    
    /// Provider-specific metadata
    pub metadata: serde_json::Value,
}

/// Tool call made by an LLM provider
#[derive(Debug, Clone)]
pub struct ProviderToolCall {
    /// Unique identifier for this tool call
    pub id: String,
    
    /// Name of the tool to call
    pub name: String,
    
    /// Parameters for the tool call
    pub parameters: serde_json::Value,
}

/// Events emitted during streaming responses
#[derive(Debug, Clone)]
pub enum ProviderEvent {
    /// A chunk of content was received
    ContentChunk {
        content: String,
    },
    
    /// A tool call was requested
    ToolCall {
        name: String,
        parameters: serde_json::Value,
    },
    
    /// The response is complete
    Complete {
        token_usage: Option<TokenUsage>,
    },
    
    /// An error occurred
    Error {
        error: String,
    },
}

/// Information about an AI model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Model {
    /// Model identifier (e.g., "gpt-4", "claude-3-5-sonnet")
    pub id: String,
    
    /// Human-readable model name
    pub name: String,
    
    /// Provider that offers this model
    pub provider: String,
    
    /// Maximum context length in tokens
    pub context_length: u32,
    
    /// Maximum output tokens
    pub max_output_tokens: u32,
    
    /// Whether the model supports tool calling
    pub supports_tools: bool,
    
    /// Whether the model supports streaming
    pub supports_streaming: bool,
    
    /// Whether the model supports vision/image inputs
    pub supports_vision: bool,
    
    /// Cost per input token (in USD)
    pub cost_per_input_token: f64,
    
    /// Cost per output token (in USD)
    pub cost_per_output_token: f64,
    
    /// Model capabilities and features
    pub capabilities: ModelCapabilities,
}

/// Model capabilities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCapabilities {
    /// Programming languages the model excels at
    pub programming_languages: Vec<String>,
    
    /// Whether the model can generate code
    pub code_generation: bool,
    
    /// Whether the model can explain code
    pub code_explanation: bool,
    
    /// Whether the model can debug code
    pub code_debugging: bool,
    
    /// Whether the model can refactor code
    pub code_refactoring: bool,
    
    /// Whether the model supports function calling
    pub function_calling: bool,
    
    /// Whether the model can work with files
    pub file_operations: bool,
}

impl Default for ModelCapabilities {
    fn default() -> Self {
        Self {
            programming_languages: vec![
                "rust".to_string(),
                "python".to_string(),
                "javascript".to_string(),
                "typescript".to_string(),
                "go".to_string(),
                "java".to_string(),
                "cpp".to_string(),
                "c".to_string(),
            ],
            code_generation: true,
            code_explanation: true,
            code_debugging: true,
            code_refactoring: true,
            function_calling: true,
            file_operations: true,
        }
    }
}

impl Model {
    /// Calculate the cost for a given token usage
    pub fn calculate_cost(&self, token_usage: &TokenUsage) -> f64 {
        let input_cost = token_usage.input_tokens as f64 * self.cost_per_input_token;
        let output_cost = token_usage.output_tokens as f64 * self.cost_per_output_token;
        input_cost + output_cost
    }
    
    /// Check if the model supports a specific capability
    pub fn supports_capability(&self, capability: &str) -> bool {
        match capability {
            "tools" => self.supports_tools,
            "streaming" => self.supports_streaming,
            "vision" => self.supports_vision,
            "code_generation" => self.capabilities.code_generation,
            "code_explanation" => self.capabilities.code_explanation,
            "code_debugging" => self.capabilities.code_debugging,
            "code_refactoring" => self.capabilities.code_refactoring,
            "function_calling" => self.capabilities.function_calling,
            "file_operations" => self.capabilities.file_operations,
            _ => false,
        }
    }
    
    /// Check if the model supports a programming language
    pub fn supports_language(&self, language: &str) -> bool {
        self.capabilities.programming_languages
            .iter()
            .any(|lang| lang.eq_ignore_ascii_case(language))
    }
}

/// Factory for creating provider instances
pub struct ProviderFactory;

impl ProviderFactory {
    /// Create a provider instance based on configuration
    pub async fn create_provider(
        provider_name: &str,
        config: &crate::core::config::ProviderConfig,
    ) -> Result<Box<dyn Provider>, ProviderError> {
        match provider_name {
            "mock" => {
                Ok(Box::new(providers::MockProvider::new()))
            }
            "openai" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("OpenAI API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let base_url = config.settings.get("base_url")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string());

                let provider = providers::OpenAIProvider::new(
                    api_key,
                    model,
                    base_url,
                    config.max_tokens,
                    Some(0.7), // Default temperature
                )?;

                Ok(Box::new(provider))
            }
            "anthropic" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Anthropic API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let base_url = config.settings.get("base_url")
                    .and_then(|v| v.as_str())
                    .map(|s| s.to_string());

                let provider = providers::AnthropicProvider::new(
                    api_key,
                    model,
                    base_url,
                    config.max_tokens,
                    Some(0.7), // Default temperature
                )?;

                Ok(Box::new(provider))
            }
            "groq" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Groq API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenAICompatibleProvider::groq(api_key, model)?;
                Ok(Box::new(provider))
            }
            "cohere" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Cohere API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenAICompatibleProvider::cohere(api_key, model)?;
                Ok(Box::new(provider))
            }
            "sambanova" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("SambaNova API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenAICompatibleProvider::sambanova(api_key, model)?;
                Ok(Box::new(provider))
            }
            "together" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Together API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenAICompatibleProvider::together(api_key, model)?;
                Ok(Box::new(provider))
            }
            "perplexity" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Perplexity API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenAICompatibleProvider::perplexity(api_key, model)?;
                Ok(Box::new(provider))
            }
            "gemini" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Gemini API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::GeminiProvider::new(api_key, model)?;
                Ok(Box::new(provider))
            }
            "azure" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Azure API key required".into()))?;

                let endpoint = config.base_url.clone()
                    .ok_or_else(|| ProviderError::Configuration("Azure endpoint required".into()))?;

                let deployment_name = config.settings.get("deployment_name")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| ProviderError::Configuration("Azure deployment name required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::AzureProvider::new(
                    api_key,
                    endpoint,
                    deployment_name.to_string(),
                    config.settings.get("api_version").and_then(|v| v.as_str()).map(|s| s.to_string()),
                    model,
                    config.max_tokens,
                    config.settings.get("temperature").and_then(|v| v.as_f64()).map(|f| f as f32),
                )?;
                Ok(Box::new(provider))
            }
            "vertex" => {
                let project_id = config.settings.get("project_id")
                    .and_then(|v| v.as_str())
                    .ok_or_else(|| ProviderError::Configuration("Vertex AI project ID required".into()))?;

                let location = config.settings.get("location")
                    .and_then(|v| v.as_str())
                    .unwrap_or("us-central1");

                let access_token = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("Vertex AI access token required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::VertexProvider::new(
                    project_id.to_string(),
                    location.to_string(),
                    access_token,
                    model,
                    config.max_tokens,
                    config.settings.get("temperature").and_then(|v| v.as_f64()).map(|f| f as f32),
                )?;
                Ok(Box::new(provider))
            }
            "openrouter" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("OpenRouter API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::OpenRouterProvider::new(
                    api_key,
                    model,
                    config.max_tokens,
                    config.settings.get("temperature").and_then(|v| v.as_f64()).map(|f| f as f32),
                    config.settings.get("app_name").and_then(|v| v.as_str()).map(|s| s.to_string()),
                    config.settings.get("site_url").and_then(|v| v.as_str()).map(|s| s.to_string()),
                )?;
                Ok(Box::new(provider))
            }
            "xai" => {
                let api_key = config.api_key.clone()
                    .ok_or_else(|| ProviderError::Configuration("xAI API key required".into()))?;

                let model = models::find_model_by_id(&config.default_model)
                    .ok_or_else(|| ProviderError::ModelNotFound(config.default_model.clone()))?;

                let provider = providers::XaiProvider::new(
                    api_key,
                    model,
                    config.max_tokens,
                    config.settings.get("temperature").and_then(|v| v.as_f64()).map(|f| f as f32),
                )?;
                Ok(Box::new(provider))
            }
            "local" => {
                let local_config = providers::local::LocalProviderConfig {
                    endpoint: config.base_url.clone()
                        .unwrap_or_else(|| std::env::var("LOCAL_ENDPOINT")
                            .unwrap_or_else(|_| "http://localhost:1234".to_string())),
                    api_key: config.api_key.clone(),
                    timeout_seconds: 60,
                    discovery_paths: vec![
                        "v1/models".to_string(),
                        "api/v0/models".to_string(),
                        "api/v1/models".to_string(),
                        "v1/internal/model/list".to_string(),
                    ],
                };

                let provider = if let Some(model_id) = config.settings.get("model_id").and_then(|v| v.as_str()) {
                    providers::LocalProvider::with_model(local_config, model_id).await?
                } else {
                    providers::LocalProvider::new(local_config).await?
                };

                Ok(Box::new(provider))
            }
            _ => Err(ProviderError::ModelNotFound(format!("Unknown provider: {}", provider_name))),
        }
    }
    
    /// List all available providers
    pub fn available_providers() -> Vec<&'static str> {
        vec![
            // Always available providers
            "mock",
            "openai",
            "anthropic",
            "groq",
            "cohere",
            "sambanova",
            "together",
            "perplexity",
            "gemini",
            "azure",
            "vertex",
            "openrouter",
            "xai",
            "local",
            // "bedrock", // Requires AWS SDK setup
        ]
    }
}

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

    #[test]
    fn test_model_cost_calculation() {
        let model = Model {
            id: "test-model".to_string(),
            name: "Test Model".to_string(),
            provider: "test".to_string(),
            context_length: 4000,
            max_output_tokens: 1000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00001,
            cost_per_output_token: 0.00003,
            capabilities: ModelCapabilities::default(),
        };
        
        let token_usage = TokenUsage {
            input_tokens: 1000,
            output_tokens: 500,
            total_tokens: 1500,
            cache_creation_tokens: 0,
            cache_read_tokens: 0,
        };
        
        let cost = model.calculate_cost(&token_usage);
        assert_eq!(cost, 0.025); // (1000 * 0.00001) + (500 * 0.00003)
    }

    #[test]
    fn test_model_capabilities() {
        let model = Model {
            id: "test-model".to_string(),
            name: "Test Model".to_string(),
            provider: "test".to_string(),
            context_length: 4000,
            max_output_tokens: 1000,
            supports_tools: true,
            supports_streaming: true,
            supports_vision: false,
            cost_per_input_token: 0.00001,
            cost_per_output_token: 0.00003,
            capabilities: ModelCapabilities::default(),
        };
        
        assert!(model.supports_capability("tools"));
        assert!(model.supports_capability("streaming"));
        assert!(!model.supports_capability("vision"));
        assert!(model.supports_capability("code_generation"));
        assert!(model.supports_language("rust"));
        assert!(model.supports_language("Python")); // Case insensitive
        assert!(!model.supports_language("cobol"));
    }

    #[test]
    fn test_provider_factory_available_providers() {
        let providers = ProviderFactory::available_providers();
        // The exact providers depend on which features are enabled
        // In default configuration, we should have at least openai and anthropic
        assert!(!providers.is_empty());
    }
}