mermaid_cli/models/
config.rs1use crate::constants::{DEFAULT_MAX_TOKENS, DEFAULT_TEMPERATURE};
7use crate::prompts;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ModelConfig {
14 pub model: String,
17
18 #[serde(default = "default_temperature")]
20 pub temperature: f32,
21
22 #[serde(default = "default_max_tokens")]
24 pub max_tokens: usize,
25
26 pub system_prompt: Option<String>,
28
29 #[serde(default = "default_thinking_enabled")]
32 pub thinking_enabled: Option<bool>,
33
34 #[serde(skip)]
37 pub is_subagent: bool,
38
39 #[serde(default)]
42 pub backend_options: HashMap<String, HashMap<String, String>>,
43}
44
45impl Default for ModelConfig {
46 fn default() -> Self {
47 Self {
48 model: "ollama/tinyllama".to_string(),
49 temperature: default_temperature(),
50 max_tokens: default_max_tokens(),
51 system_prompt: Some(prompts::get_system_prompt()),
52 thinking_enabled: Some(true),
53 is_subagent: false,
54 backend_options: HashMap::new(),
55 }
56 }
57}
58
59impl ModelConfig {
60 pub fn get_backend_option(&self, backend: &str, key: &str) -> Option<&String> {
62 self.backend_options.get(backend)?.get(key)
63 }
64
65 pub fn get_backend_option_i32(&self, backend: &str, key: &str) -> Option<i32> {
67 self.get_backend_option(backend, key)?.parse::<i32>().ok()
68 }
69
70 pub fn get_backend_option_bool(&self, backend: &str, key: &str) -> Option<bool> {
72 self.get_backend_option(backend, key)?.parse::<bool>().ok()
73 }
74
75 pub fn set_backend_option(&mut self, backend: String, key: String, value: String) {
77 self.backend_options
78 .entry(backend)
79 .or_default()
80 .insert(key, value);
81 }
82
83 pub fn ollama_options(&self) -> OllamaOptions {
85 OllamaOptions {
86 num_gpu: self.get_backend_option_i32("ollama", "num_gpu"),
87 num_thread: self.get_backend_option_i32("ollama", "num_thread"),
88 num_ctx: self.get_backend_option_i32("ollama", "num_ctx"),
89 numa: self.get_backend_option_bool("ollama", "numa"),
90 }
91 }
92}
93
94#[derive(Debug, Clone, Default)]
96pub struct OllamaOptions {
97 pub num_gpu: Option<i32>,
98 pub num_thread: Option<i32>,
99 pub num_ctx: Option<i32>,
100 pub numa: Option<bool>,
101}
102
103#[derive(Debug, Clone, Serialize, Deserialize)]
105pub struct BackendConfig {
106 #[serde(default = "default_ollama_url")]
108 pub ollama_url: String,
109
110 #[serde(default = "default_timeout")]
112 pub timeout_secs: u64,
113
114 #[serde(default = "default_max_idle")]
116 pub max_idle_per_host: usize,
117}
118
119impl Default for BackendConfig {
120 fn default() -> Self {
121 Self {
122 ollama_url: default_ollama_url(),
123 timeout_secs: default_timeout(),
124 max_idle_per_host: default_max_idle(),
125 }
126 }
127}
128
129fn default_temperature() -> f32 {
131 DEFAULT_TEMPERATURE
132}
133
134fn default_max_tokens() -> usize {
135 DEFAULT_MAX_TOKENS
136}
137
138fn default_ollama_url() -> String {
139 std::env::var("OLLAMA_HOST").unwrap_or_else(|_| "http://localhost:11434".to_string())
140}
141
142fn default_timeout() -> u64 {
143 10
144}
145
146fn default_max_idle() -> usize {
147 10
148}
149
150fn default_thinking_enabled() -> Option<bool> {
151 Some(true)
152}