ck_models/
lib.rs

1use anyhow::Result;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::path::Path;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ModelConfig {
8    pub name: String,
9    pub provider: String,
10    pub dimensions: usize,
11    pub max_tokens: usize,
12    pub description: String,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct ModelRegistry {
17    pub models: HashMap<String, ModelConfig>,
18    pub default_model: String,
19}
20
21impl Default for ModelRegistry {
22    fn default() -> Self {
23        let mut models = HashMap::new();
24        
25        models.insert(
26            "bge-small".to_string(),
27            ModelConfig {
28                name: "BAAI/bge-small-en-v1.5".to_string(),
29                provider: "fastembed".to_string(),
30                dimensions: 384,
31                max_tokens: 512,
32                description: "Small, fast English embedding model".to_string(),
33            },
34        );
35        
36        models.insert(
37            "minilm".to_string(),
38            ModelConfig {
39                name: "sentence-transformers/all-MiniLM-L6-v2".to_string(),
40                provider: "fastembed".to_string(),
41                dimensions: 384,
42                max_tokens: 256,
43                description: "Lightweight English embedding model".to_string(),
44            },
45        );
46        
47        Self {
48            models,
49            default_model: "bge-small".to_string(),
50        }
51    }
52}
53
54impl ModelRegistry {
55    pub fn load(path: &Path) -> Result<Self> {
56        if path.exists() {
57            let data = std::fs::read_to_string(path)?;
58            Ok(serde_json::from_str(&data)?)
59        } else {
60            Ok(Self::default())
61        }
62    }
63    
64    pub fn save(&self, path: &Path) -> Result<()> {
65        let data = serde_json::to_string_pretty(self)?;
66        std::fs::write(path, data)?;
67        Ok(())
68    }
69    
70    pub fn get_model(&self, name: &str) -> Option<&ModelConfig> {
71        self.models.get(name)
72    }
73    
74    pub fn get_default_model(&self) -> Option<&ModelConfig> {
75        self.models.get(&self.default_model)
76    }
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct ProjectConfig {
81    pub model: String,
82    pub chunk_size: usize,
83    pub chunk_overlap: usize,
84    pub index_backend: String,
85}
86
87impl Default for ProjectConfig {
88    fn default() -> Self {
89        Self {
90            model: "bge-small".to_string(),
91            chunk_size: 512,
92            chunk_overlap: 128,
93            index_backend: "hnsw".to_string(),
94        }
95    }
96}
97
98impl ProjectConfig {
99    pub fn load(path: &Path) -> Result<Self> {
100        if path.exists() {
101            let data = std::fs::read_to_string(path)?;
102            Ok(serde_json::from_str(&data)?)
103        } else {
104            Ok(Self::default())
105        }
106    }
107    
108    pub fn save(&self, path: &Path) -> Result<()> {
109        let data = serde_json::to_string_pretty(self)?;
110        std::fs::write(path, data)?;
111        Ok(())
112    }
113}
ck_models/lib.rs

ck_models/
lib.rs
