Skip to main content

mermaid_cli/models/
traits.rs

1/// Core Model trait - the single interface for model interactions
2///
3/// Adapters implement this trait directly. No intermediate layers.
4
5use async_trait::async_trait;
6
7use super::config::ModelConfig;
8use super::error::Result;
9use super::types::{ChatMessage, ModelResponse, StreamCallback};
10
11/// Core trait that all model adapters implement
12///
13/// This is the only abstraction layer between user code and model providers.
14#[async_trait]
15pub trait Model: Send + Sync {
16    /// Send a chat conversation to the model and get a response
17    async fn chat(
18        &self,
19        messages: &[ChatMessage],
20        config: &ModelConfig,
21        stream_callback: Option<StreamCallback>,
22    ) -> Result<ModelResponse>;
23
24    /// Get the model identifier (e.g., "ollama/tinyllama")
25    fn name(&self) -> &str;
26
27    /// Check if this is a local model (no external API calls)
28    fn is_local(&self) -> bool;
29
30    /// Check if the model backend is available and healthy
31    async fn health_check(&self) -> Result<()>;
32
33    /// List available models from this backend
34    async fn list_models(&self) -> Result<Vec<String>>;
35
36    /// Check if a specific model is available
37    async fn has_model(&self, model_name: &str) -> Result<bool> {
38        let models = self.list_models().await?;
39        // Check for exact match or prefix match (model:tag format)
40        Ok(models.iter().any(|m| {
41            m == model_name
42                || m.starts_with(&format!("{}:", model_name))
43                || model_name.starts_with(&format!("{}:", m))
44        }))
45    }
46
47    /// Get model capabilities
48    fn capabilities(&self) -> ModelCapabilities {
49        ModelCapabilities::default()
50    }
51}
52
53/// Model capabilities (what the model/backend supports)
54#[derive(Debug, Clone)]
55pub struct ModelCapabilities {
56    /// Maximum context length supported
57    pub max_context_length: usize,
58    /// Supports streaming responses
59    pub supports_streaming: bool,
60    /// Supports function/tool calling
61    pub supports_functions: bool,
62    /// Supports vision/images
63    pub supports_vision: bool,
64}
65
66impl Default for ModelCapabilities {
67    fn default() -> Self {
68        Self {
69            max_context_length: 4096,
70            supports_streaming: true,
71            supports_functions: false,
72            supports_vision: false,
73        }
74    }
75}