razeri-plugin-api 0.22.0

// Copyright (C) 2026 Gabriel Shea
// Licensed under either of the Apache License, Version 2.0 or the MIT
// license, at your option.

//! LLM model adapter trait and registry.
//!
//! [`LlmModelAdapter`] decouples model-specific capabilities and routing
//! from provider transport code. [`ModelRegistry`] stores adapters keyed
//! by (provider, `model_id`) pairs.

use std::collections::HashMap;
use std::sync::Arc;

// ── Model capabilities ────────────────────────────────────────────────

/// Machine-readable capability metadata for a concrete model.
#[derive(Debug, Clone)]
pub struct ModelCapabilities {
    /// Maximum context window in tokens.
    pub max_context_tokens: u32,
    /// Whether the model supports structured JSON output.
    pub supports_json_mode: bool,
    /// Whether the model supports tool/function calling.
    pub supports_tools: bool,
    /// Whether the model supports streaming responses.
    pub supports_streaming: bool,
}

// ── LlmModelAdapter trait ─────────────────────────────────────────────

/// Adapter for model-specific behavior and capabilities under a provider.
///
/// Separates model identity/capabilities from transport so that new models
/// can be added without changing the provider client code.
pub trait LlmModelAdapter: Send + Sync {
    /// Unique model identifier (e.g., "llama3.2:3b", "gpt-4o").
    fn model_id(&self) -> &str;

    /// Provider name this model belongs to (e.g., "ollama", "openai").
    fn provider_name(&self) -> &str;

    /// Human-readable display name.
    fn display_name(&self) -> &str;

    /// Capability metadata for this model.
    fn capabilities(&self) -> ModelCapabilities;
}

// ── ModelRegistry ─────────────────────────────────────────────────────

/// Registry for model adapters keyed by (provider, `model_id`) pairs.
pub struct ModelRegistry {
    adapters: HashMap<(String, String), Arc<dyn LlmModelAdapter>>,
}

impl ModelRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self { adapters: HashMap::new() }
    }

    /// Register a model adapter, keyed by its (`provider_name`, `model_id`).
    pub fn register(&mut self, adapter: Arc<dyn LlmModelAdapter>) {
        let key = (adapter.provider_name().to_string(), adapter.model_id().to_string());
        self.adapters.insert(key, adapter);
    }

    /// Look up an adapter by provider and model ID.
    #[must_use]
    pub fn get(&self, provider_name: &str, model_id: &str) -> Option<&Arc<dyn LlmModelAdapter>> {
        self.adapters.get(&(provider_name.to_string(), model_id.to_string()))
    }

    /// List all registered adapters for a specific provider.
    #[must_use]
    pub fn list_for_provider(&self, provider_name: &str) -> Vec<Arc<dyn LlmModelAdapter>> {
        self.adapters
            .iter()
            .filter(|((p, _), _)| p == provider_name)
            .map(|(_, adapter)| Arc::clone(adapter))
            .collect()
    }
}

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

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

    struct TestModelAdapter {
        provider: &'static str,
        model: &'static str,
        display: &'static str,
        caps: ModelCapabilities,
    }

    impl LlmModelAdapter for TestModelAdapter {
        fn model_id(&self) -> &str {
            self.model
        }

        fn provider_name(&self) -> &str {
            self.provider
        }

        fn display_name(&self) -> &str {
            self.display
        }

        fn capabilities(&self) -> ModelCapabilities {
            self.caps.clone()
        }
    }

    fn default_caps() -> ModelCapabilities {
        ModelCapabilities {
            max_context_tokens: 8192,
            supports_json_mode: true,
            supports_tools: false,
            supports_streaming: true,
        }
    }

    fn make_adapter(
        provider: &'static str,
        model: &'static str,
        display: &'static str,
    ) -> Arc<dyn LlmModelAdapter> {
        Arc::new(TestModelAdapter { provider, model, display, caps: default_caps() })
    }

    #[test]
    fn model_adapter_capabilities_exposed() {
        let adapter = make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B");
        let caps = adapter.capabilities();
        assert_eq!(caps.max_context_tokens, 8192);
        assert!(caps.supports_json_mode);
        assert!(!caps.supports_tools);
        assert!(caps.supports_streaming);
    }

    #[test]
    fn model_registry_provider_model_lookup() {
        let mut registry = ModelRegistry::new();
        registry.register(make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B"));

        assert!(registry.get("ollama", "llama3.2:3b").is_some());
        assert!(registry.get("ollama", "nonexistent").is_none());
        assert!(registry.get("openai", "llama3.2:3b").is_none());
    }

    #[test]
    fn model_registry_lists_models_by_provider() {
        let mut registry = ModelRegistry::new();
        registry.register(make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B"));
        registry.register(make_adapter("ollama", "qwen2.5:7b", "Qwen 2.5 7B"));
        registry.register(make_adapter("openai", "gpt-4o", "GPT-4o"));

        let ollama_models = registry.list_for_provider("ollama");
        assert_eq!(ollama_models.len(), 2);

        let openai_models = registry.list_for_provider("openai");
        assert_eq!(openai_models.len(), 1);
    }
}