1use serde::{Deserialize, Serialize};
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct HanzoConfig {
15 pub hanzo_home: PathBuf,
17 pub models_dir: PathBuf,
18 pub embeddings_dir: PathBuf,
19 pub rerankers_dir: PathBuf,
20 pub llms_dir: PathBuf,
21 pub cache_dir: PathBuf,
22 pub data_dir: PathBuf,
23 pub logs_dir: PathBuf,
24 pub config_dir: PathBuf,
25
26 pub max_cache_size_gb: f64,
28 pub max_models_size_gb: f64,
29 pub auto_cleanup: bool,
30
31 pub lancedb_path: PathBuf,
33 pub lancedb_max_size_gb: f64,
34 pub enable_vector_search: bool,
35 pub enable_full_text_search: bool,
36
37 pub web_enabled: bool,
39 pub web_host: String,
40 pub web_port: u16,
41 pub api_host: String,
42 pub api_port: u16,
43 pub ws_enabled: bool, pub ws_port: Option<u16>, pub p2p_port: u16, pub public_url: Option<String>,
47 pub enable_cors: bool,
48 pub allowed_origins: Vec<String>,
49
50 pub engine_binary: PathBuf,
52 pub engine_threads: usize,
53 pub engine_gpu_layers: Option<u32>,
54 pub engine_batch_size: usize,
55
56 pub default_embedding_model: String,
58 pub default_reranker_model: String,
59 pub default_llm_model: String,
60
61 pub api_keys_file: PathBuf,
63}
64
65impl Default for HanzoConfig {
66 fn default() -> Self {
67 let home = dirs::home_dir()
68 .expect("Could not find home directory")
69 .join(".hanzo");
70
71 Self {
72 hanzo_home: home.clone(),
73 models_dir: home.join("models"),
74 embeddings_dir: home.join("models/embeddings"),
75 rerankers_dir: home.join("models/rerankers"),
76 llms_dir: home.join("models/llms"),
77 cache_dir: home.join("cache"),
78 data_dir: home.join("data"),
79 logs_dir: home.join("logs"),
80 config_dir: home.join("config"),
81
82 max_cache_size_gb: 10.0,
84 max_models_size_gb: 100.0,
85 auto_cleanup: true,
86
87 lancedb_path: home.join("data/lancedb"),
89 lancedb_max_size_gb: 50.0,
90 enable_vector_search: true,
91 enable_full_text_search: true,
92
93 web_enabled: true,
95 web_host: "0.0.0.0".to_string(),
96 web_port: 3692, api_host: "0.0.0.0".to_string(),
98 api_port: 3690, ws_enabled: true,
100 ws_port: None, p2p_port: 3691, public_url: None,
103 enable_cors: true,
104 allowed_origins: vec!["*".to_string()],
105
106 engine_binary: home.join("bin/hanzo-engine"),
108 engine_threads: num_cpus::get(),
109 engine_gpu_layers: None,
110 engine_batch_size: 32,
111
112 default_embedding_model: "qwen3-embedding-8b".to_string(),
114 default_reranker_model: "qwen3-reranker-8b".to_string(),
115 default_llm_model: "qwen3-8b-instruct".to_string(),
116
117 api_keys_file: home.join("config/api_keys.encrypted"),
119 }
120 }
121}
122
123impl HanzoConfig {
124 pub fn load() -> Result<Self, Box<dyn std::error::Error>> {
126 let config_path = Self::default().config_dir.join("hanzo.toml");
127
128 if config_path.exists() {
129 let contents = fs::read_to_string(&config_path)?;
130 let config: Self = toml::from_str(&contents)?;
131 Ok(config)
132 } else {
133 let config = Self::default();
134 config.save()?;
135 Ok(config)
136 }
137 }
138
139 pub fn save(&self) -> Result<(), Box<dyn std::error::Error>> {
141 fs::create_dir_all(&self.config_dir)?;
142 let config_path = self.config_dir.join("hanzo.toml");
143 let contents = toml::to_string_pretty(self)?;
144 fs::write(&config_path, contents)?;
145 Ok(())
146 }
147
148 pub fn ensure_directories(&self) -> Result<(), Box<dyn std::error::Error>> {
150 let dirs = vec![
151 &self.hanzo_home,
152 &self.models_dir,
153 &self.embeddings_dir,
154 &self.rerankers_dir,
155 &self.llms_dir,
156 &self.cache_dir,
157 &self.data_dir,
158 &self.logs_dir,
159 &self.config_dir,
160 &self.lancedb_path,
161 ];
162
163 for dir in dirs {
164 fs::create_dir_all(dir)?;
165 }
166
167 Ok(())
168 }
169
170 pub fn get_storage_stats(&self) -> StorageStats {
172 StorageStats {
173 models_size_gb: get_dir_size_gb(&self.models_dir),
174 cache_size_gb: get_dir_size_gb(&self.cache_dir),
175 lancedb_size_gb: get_dir_size_gb(&self.lancedb_path),
176 total_size_gb: get_dir_size_gb(&self.hanzo_home),
177 }
178 }
179
180 pub fn cleanup_cache(&self) -> Result<usize, Box<dyn std::error::Error>> {
182 if !self.auto_cleanup {
183 return Ok(0);
184 }
185
186 let current_size = get_dir_size_gb(&self.cache_dir);
187 if current_size <= self.max_cache_size_gb {
188 return Ok(0);
189 }
190
191 let mut entries: Vec<_> = fs::read_dir(&self.cache_dir)?
193 .filter_map(|e| e.ok())
194 .collect();
195
196 entries.sort_by_key(|e| {
197 e.metadata()
198 .and_then(|m| m.modified())
199 .unwrap_or_else(|_| std::time::SystemTime::UNIX_EPOCH)
200 });
201
202 let mut deleted = 0;
203 let mut current_size = current_size;
204
205 for entry in entries {
206 if current_size <= self.max_cache_size_gb * 0.8 {
207 break;
208 }
209
210 if let Ok(metadata) = entry.metadata() {
211 let size_gb = metadata.len() as f64 / 1_073_741_824.0;
212 fs::remove_file(entry.path())?;
213 current_size -= size_gb;
214 deleted += 1;
215 }
216 }
217
218 Ok(deleted)
219 }
220
221 pub fn get_model_path(&self, model_name: &str) -> PathBuf {
223 if model_name.contains("embed") {
224 self.embeddings_dir.join(model_name)
225 } else if model_name.contains("rerank") {
226 self.rerankers_dir.join(model_name)
227 } else {
228 self.llms_dir.join(model_name)
229 }
230 }
231
232 pub fn is_model_downloaded(&self, model_name: &str) -> bool {
234 let path = self.get_model_path(model_name);
235 path.exists() && path.is_dir()
236 }
237
238 pub fn get_public_url(&self) -> String {
240 self.public_url
241 .clone()
242 .unwrap_or_else(|| format!("http://{}:{}", self.web_host, self.web_port))
243 }
244
245 pub fn get_api_url(&self) -> String {
247 format!("http://{}:{}", self.api_host, self.api_port)
248 }
249
250 pub fn get_ws_url(&self) -> String {
252 let port = self.ws_port.unwrap_or(self.api_port);
253 format!("ws://{}:{}", self.api_host, port)
254 }
255
256 pub fn get_public_ws_url(&self) -> String {
258 if let Some(ref public_url) = self.public_url {
259 public_url
261 .replace("https://", "wss://")
262 .replace("http://", "ws://")
263 } else {
264 self.get_ws_url()
265 }
266 }
267}
268
269#[derive(Debug, Clone, Serialize, Deserialize)]
270pub struct StorageStats {
271 pub models_size_gb: f64,
272 pub cache_size_gb: f64,
273 pub lancedb_size_gb: f64,
274 pub total_size_gb: f64,
275}
276
277fn get_dir_size_gb(path: &Path) -> f64 {
278 if !path.exists() {
279 return 0.0;
280 }
281
282 let size = walkdir::WalkDir::new(path)
283 .into_iter()
284 .filter_map(|e| e.ok())
285 .filter_map(|e| e.metadata().ok())
286 .filter(|m| m.is_file())
287 .map(|m| m.len())
288 .sum::<u64>();
289
290 size as f64 / 1_073_741_824.0
291}
292
293pub struct GlobalConfig {
295 inner: Arc<RwLock<HanzoConfig>>,
296}
297
298impl GlobalConfig {
299 pub fn new() -> Self {
300 Self {
301 inner: Arc::new(RwLock::new(HanzoConfig::default())),
302 }
303 }
304
305 pub async fn load() -> Result<Self, Box<dyn std::error::Error>> {
306 let config = HanzoConfig::load()?;
307 Ok(Self {
308 inner: Arc::new(RwLock::new(config)),
309 })
310 }
311
312 pub async fn get(&self) -> HanzoConfig {
313 self.inner.read().await.clone()
314 }
315
316 pub async fn update<F>(&self, f: F) -> Result<(), Box<dyn std::error::Error>>
317 where
318 F: FnOnce(&mut HanzoConfig),
319 {
320 let mut config = self.inner.write().await;
321 f(&mut *config);
322 config.save()?;
323 Ok(())
324 }
325}
326
327pub async fn init_hanzo_environment() -> Result<GlobalConfig, Box<dyn std::error::Error>> {
329 let config = GlobalConfig::load().await?;
330 let cfg = config.get().await;
331 cfg.ensure_directories()?;
332
333 let hanzo_toml = cfg.config_dir.join("hanzo.toml");
335 if !hanzo_toml.exists() {
336 cfg.save()?;
337 }
338
339 if cfg.enable_vector_search {
341 fs::create_dir_all(&cfg.lancedb_path)?;
342 }
343
344 if cfg.auto_cleanup {
346 let _ = cfg.cleanup_cache();
347 }
348
349 Ok(config)
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355
356 #[test]
357 fn test_default_config() {
358 let config = HanzoConfig::default();
359 assert_eq!(config.default_embedding_model, "qwen3-embedding-8b");
360 assert_eq!(config.default_reranker_model, "qwen3-reranker-8b");
361 assert!(config.hanzo_home.ends_with(".hanzo"));
362 }
363
364 #[tokio::test]
365 async fn test_global_config() {
366 let config = GlobalConfig::new();
367 let cfg = config.get().await;
368 assert_eq!(cfg.web_port, 3692);
369 assert_eq!(cfg.api_port, 3690);
370 }
371}