edgequake_llm/discovery/
types.rs

1//! Core types for the model discovery system.
2//!
3//! Every type in this module follows the Anti-Heuristic Manifesto:
4//! capabilities come from API responses, verified documentation, or
5//! explicit "unknown" — never from name-pattern inference.
6
7use chrono::{DateTime, Utc};
8use serde::{Deserialize, Serialize};
9
10use crate::model_config::{ModelCapabilities, ModelType};
11
12/// A discovered model with normalized capabilities.
13///
14/// Single source of truth for model metadata regardless of whether
15/// it came from a dynamic API call or a static registry.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct DiscoveredModel {
18    /// Provider-specific model ID (e.g., "gpt-4.1", "claude-opus-4-8").
19    pub id: String,
20    /// Human-readable display name.
21    pub name: String,
22    /// Provider that owns this model (string, not enum — avoids coupling to ProviderType).
23    pub provider: String,
24    /// Maximum input context window in tokens.
25    pub context_length: usize,
26    /// Maximum output tokens the model can generate.
27    pub max_output_tokens: usize,
28    /// Normalized capabilities.
29    pub capabilities: ModelCapabilities,
30    /// How this model was discovered.
31    pub source: DiscoverySource,
32    /// When this information was last verified.
33    pub discovered_at: DateTime<Utc>,
34    /// Whether the model is currently available/loaded.
35    pub available: bool,
36    /// Cost per million input tokens (USD).
37    pub cost_per_m_input: Option<f64>,
38    /// Cost per million output tokens (USD).
39    pub cost_per_m_output: Option<f64>,
40    /// Model type (LLM, Embedding, Multimodal).
41    pub model_type: ModelType,
42    /// Tags for filtering (e.g., "reasoning", "coding", "fast").
43    pub tags: Vec<String>,
44    /// Whether the model is deprecated.
45    pub deprecated: bool,
46}
47
48impl Default for DiscoveredModel {
49    fn default() -> Self {
50        Self {
51            id: String::new(),
52            name: String::new(),
53            provider: String::new(),
54            context_length: 0,
55            max_output_tokens: 0,
56            capabilities: ModelCapabilities::default(),
57            source: DiscoverySource::Unknown,
58            discovered_at: Utc::now(),
59            available: true,
60            cost_per_m_input: None,
61            cost_per_m_output: None,
62            model_type: ModelType::Llm,
63            tags: Vec::new(),
64            deprecated: false,
65        }
66    }
67}
68
69/// How a model was discovered.
70#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
71pub enum DiscoverySource {
72    /// Discovered via live API call (highest freshness).
73    DynamicApi,
74    /// From built-in static registry (manually verified, may be stale).
75    StaticRegistry,
76    /// Dynamic API enriched with static metadata.
77    Hybrid,
78    /// User-provided configuration (models.toml override).
79    UserConfig,
80    /// Model not found in any source — all capabilities are unknown.
81    /// Caller MUST handle this explicitly (never guess).
82    #[default]
83    Unknown,
84}
85
86/// How a provider discovers its models.
87#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
88pub enum DiscoveryStrategy {
89    /// Models discovered via live API call (Ollama, LM Studio, OpenRouter, Anthropic, Gemini, Mistral).
90    Dynamic,
91    /// Models from a built-in static registry (xAI, HuggingFace).
92    Static,
93    /// API call enriched with static metadata (OpenAI, NVIDIA, Bedrock).
94    Hybrid,
95}
96
97/// Filter for discovering models by capability.
98///
99/// All fields use AND logic — a model must satisfy every non-None constraint.
100#[derive(Debug, Clone, Default, Serialize, Deserialize)]
101pub struct CapabilityFilter {
102    /// Minimum input context window (tokens).
103    pub min_context_length: Option<usize>,
104    /// Maximum input context window (tokens).
105    pub max_context_length: Option<usize>,
106    /// Minimum max output tokens the model can generate.
107    pub min_output_tokens: Option<usize>,
108    /// Maximum max output tokens the model can generate.
109    pub max_output_tokens: Option<usize>,
110    pub requires_vision: Option<bool>,
111    pub requires_tools: Option<bool>,
112    pub requires_thinking: Option<bool>,
113    pub requires_streaming: Option<bool>,
114    pub requires_json_mode: Option<bool>,
115    pub model_type: Option<ModelType>,
116    pub provider: Option<String>,
117    pub tags: Option<Vec<String>>,
118    pub max_cost_per_m_input: Option<f64>,
119    pub exclude_deprecated: Option<bool>,
120}
121
122/// Typed model capability for ergonomic filter construction.
123#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
124pub enum ModelCapability {
125    Vision,
126    Tools,
127    Thinking,
128    Streaming,
129    JsonMode,
130}
131
132impl CapabilityFilter {
133    /// Require a specific capability (sets the corresponding `requires_*` field to `true`).
134    pub fn requiring(mut self, capability: ModelCapability) -> Self {
135        match capability {
136            ModelCapability::Vision => self.requires_vision = Some(true),
137            ModelCapability::Tools => self.requires_tools = Some(true),
138            ModelCapability::Thinking => self.requires_thinking = Some(true),
139            ModelCapability::Streaming => self.requires_streaming = Some(true),
140            ModelCapability::JsonMode => self.requires_json_mode = Some(true),
141        }
142        self
143    }
144
145    /// Require multiple capabilities (AND logic).
146    pub fn requiring_all(
147        mut self,
148        capabilities: impl IntoIterator<Item = ModelCapability>,
149    ) -> Self {
150        for capability in capabilities {
151            self = self.requiring(capability);
152        }
153        self
154    }
155
156    pub fn with_min_context_length(mut self, tokens: usize) -> Self {
157        self.min_context_length = Some(tokens);
158        self
159    }
160
161    pub fn with_max_context_length(mut self, tokens: usize) -> Self {
162        self.max_context_length = Some(tokens);
163        self
164    }
165
166    pub fn with_min_output_tokens(mut self, tokens: usize) -> Self {
167        self.min_output_tokens = Some(tokens);
168        self
169    }
170
171    pub fn with_max_output_tokens(mut self, tokens: usize) -> Self {
172        self.max_output_tokens = Some(tokens);
173        self
174    }
175
176    pub fn with_provider(mut self, provider: impl Into<String>) -> Self {
177        self.provider = Some(provider.into());
178        self
179    }
180
181    pub fn with_max_cost_per_m_input(mut self, usd: f64) -> Self {
182        self.max_cost_per_m_input = Some(usd);
183        self
184    }
185
186    pub fn excluding_deprecated(mut self) -> Self {
187        self.exclude_deprecated = Some(true);
188        self
189    }
190
191    pub fn with_model_type(mut self, model_type: ModelType) -> Self {
192        self.model_type = Some(model_type);
193        self
194    }
195
196    pub fn with_tags(mut self, tags: impl IntoIterator<Item = impl Into<String>>) -> Self {
197        self.tags = Some(tags.into_iter().map(Into::into).collect());
198        self
199    }
200
201    /// Test whether a model matches all filter constraints.
202    pub fn matches(&self, model: &DiscoveredModel) -> bool {
203        if let Some(min) = self.min_context_length {
204            if model.context_length < min {
205                return false;
206            }
207        }
208        if let Some(max) = self.max_context_length {
209            if model.context_length > max {
210                return false;
211            }
212        }
213        if let Some(min) = self.min_output_tokens {
214            if model.max_output_tokens < min {
215                return false;
216            }
217        }
218        if let Some(max) = self.max_output_tokens {
219            if model.max_output_tokens > max {
220                return false;
221            }
222        }
223        if let Some(true) = self.requires_vision {
224            if !model.capabilities.supports_vision {
225                return false;
226            }
227        }
228        if let Some(true) = self.requires_tools {
229            if !model.capabilities.supports_function_calling {
230                return false;
231            }
232        }
233        if let Some(true) = self.requires_thinking {
234            if !model.capabilities.supports_thinking {
235                return false;
236            }
237        }
238        if let Some(true) = self.requires_streaming {
239            if !model.capabilities.supports_streaming {
240                return false;
241            }
242        }
243        if let Some(true) = self.requires_json_mode {
244            if !model.capabilities.supports_json_mode {
245                return false;
246            }
247        }
248        if let Some(ref mt) = self.model_type {
249            if &model.model_type != mt {
250                return false;
251            }
252        }
253        if let Some(ref p) = self.provider {
254            if &model.provider != p {
255                return false;
256            }
257        }
258        if let Some(ref tags) = self.tags {
259            if !tags.iter().any(|t| model.tags.contains(t)) {
260                return false;
261            }
262        }
263        if let Some(max) = self.max_cost_per_m_input {
264            match model.cost_per_m_input {
265                Some(cost) if cost > max => return false,
266                None => return false,
267                _ => {}
268            }
269        }
270        if let Some(true) = self.exclude_deprecated {
271            if model.deprecated {
272                return false;
273            }
274        }
275        true
276    }
277}
278
279/// Errors specific to model discovery.
280#[derive(Debug, thiserror::Error)]
281pub enum DiscoveryError {
282    #[error("Provider '{provider}' unreachable: {source}")]
283    ProviderUnreachable {
284        provider: String,
285        source: Box<dyn std::error::Error + Send + Sync>,
286    },
287
288    #[error("Failed to parse discovery response from '{provider}': {details}")]
289    ParseError { provider: String, details: String },
290
291    #[error("Authentication required for '{provider}' model discovery")]
292    AuthRequired { provider: String },
293
294    #[error("Discovery timed out for '{provider}' after {timeout_ms}ms")]
295    Timeout { provider: String, timeout_ms: u64 },
296
297    #[error("Provider '{provider}' not registered")]
298    ProviderNotRegistered { provider: String },
299
300    #[error("{0}")]
301    Other(String),
302}
303
304impl From<DiscoveryError> for crate::error::LlmError {
305    fn from(e: DiscoveryError) -> Self {
306        crate::error::LlmError::Unknown(e.to_string())
307    }
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313
314    fn make_model(
315        id: &str,
316        provider: &str,
317        ctx: usize,
318        vision: bool,
319        tools: bool,
320        thinking: bool,
321    ) -> DiscoveredModel {
322        DiscoveredModel {
323            id: id.into(),
324            name: id.into(),
325            provider: provider.into(),
326            context_length: ctx,
327            max_output_tokens: 4096,
328            capabilities: ModelCapabilities {
329                supports_vision: vision,
330                supports_function_calling: tools,
331                supports_thinking: thinking,
332                supports_streaming: true,
333                supports_json_mode: true,
334                ..Default::default()
335            },
336            cost_per_m_input: Some(2.0),
337            ..Default::default()
338        }
339    }
340
341    #[test]
342    fn test_filter_matches_all() {
343        let model = make_model("gpt-4.1", "openai", 1_000_000, true, true, false);
344        let filter = CapabilityFilter::default();
345        assert!(filter.matches(&model));
346    }
347
348    #[test]
349    fn test_filter_max_context_length() {
350        let model = make_model("big", "openai", 1_000_000, true, true, false);
351        let filter = CapabilityFilter {
352            max_context_length: Some(500_000),
353            ..Default::default()
354        };
355        assert!(!filter.matches(&model));
356    }
357
358    #[test]
359    fn test_filter_max_output_tokens() {
360        let mut model = make_model("small-out", "openai", 128_000, true, true, false);
361        model.max_output_tokens = 4096;
362        let filter = CapabilityFilter {
363            max_output_tokens: Some(2048),
364            ..Default::default()
365        };
366        assert!(!filter.matches(&model));
367    }
368
369    #[test]
370    fn test_filter_min_context() {
371        let model = make_model("small", "openai", 8_000, false, false, false);
372        let filter = CapabilityFilter {
373            min_context_length: Some(100_000),
374            ..Default::default()
375        };
376        assert!(!filter.matches(&model));
377    }
378
379    #[test]
380    fn test_filter_requires_vision() {
381        let model = make_model("no-vision", "openai", 128_000, false, true, false);
382        let filter = CapabilityFilter {
383            requires_vision: Some(true),
384            ..Default::default()
385        };
386        assert!(!filter.matches(&model));
387    }
388
389    #[test]
390    fn test_filter_requires_tools() {
391        let model = make_model("no-tools", "openai", 128_000, true, false, false);
392        let filter = CapabilityFilter {
393            requires_tools: Some(true),
394            ..Default::default()
395        };
396        assert!(!filter.matches(&model));
397    }
398
399    #[test]
400    fn test_filter_requires_thinking() {
401        let model = make_model("thinker", "anthropic", 200_000, true, true, true);
402        let filter = CapabilityFilter {
403            requires_thinking: Some(true),
404            ..Default::default()
405        };
406        assert!(filter.matches(&model));
407    }
408
409    #[test]
410    fn test_filter_provider() {
411        let model = make_model("claude", "anthropic", 200_000, true, true, true);
412        let filter = CapabilityFilter {
413            provider: Some("openai".into()),
414            ..Default::default()
415        };
416        assert!(!filter.matches(&model));
417    }
418
419    #[test]
420    fn test_filter_max_cost() {
421        let model = make_model("expensive", "openai", 128_000, true, true, false);
422        let filter = CapabilityFilter {
423            max_cost_per_m_input: Some(1.0),
424            ..Default::default()
425        };
426        assert!(!filter.matches(&model));
427    }
428
429    #[test]
430    fn test_filter_exclude_deprecated() {
431        let mut model = make_model("old", "openai", 128_000, true, true, false);
432        model.deprecated = true;
433        let filter = CapabilityFilter {
434            exclude_deprecated: Some(true),
435            ..Default::default()
436        };
437        assert!(!filter.matches(&model));
438    }
439
440    #[test]
441    fn test_filter_builder_requires_capabilities() {
442        let filter = CapabilityFilter::default()
443            .requiring_all([ModelCapability::Vision, ModelCapability::Tools]);
444        let model = make_model("vision-tools", "openai", 128_000, true, true, false);
445        assert!(filter.matches(&model));
446        let no_tools = make_model("vision-only", "openai", 128_000, true, false, false);
447        assert!(!filter.matches(&no_tools));
448    }
449
450    #[test]
451    fn test_filter_combined() {
452        let model = make_model("claude-opus-4-8", "anthropic", 1_000_000, true, true, true);
453        let filter = CapabilityFilter {
454            min_context_length: Some(500_000),
455            requires_vision: Some(true),
456            requires_tools: Some(true),
457            requires_thinking: Some(true),
458            provider: Some("anthropic".into()),
459            ..Default::default()
460        };
461        assert!(filter.matches(&model));
462    }
463
464    #[test]
465    fn test_discovery_source_default_is_unknown() {
466        assert_eq!(DiscoverySource::default(), DiscoverySource::Unknown);
467    }
468
469    #[test]
470    fn test_unknown_model_has_zero_capabilities() {
471        let model = DiscoveredModel::default();
472        assert_eq!(model.context_length, 0);
473        assert_eq!(model.max_output_tokens, 0);
474        assert!(!model.capabilities.supports_vision);
475        assert!(!model.capabilities.supports_function_calling);
476        assert!(!model.capabilities.supports_thinking);
477        assert_eq!(model.source, DiscoverySource::Unknown);
478    }
479}
edgequake_llm/discovery/types.rs

edgequake_llm/discovery/
types.rs
