mermaid_cli/models/
traits.rs1use async_trait::async_trait;
6
7use super::config::ModelConfig;
8use super::error::Result;
9use super::types::{ChatMessage, ModelResponse, StreamCallback};
10
11#[async_trait]
15pub trait Model: Send + Sync {
16 async fn chat(
18 &self,
19 messages: &[ChatMessage],
20 config: &ModelConfig,
21 stream_callback: Option<StreamCallback>,
22 ) -> Result<ModelResponse>;
23
24 fn name(&self) -> &str;
26
27 fn is_local(&self) -> bool;
29
30 async fn health_check(&self) -> Result<()>;
32
33 async fn list_models(&self) -> Result<Vec<String>>;
35
36 async fn has_model(&self, model_name: &str) -> Result<bool> {
38 let models = self.list_models().await?;
39 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 fn capabilities(&self) -> ModelCapabilities {
49 ModelCapabilities::default()
50 }
51}
52
53#[derive(Debug, Clone)]
55pub struct ModelCapabilities {
56 pub max_context_length: usize,
58 pub supports_streaming: bool,
60 pub supports_functions: bool,
62 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}