razeri_plugin_api/
model.rs

1// Copyright (C) 2026 Gabriel Shea <gshea@ideapivot.com> - All Rights Reserved
2
3//! LLM model adapter trait and registry.
4//!
5//! [`LlmModelAdapter`] decouples model-specific capabilities and routing
6//! from provider transport code. [`ModelRegistry`] stores adapters keyed
7//! by (provider, `model_id`) pairs.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12// ── Model capabilities ────────────────────────────────────────────────
13
14/// Machine-readable capability metadata for a concrete model.
15#[derive(Debug, Clone)]
16pub struct ModelCapabilities {
17    /// Maximum context window in tokens.
18    pub max_context_tokens: u32,
19    /// Whether the model supports structured JSON output.
20    pub supports_json_mode: bool,
21    /// Whether the model supports tool/function calling.
22    pub supports_tools: bool,
23    /// Whether the model supports streaming responses.
24    pub supports_streaming: bool,
25}
26
27// ── LlmModelAdapter trait ─────────────────────────────────────────────
28
29/// Adapter for model-specific behavior and capabilities under a provider.
30///
31/// Separates model identity/capabilities from transport so that new models
32/// can be added without changing the provider client code.
33pub trait LlmModelAdapter: Send + Sync {
34    /// Unique model identifier (e.g., "llama3.2:3b", "gpt-4o").
35    fn model_id(&self) -> &str;
36
37    /// Provider name this model belongs to (e.g., "ollama", "openai").
38    fn provider_name(&self) -> &str;
39
40    /// Human-readable display name.
41    fn display_name(&self) -> &str;
42
43    /// Capability metadata for this model.
44    fn capabilities(&self) -> ModelCapabilities;
45}
46
47// ── ModelRegistry ─────────────────────────────────────────────────────
48
49/// Registry for model adapters keyed by (provider, `model_id`) pairs.
50pub struct ModelRegistry {
51    adapters: HashMap<(String, String), Arc<dyn LlmModelAdapter>>,
52}
53
54impl ModelRegistry {
55    /// Create an empty registry.
56    #[must_use]
57    pub fn new() -> Self {
58        Self { adapters: HashMap::new() }
59    }
60
61    /// Register a model adapter, keyed by its (`provider_name`, `model_id`).
62    pub fn register(&mut self, adapter: Arc<dyn LlmModelAdapter>) {
63        let key = (adapter.provider_name().to_string(), adapter.model_id().to_string());
64        self.adapters.insert(key, adapter);
65    }
66
67    /// Look up an adapter by provider and model ID.
68    #[must_use]
69    pub fn get(&self, provider_name: &str, model_id: &str) -> Option<&Arc<dyn LlmModelAdapter>> {
70        self.adapters.get(&(provider_name.to_string(), model_id.to_string()))
71    }
72
73    /// List all registered adapters for a specific provider.
74    #[must_use]
75    pub fn list_for_provider(&self, provider_name: &str) -> Vec<Arc<dyn LlmModelAdapter>> {
76        self.adapters
77            .iter()
78            .filter(|((p, _), _)| p == provider_name)
79            .map(|(_, adapter)| Arc::clone(adapter))
80            .collect()
81    }
82}
83
84impl Default for ModelRegistry {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    struct TestModelAdapter {
95        provider: &'static str,
96        model: &'static str,
97        display: &'static str,
98        caps: ModelCapabilities,
99    }
100
101    impl LlmModelAdapter for TestModelAdapter {
102        fn model_id(&self) -> &str {
103            self.model
104        }
105
106        fn provider_name(&self) -> &str {
107            self.provider
108        }
109
110        fn display_name(&self) -> &str {
111            self.display
112        }
113
114        fn capabilities(&self) -> ModelCapabilities {
115            self.caps.clone()
116        }
117    }
118
119    fn default_caps() -> ModelCapabilities {
120        ModelCapabilities {
121            max_context_tokens: 8192,
122            supports_json_mode: true,
123            supports_tools: false,
124            supports_streaming: true,
125        }
126    }
127
128    fn make_adapter(
129        provider: &'static str,
130        model: &'static str,
131        display: &'static str,
132    ) -> Arc<dyn LlmModelAdapter> {
133        Arc::new(TestModelAdapter { provider, model, display, caps: default_caps() })
134    }
135
136    #[test]
137    fn model_adapter_capabilities_exposed() {
138        let adapter = make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B");
139        let caps = adapter.capabilities();
140        assert_eq!(caps.max_context_tokens, 8192);
141        assert!(caps.supports_json_mode);
142        assert!(!caps.supports_tools);
143        assert!(caps.supports_streaming);
144    }
145
146    #[test]
147    fn model_registry_provider_model_lookup() {
148        let mut registry = ModelRegistry::new();
149        registry.register(make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B"));
150
151        assert!(registry.get("ollama", "llama3.2:3b").is_some());
152        assert!(registry.get("ollama", "nonexistent").is_none());
153        assert!(registry.get("openai", "llama3.2:3b").is_none());
154    }
155
156    #[test]
157    fn model_registry_lists_models_by_provider() {
158        let mut registry = ModelRegistry::new();
159        registry.register(make_adapter("ollama", "llama3.2:3b", "Llama 3.2 3B"));
160        registry.register(make_adapter("ollama", "qwen2.5:7b", "Qwen 2.5 7B"));
161        registry.register(make_adapter("openai", "gpt-4o", "GPT-4o"));
162
163        let ollama_models = registry.list_for_provider("ollama");
164        assert_eq!(ollama_models.len(), 2);
165
166        let openai_models = registry.list_for_provider("openai");
167        assert_eq!(openai_models.len(), 1);
168    }
169}
razeri_plugin_api/model.rs

razeri_plugin_api/
model.rs
