mur_core/model/

router.rs

//! Model Router — selects the appropriate AI model based on step type and role.
//!
//! Routes requests to different providers (Anthropic, Ollama, OpenRouter)
//! based on the 5 model roles: Thinking, Coding, Task, Embedding, Auditor.

use super::cli_provider::{CliProvider, CliProviderConfig};
use super::ollama::OllamaProvider;
use super::openrouter::OpenRouterProvider;
use super::provider::{AnthropicProvider, ModelResponse};
use super::{ModelConfig, ModelRole, Provider};
use crate::types::StepType;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Configuration for all model roles.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterConfig {
    pub roles: HashMap<String, ModelConfig>,
}

impl Default for RouterConfig {
    fn default() -> Self {
        let mut roles = HashMap::new();
        roles.insert(
            "thinking".into(),
            ModelConfig {
                provider: Provider::Anthropic,
                model: "claude-opus-4-20250514".into(),
                api_key_env: Some("ANTHROPIC_API_KEY".into()),
                max_cost_per_call: 5.0,
                temperature: 0.3,
                max_tokens: 8192,
            },
        );
        roles.insert(
            "coding".into(),
            ModelConfig {
                provider: Provider::Anthropic,
                model: "claude-sonnet-4-20250514".into(),
                api_key_env: Some("ANTHROPIC_API_KEY".into()),
                max_cost_per_call: 2.0,
                temperature: 0.2,
                max_tokens: 8192,
            },
        );
        roles.insert(
            "task".into(),
            ModelConfig {
                provider: Provider::OpenRouter,
                model: "google/gemini-2.0-flash-001".into(),
                api_key_env: Some("OPENROUTER_API_KEY".into()),
                max_cost_per_call: 0.5,
                temperature: 0.3,
                max_tokens: 4096,
            },
        );
        roles.insert(
            "embedding".into(),
            ModelConfig {
                provider: Provider::Ollama,
                model: "nomic-embed-text".into(),
                api_key_env: None,
                max_cost_per_call: 0.0,
                temperature: 0.0,
                max_tokens: 0,
            },
        );
        roles.insert(
            "auditor".into(),
            ModelConfig {
                provider: Provider::Ollama,
                model: "llama3.2:3b".into(),
                api_key_env: None,
                max_cost_per_call: 0.0,
                temperature: 0.1,
                max_tokens: 2048,
            },
        );
        Self { roles }
    }
}

/// Commander config file structure (subset relevant to model routing).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct CommanderFileConfig {
    #[serde(default)]
    anthropic_api_key: Option<String>,
    #[serde(default)]
    openrouter_api_key: Option<String>,
    #[serde(default)]
    ollama_url: Option<String>,
    /// AI section with CLI providers and default provider.
    #[serde(default)]
    ai: Option<AiConfig>,
}

/// AI configuration section.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
struct AiConfig {
    /// Default provider name (e.g., "claude-cli", "ollama").
    #[serde(default)]
    pub default_provider: Option<String>,
    /// CLI provider configurations.
    #[serde(default)]
    pub cli_providers: Vec<CliProviderConfig>,
}

/// Load API keys from ~/.mur/commander/config.toml if it exists.
fn load_config_file() -> CommanderFileConfig {
    let config_path = if let Some(base_dirs) = directories::BaseDirs::new() {
        base_dirs
            .home_dir()
            .join(".mur")
            .join("commander")
            .join("config.toml")
    } else {
        return CommanderFileConfig::default();
    };

    if !config_path.exists() {
        return CommanderFileConfig::default();
    }

    std::fs::read_to_string(&config_path)
        .ok()
        .and_then(|content| toml::from_str(&content).ok())
        .unwrap_or_default()
}

/// Multi-model router that dispatches to the right provider.
pub struct ModelRouter {
    config: RouterConfig,
    anthropic: Option<AnthropicProvider>,
    ollama: OllamaProvider,
    openrouter: Option<OpenRouterProvider>,
    /// CLI-based providers (claude -p, gemini, etc.).
    cli_providers: Vec<CliProvider>,
    /// Default provider name from config.
    default_provider: Option<String>,
}

impl ModelRouter {
    /// Create a router with the given config.
    /// Loads API keys from config file (~/.mur/commander/config.toml),
    /// falling back to environment variables.
    pub fn new(config: RouterConfig) -> Self {
        let file_config = load_config_file();

        // Anthropic: config file key → env var
        let anthropic_key = file_config
            .anthropic_api_key
            .filter(|k| !k.is_empty())
            .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
            .filter(|k| !k.is_empty());
        let anthropic = anthropic_key.and_then(|k| AnthropicProvider::new(Some(k)).ok());

        // OpenRouter: config file key → env var
        let openrouter_key = file_config
            .openrouter_api_key
            .filter(|k| !k.is_empty())
            .or_else(|| std::env::var("OPENROUTER_API_KEY").ok())
            .filter(|k| !k.is_empty());
        let openrouter = openrouter_key.and_then(|k| OpenRouterProvider::new(Some(k)).ok());

        // Ollama: config file URL → default localhost
        let ollama = OllamaProvider::new(file_config.ollama_url);

        // CLI providers: from config file, or auto-detect
        let (cli_providers, default_provider) = if let Some(ref ai) = file_config.ai {
            let providers = ai
                .cli_providers
                .iter()
                .map(|c| CliProvider::new(c.clone()))
                .collect();
            (providers, ai.default_provider.clone())
        } else {
            // Auto-detect available CLI tools
            (CliProvider::detect_all(), None)
        };

        Self {
            config,
            anthropic,
            ollama,
            openrouter,
            cli_providers,
            default_provider,
        }
    }

    /// Create a router with default config.
    pub fn default_router() -> Self {
        Self::new(RouterConfig::default())
    }

    /// Select the model role for a given step type.
    pub fn role_for_step(&self, step_type: &StepType) -> ModelRole {
        ModelRole::for_step_type(step_type)
    }

    /// Get the model config for a given role.
    pub fn config_for_role(&self, role: &ModelRole) -> Option<&ModelConfig> {
        let key = match role {
            ModelRole::Thinking => "thinking",
            ModelRole::Coding => "coding",
            ModelRole::Task => "task",
            ModelRole::Embedding => "embedding",
            ModelRole::Auditor => "auditor",
        };
        self.config.roles.get(key)
    }

    /// Get a CLI provider by name.
    pub fn cli_provider(&self, name: &str) -> Option<&CliProvider> {
        self.cli_providers
            .iter()
            .find(|p| p.config.name == name && p.is_available())
    }

    /// List detected CLI providers.
    pub fn detected_cli_providers(&self) -> &[CliProvider] {
        &self.cli_providers
    }

    /// Complete using a CLI provider directly.
    pub async fn complete_with_cli(&self, provider_name: &str, prompt: &str) -> Result<ModelResponse> {
        let provider = self
            .cli_provider(provider_name)
            .context(format!("CLI provider '{}' not found or not available", provider_name))?;

        let text = provider.call(prompt).await?;
        Ok(ModelResponse {
            content: text,
            model: provider_name.to_string(),
            input_tokens: 0,
            output_tokens: 0,
            cost: 0.0,
        })
    }

    /// Route a completion request to the appropriate provider.
    /// If a default CLI provider is configured and available, uses it first.
    pub async fn complete(&self, role: &ModelRole, prompt: &str) -> Result<ModelResponse> {
        // Try default CLI provider first if configured
        if let Some(ref default) = self.default_provider {
            if let Some(provider) = self.cli_provider(default) {
                match provider.call(prompt).await {
                    Ok(text) => {
                        return Ok(ModelResponse {
                            content: text,
                            model: default.clone(),
                            input_tokens: 0,
                            output_tokens: 0,
                            cost: 0.0,
                        });
                    }
                    Err(e) => {
                        tracing::warn!("CLI provider '{}' failed: {}, falling back to API", default, e);
                    }
                }
            }
        }

        let config = self
            .config_for_role(role)
            .context(format!("No config for role {:?}", role))?;

        match config.provider {
            Provider::Anthropic => {
                if let Some(provider) = self.anthropic.as_ref() {
                    provider
                        .complete(prompt, &config.model, config.temperature, config.max_tokens)
                        .await
                } else {
                    // Fallback to Ollama when Anthropic key not available
                    tracing::info!("Anthropic not available, falling back to Ollama");
                    self.ollama
                        .complete(prompt, "llama3.2:3b", config.temperature, config.max_tokens)
                        .await
                }
            }
            Provider::Ollama => {
                self.ollama
                    .complete(prompt, &config.model, config.temperature, config.max_tokens)
                    .await
            }
            Provider::OpenRouter => {
                if let Some(provider) = self.openrouter.as_ref() {
                    provider
                        .complete(prompt, &config.model, config.temperature, config.max_tokens)
                        .await
                } else {
                    // Fallback to Ollama when OpenRouter key not available
                    tracing::info!("OpenRouter not available, falling back to Ollama");
                    self.ollama
                        .complete(prompt, "llama3.2:3b", config.temperature, config.max_tokens)
                        .await
                }
            }
            Provider::OpenAI => {
                anyhow::bail!("OpenAI provider not yet implemented")
            }
        }
    }

    /// Route a step to the appropriate model and get a completion.
    pub async fn complete_for_step(
        &self,
        step_type: &StepType,
        prompt: &str,
    ) -> Result<ModelResponse> {
        let role = self.role_for_step(step_type);
        self.complete(&role, prompt).await
    }

    /// List all configured roles and their models.
    pub fn list_models(&self) -> Vec<(String, &ModelConfig)> {
        let mut models: Vec<_> = self
            .config
            .roles
            .iter()
            .map(|(k, v)| (k.clone(), v))
            .collect();
        models.sort_by(|a, b| a.0.cmp(&b.0));
        models
    }
}

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

    #[test]
    fn test_default_config() {
        let config = RouterConfig::default();
        assert_eq!(config.roles.len(), 5);
        assert!(config.roles.contains_key("thinking"));
        assert!(config.roles.contains_key("coding"));
        assert!(config.roles.contains_key("task"));
        assert!(config.roles.contains_key("embedding"));
        assert!(config.roles.contains_key("auditor"));
    }

    #[test]
    fn test_role_for_step() {
        let router = ModelRouter::new(RouterConfig::default());
        assert_eq!(
            router.role_for_step(&StepType::Analyze),
            ModelRole::Thinking
        );
        assert_eq!(router.role_for_step(&StepType::Code), ModelRole::Coding);
        assert_eq!(router.role_for_step(&StepType::Search), ModelRole::Task);
        assert_eq!(
            router.role_for_step(&StepType::SecurityCheck),
            ModelRole::Auditor
        );
    }

    #[test]
    fn test_config_for_role() {
        let router = ModelRouter::new(RouterConfig::default());
        let thinking = router.config_for_role(&ModelRole::Thinking).unwrap();
        assert!(thinking.model.contains("opus"));

        let coding = router.config_for_role(&ModelRole::Coding).unwrap();
        assert!(coding.model.contains("sonnet"));
    }

    #[test]
    fn test_list_models() {
        let router = ModelRouter::new(RouterConfig::default());
        let models = router.list_models();
        assert_eq!(models.len(), 5);
        // Should be sorted alphabetically
        assert_eq!(models[0].0, "auditor");
        assert_eq!(models[4].0, "thinking");
    }
}