Skip to main content

hanzo_config/
lib.rs

1use serde::{Deserialize, Serialize};
2use std::fs;
3use std::path::{Path, PathBuf};
4use std::sync::Arc;
5use tokio::sync::RwLock;
6
7/// Unified configuration for all Hanzo components
8/// This ensures consistency across:
9/// - hanzoai (engine CLI)
10/// - hanzod (network node with web exposure)
11/// - app (~/work/hanzo/app)
12/// - All other Hanzo tools
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct HanzoConfig {
15    // Paths
16    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    // Storage settings
27    pub max_cache_size_gb: f64,
28    pub max_models_size_gb: f64,
29    pub auto_cleanup: bool,
30
31    // Vector DB settings (LanceDB for RAG)
32    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    // Web exposure settings (for hanzod)
38    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,     // WebSocket support
44    pub ws_port: Option<u16>, // WebSocket port (if different from api_port)
45    pub p2p_port: u16,        // P2P consensus port
46    pub public_url: Option<String>,
47    pub enable_cors: bool,
48    pub allowed_origins: Vec<String>,
49
50    // Engine settings
51    pub engine_binary: PathBuf,
52    pub engine_threads: usize,
53    pub engine_gpu_layers: Option<u32>,
54    pub engine_batch_size: usize,
55
56    // Model defaults
57    pub default_embedding_model: String,
58    pub default_reranker_model: String,
59    pub default_llm_model: String,
60
61    // API keys (encrypted storage)
62    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            // Storage: Default to 100GB for models, 10GB for cache
83            max_cache_size_gb: 10.0,
84            max_models_size_gb: 100.0,
85            auto_cleanup: true,
86
87            // LanceDB for vector storage (RAG backend)
88            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 exposure (for public access via hanzod)
94            web_enabled: true,
95            web_host: "0.0.0.0".to_string(),
96            web_port: 3692, // Web interface port (3690 + 2)
97            api_host: "0.0.0.0".to_string(),
98            api_port: 3690, // Main hanzod port (API + WebSocket)
99            ws_enabled: true,
100            ws_port: None,  // Use same port as API (3690) for WebSocket
101            p2p_port: 3691, // P2P consensus port (3690 + 1)
102            public_url: None,
103            enable_cors: true,
104            allowed_origins: vec!["*".to_string()],
105
106            // Engine settings
107            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            // Model defaults - prioritize 8B models
113            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            // Security
118            api_keys_file: home.join("config/api_keys.encrypted"),
119        }
120    }
121}
122
123impl HanzoConfig {
124    /// Load config from ~/.hanzo/config/hanzo.toml or create default
125    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    /// Save config to ~/.hanzo/config/hanzo.toml
140    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    /// Ensure all directories exist
149    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    /// Get storage usage statistics
171    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    /// Clean up old cache files if over limit
181    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        // Clean oldest files first
192        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    /// Get path for a specific model
222    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    /// Check if model is downloaded
233    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    /// Get public URL for web exposure
239    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    /// Get API URL
246    pub fn get_api_url(&self) -> String {
247        format!("http://{}:{}", self.api_host, self.api_port)
248    }
249
250    /// Get WebSocket URL
251    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    /// Get public WebSocket URL
257    pub fn get_public_ws_url(&self) -> String {
258        if let Some(ref public_url) = self.public_url {
259            // Convert http/https to ws/wss
260            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
293/// Global config instance for all Hanzo components
294pub 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
327/// Initialize Hanzo environment for all tools
328pub 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    // Create initial config files if they don't exist
334    let hanzo_toml = cfg.config_dir.join("hanzo.toml");
335    if !hanzo_toml.exists() {
336        cfg.save()?;
337    }
338
339    // Initialize LanceDB directory structure
340    if cfg.enable_vector_search {
341        fs::create_dir_all(&cfg.lancedb_path)?;
342    }
343
344    // Clean cache if needed
345    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}