microscope_memory/
zen_keystore.rs

1//! Binary Zen key store — `zen_keys.bin`
2//!
3//! Zero-JSON, pure bincode. ORA kulcsai binárisan.
4//! Formátum: ZKEY magic + version + Provider[] + Model[].
5//!
6//! Váltja ki a zen_keys.json-t — gyorsabb, kisebb, biztonságosabb.
7
8use serde::{Deserialize, Serialize};
9use std::fs;
10use std::path::Path;
11use std::str::FromStr;
12
13/// ZKEY magic bytes
14const MAGIC: [u8; 4] = [0x5A, 0x4B, 0x45, 0x59]; // "ZKEY"
15const VERSION: u8 = 1;
16
17/// A teljes Zen kulcstár
18#[derive(Debug, Serialize, Deserialize, Clone)]
19pub struct ZenKeyStore {
20    pub magic: [u8; 4],
21    pub version: u8,
22    pub providers: Vec<Provider>,
23    pub models: Vec<Model>,
24}
25
26/// Egy API provider (openai, groq, sambanova, ...)
27#[derive(Debug, Serialize, Deserialize, Clone)]
28pub struct Provider {
29    /// "openai" | "groq" | "sambanova" | ...
30    pub name: String,
31    /// API base URL
32    pub base_url: String,
33    /// Kulcs rotációs stratégia
34    pub rotation: RotationStrategy,
35    /// API kulcsok (priority sorrendben)
36    pub keys: Vec<KeyEntry>,
37}
38
39/// Kulcs rotációs stratégia
40#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq)]
41pub enum RotationStrategy {
42    RoundRobin,
43    Priority,
44}
45
46impl RotationStrategy {
47    pub fn as_str(&self) -> &'static str {
48        match self {
49            Self::RoundRobin => "round-robin",
50            Self::Priority => "priority",
51        }
52    }
53}
54
55impl FromStr for RotationStrategy {
56    type Err = ();
57
58    fn from_str(s: &str) -> Result<Self, Self::Err> {
59        match s.to_lowercase().as_str() {
60            "priority" => Ok(Self::Priority),
61            _ => Ok(Self::RoundRobin),
62        }
63    }
64}
65
66/// Egy API kulcs
67#[derive(Debug, Serialize, Deserialize, Clone)]
68pub struct KeyEntry {
69    /// Maga a kulcs
70    pub key: String,
71    /// 0 = elsődleges, 1 = másodlagos, ...
72    pub priority: u8,
73    /// Ha tudjuk, mennyi quota maradt
74    pub quota_remaining: Option<f64>,
75    /// Utolsó hibaüzenet
76    #[serde(default)]
77    pub last_error: Option<String>,
78    /// Letiltva (pl. 429 miatt)
79    #[serde(default)]
80    pub disabled: bool,
81}
82
83/// Egy modell definíció
84#[derive(Debug, Serialize, Deserialize, Clone)]
85pub struct Model {
86    /// Modell ID (pl. "deepseek-v4-flash-free")
87    pub id: String,
88    /// Provider neve (None = default / openai)
89    pub provider: Option<String>,
90    /// Endpoint típus: "chat/completions" vagy "messages"
91    pub endpoint: String,
92    /// Ingyenes modell?
93    pub free: bool,
94    /// Priority (0 = legmagasabb)
95    pub priority: u8,
96}
97
98impl ZenKeyStore {
99    /// Betöltés bináris fájlból
100    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, Box<dyn std::error::Error>> {
101        if !path.as_ref().exists() {
102            return Ok(Self {
103                magic: MAGIC,
104                version: VERSION,
105                providers: Vec::new(),
106                models: Vec::new(),
107            });
108        }
109        let bytes = fs::read(path.as_ref())?;
110        let store: ZenKeyStore = bincode::deserialize(&bytes)?;
111        if store.magic != MAGIC {
112            return Err("Invalid zen_keys.bin magic bytes".into());
113        }
114        if store.version != VERSION {
115            return Err(format!("Unsupported zen_keys.bin version: {}", store.version).into());
116        }
117        Ok(store)
118    }
119
120    /// Mentés bincode formátumban
121    pub fn save<P: AsRef<Path>>(&self, path: P) -> Result<(), Box<dyn std::error::Error>> {
122        let bytes = bincode::serialize(self)?;
123        let tmp = path.as_ref().with_extension("bin.tmp");
124        fs::write(&tmp, &bytes)?;
125        fs::rename(&tmp, path.as_ref())?;
126        Ok(())
127    }
128
129    /// Kulcs lekérése provider + priority szerint (legkisebb priority = első)
130    pub fn get_key(&self, provider: &str) -> Option<&KeyEntry> {
131        self.providers
132            .iter()
133            .find(|p| p.name == provider)
134            .and_then(|p| p.keys.iter().find(|k| !k.disabled))
135    }
136
137    /// Összes aktív kulcs egy providerhez
138    pub fn get_all_keys(&self, provider: &str) -> Vec<&KeyEntry> {
139        self.providers
140            .iter()
141            .find(|p| p.name == provider)
142            .map(|p| {
143                let mut v: Vec<_> = p.keys.iter().filter(|k| !k.disabled).collect();
144                v.sort_by_key(|k| k.priority);
145                v
146            })
147            .unwrap_or_default()
148    }
149
150    /// Modell lekérése ID alapján
151    pub fn get_model(&self, id: &str) -> Option<&Model> {
152        self.models.iter().find(|m| m.id == id)
153    }
154
155    /// Alapértelmezett modell (legmagasabb priority, free)
156    pub fn default_model(&self) -> Option<&Model> {
157        self.models.iter().min_by_key(|m| m.priority)
158    }
159
160    /// Ingyenes modellek listája
161    pub fn free_models(&self) -> Vec<&Model> {
162        let mut v: Vec<_> = self.models.iter().filter(|m| m.free).collect();
163        v.sort_by_key(|m| m.priority);
164        v
165    }
166
167    /// Hiba bejegyzése — letiltja a kulcsot 429/401/403 esetén
168    pub fn record_error(&mut self, provider: &str, key_idx: usize, error: &str) {
169        if let Some(p) = self.providers.iter_mut().find(|p| p.name == provider) {
170            if let Some(entry) = p.keys.get_mut(key_idx) {
171                entry.last_error = Some(error.to_string());
172                if error.contains("429") || error.contains("401") || error.contains("403") {
173                    entry.disabled = true;
174                }
175            }
176        }
177    }
178
179    /// Sikeres hívás után
180    pub fn record_success(&mut self, provider: &str, key_idx: usize) {
181        if let Some(p) = self.providers.iter_mut().find(|p| p.name == provider) {
182            if let Some(entry) = p.keys.get_mut(key_idx) {
183                entry.last_error = None;
184                entry.disabled = false;
185            }
186        }
187    }
188
189    /// JSON-ból importálás (zen_keys.json formátum)
190    pub fn import_json(json_str: &str) -> Result<Self, Box<dyn std::error::Error>> {
191        let json: serde_json::Value = serde_json::from_str(json_str)?;
192
193        let mut providers = Vec::new();
194
195        // OpenAI/Zen provider
196        if let Some(base_url) = json["base_url"].as_str() {
197            let rotation = json["rotation"]
198                .as_str()
199                .unwrap_or("round-robin")
200                .parse()
201                .unwrap_or(RotationStrategy::RoundRobin);
202            let keys: Vec<KeyEntry> = json["keys"]
203                .as_array()
204                .map(|arr| {
205                    arr.iter()
206                        .enumerate()
207                        .map(|(i, k)| KeyEntry {
208                            key: k.as_str().unwrap_or("").to_string(),
209                            priority: i as u8,
210                            quota_remaining: None,
211                            last_error: None,
212                            disabled: false,
213                        })
214                        .collect()
215                })
216                .unwrap_or_default();
217            providers.push(Provider {
218                name: "openai".to_string(),
219                base_url: base_url.to_string(),
220                rotation,
221                keys,
222            });
223        }
224
225        // Groq
226        if let Some(groq) = json["groq"].as_object() {
227            let base_url = groq["base_url"]
228                .as_str()
229                .unwrap_or("https://api.groq.com/openai/v1")
230                .to_string();
231            let rotation = groq["rotation"]
232                .as_str()
233                .unwrap_or("round-robin")
234                .parse()
235                .unwrap_or(RotationStrategy::RoundRobin);
236            let keys: Vec<KeyEntry> = groq["keys"]
237                .as_array()
238                .map(|arr| {
239                    arr.iter()
240                        .enumerate()
241                        .map(|(i, k)| KeyEntry {
242                            key: k.as_str().unwrap_or("").to_string(),
243                            priority: i as u8,
244                            quota_remaining: None,
245                            last_error: None,
246                            disabled: false,
247                        })
248                        .collect()
249                })
250                .unwrap_or_default();
251            providers.push(Provider {
252                name: "groq".to_string(),
253                base_url,
254                rotation,
255                keys,
256            });
257        }
258
259        // SambaNova
260        if let Some(sn) = json["sambanova"].as_object() {
261            let base_url = sn["base_url"]
262                .as_str()
263                .unwrap_or("https://api.sambanova.ai/v1")
264                .to_string();
265            let rotation = sn["rotation"]
266                .as_str()
267                .unwrap_or("round-robin")
268                .parse()
269                .unwrap_or(RotationStrategy::RoundRobin);
270            let keys: Vec<KeyEntry> = sn["keys"]
271                .as_array()
272                .map(|arr| {
273                    arr.iter()
274                        .enumerate()
275                        .map(|(i, k)| KeyEntry {
276                            key: k.as_str().unwrap_or("").to_string(),
277                            priority: i as u8,
278                            quota_remaining: None,
279                            last_error: None,
280                            disabled: false,
281                        })
282                        .collect()
283                })
284                .unwrap_or_default();
285            providers.push(Provider {
286                name: "sambanova".to_string(),
287                base_url,
288                rotation,
289                keys,
290            });
291        }
292
293        // Modellek
294        let models: Vec<Model> = json["models"]
295            .as_array()
296            .map(|arr| {
297                arr.iter()
298                    .map(|m| Model {
299                        id: m["id"].as_str().unwrap_or("unknown").to_string(),
300                        provider: m["provider"].as_str().map(|s| s.to_string()),
301                        endpoint: m["endpoint"]
302                            .as_str()
303                            .unwrap_or("chat/completions")
304                            .to_string(),
305                        free: m["free"].as_bool().unwrap_or(false),
306                        priority: m["priority"].as_u64().unwrap_or(99) as u8,
307                    })
308                    .collect()
309            })
310            .unwrap_or_default();
311
312        Ok(Self {
313            magic: MAGIC,
314            version: VERSION,
315            providers,
316            models,
317        })
318    }
319
320    /// Szöveges statisztika
321    pub fn stats(&self) -> String {
322        let mut s = String::new();
323        s.push_str(&format!("Providers: {}\n", self.providers.len()));
324        for p in &self.providers {
325            let active = p.keys.iter().filter(|k| !k.disabled).count();
326            let total = p.keys.len();
327            s.push_str(&format!(
328                "  {}: {}/{} keys active, rotation={}\n",
329                p.name,
330                active,
331                total,
332                p.rotation.as_str()
333            ));
334        }
335        s.push_str(&format!("Models: {}\n", self.models.len()));
336        for m in &self.models {
337            let prov = m.provider.as_deref().unwrap_or("openai");
338            s.push_str(&format!(
339                "  #{} {} [{}] {} {}\n",
340                m.priority,
341                m.id,
342                prov,
343                m.endpoint,
344                if m.free { "FREE" } else { "PAID" }
345            ));
346        }
347        s
348    }
349}
350
351/// Alapértelmezett útvonal a zen_keys.bin-hez
352pub fn default_zen_keys_path(output_dir: &str) -> String {
353    format!("{}/zen_keys.bin", output_dir)
354}
microscope_memory/zen_keystore.rs

microscope_memory/
zen_keystore.rs
