razeri_plugin_api/
model.rs

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

razeri_plugin_api/
model.rs
