jamjet_models/

registry.rs

//! Model registry — resolves model names to adapters.
//!
//! Supports routing rules: e.g. "claude-*" → Anthropic, "gpt-*" → OpenAI.
//! Falls back to a default adapter if no rule matches.

use crate::adapter::{ModelAdapter, ModelError, ModelRequest, ModelResponse, StructuredRequest};
use std::collections::HashMap;
use std::sync::Arc;

/// Routes model requests to the appropriate adapter.
///
/// Register adapters by `system_name()` (e.g. "anthropic", "openai").
/// The registry selects an adapter based on the model prefix in the request config,
/// or falls back to the default adapter.
pub struct ModelRegistry {
    adapters: HashMap<String, Arc<dyn ModelAdapter>>,
    /// Prefix routing: model name prefix → system name (e.g. "claude-" → "anthropic").
    prefix_routes: Vec<(String, String)>,
    default: Option<String>,
}

impl ModelRegistry {
    pub fn new() -> Self {
        Self {
            adapters: HashMap::new(),
            prefix_routes: Vec::new(),
            default: None,
        }
    }

    /// Register an adapter under its system name.
    pub fn register(mut self, adapter: Arc<dyn ModelAdapter>) -> Self {
        let name = adapter.system_name().to_string();
        self.adapters.insert(name, adapter);
        self
    }

    /// Route model name prefix to a system (e.g. "claude-" → "anthropic").
    pub fn route_prefix(mut self, prefix: impl Into<String>, system: impl Into<String>) -> Self {
        self.prefix_routes.push((prefix.into(), system.into()));
        self
    }

    /// Set the default adapter to use when no prefix matches.
    pub fn with_default(mut self, system: impl Into<String>) -> Self {
        self.default = Some(system.into());
        self
    }

    /// Resolve an adapter for the given model name.
    fn resolve(&self, model: &str) -> Option<Arc<dyn ModelAdapter>> {
        // Check prefix routes first.
        for (prefix, system) in &self.prefix_routes {
            if model.starts_with(prefix.as_str()) {
                if let Some(adapter) = self.adapters.get(system) {
                    return Some(Arc::clone(adapter));
                }
            }
        }
        // Fall back to default.
        if let Some(default) = &self.default {
            return self.adapters.get(default).map(Arc::clone);
        }
        // Only one adapter registered — use it.
        if self.adapters.len() == 1 {
            return self.adapters.values().next().map(Arc::clone);
        }
        None
    }

    /// Send a chat request, routing to the appropriate adapter.
    pub async fn chat(&self, request: ModelRequest) -> Result<ModelResponse, ModelError> {
        let model = request.config.model.clone().unwrap_or_default();
        let adapter = self
            .resolve(&model)
            .ok_or_else(|| ModelError::Network(format!("no adapter for model: {model}")))?;
        adapter.chat(request).await
    }

    /// Send a structured output request, routing to the appropriate adapter.
    pub async fn structured_output(
        &self,
        request: StructuredRequest,
    ) -> Result<ModelResponse, ModelError> {
        let model = request.config.model.clone().unwrap_or_default();
        let adapter = self
            .resolve(&model)
            .ok_or_else(|| ModelError::Network(format!("no adapter for model: {model}")))?;
        adapter.structured_output(request).await
    }
}

impl Default for ModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Build a `ModelRegistry` from environment variables.
///
/// Registers adapters based on available API keys / services:
/// - Anthropic if `ANTHROPIC_API_KEY` is set
/// - OpenAI if `OPENAI_API_KEY` is set
/// - Google if `GOOGLE_API_KEY` or `GEMINI_API_KEY` is set
/// - Ollama if `OLLAMA_HOST` is set or defaults to localhost:11434
///
/// Sets up standard prefix routing:
///   claude-* → anthropic, gpt-*/o1-*/o3-* → openai,
///   gemini-* → google, ollama model names → ollama.
pub fn registry_from_env() -> ModelRegistry {
    use crate::{
        anthropic::AnthropicAdapter, google::GoogleAdapter, ollama::OllamaAdapter,
        openai::OpenAiAdapter,
    };

    let mut registry = ModelRegistry::new()
        .route_prefix("claude-", "anthropic")
        .route_prefix("gpt-", "openai")
        .route_prefix("o1-", "openai")
        .route_prefix("o3-", "openai")
        .route_prefix("gemini-", "google")
        .route_prefix("google/", "google")
        // Common Ollama model name patterns.
        .route_prefix("llama", "ollama")
        .route_prefix("qwen", "ollama")
        .route_prefix("gemma", "ollama")
        .route_prefix("phi", "ollama")
        .route_prefix("mistral", "ollama")
        .route_prefix("codellama", "ollama")
        .route_prefix("deepseek", "ollama")
        .route_prefix("nomic-", "ollama");

    if let Ok(adapter) = AnthropicAdapter::from_env() {
        registry = registry.register(Arc::new(adapter));
        registry = registry.with_default("anthropic");
    }

    if let Ok(adapter) = OpenAiAdapter::from_env() {
        registry = registry.register(Arc::new(adapter));
        if registry.default.is_none() {
            registry = registry.with_default("openai");
        }
    }

    if let Ok(adapter) = GoogleAdapter::from_env() {
        registry = registry.register(Arc::new(adapter));
        if registry.default.is_none() {
            registry = registry.with_default("google");
        }
    }

    // Ollama is always available if the server is running (no API key needed).
    // Register it but don't set as default — cloud providers take priority.
    if let Ok(adapter) = OllamaAdapter::from_env() {
        registry = registry.register(Arc::new(adapter));
        if registry.default.is_none() {
            registry = registry.with_default("ollama");
        }
    }

    registry
}