manx_cli/rag/
model_metadata.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::path::PathBuf;
6
7/// Metadata about an installed embedding model
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct ModelMetadata {
10    pub model_name: String,
11    pub provider_type: String,
12    pub dimension: usize,
13    pub size_mb: f64,
14    pub model_path: Option<PathBuf>,
15    pub api_endpoint: Option<String>,
16    pub installed_date: DateTime<Utc>,
17    pub last_used: Option<DateTime<Utc>>,
18    pub checksum: Option<String>,
19    pub description: Option<String>,
20    pub max_input_length: Option<usize>,
21}
22
23/// Manager for model metadata storage and retrieval
24pub struct ModelMetadataManager {
25    metadata_file: PathBuf,
26    models: HashMap<String, ModelMetadata>,
27}
28
29impl ModelMetadataManager {
30    /// Create a new metadata manager
31    pub fn new() -> Result<Self> {
32        let cache_dir = dirs::cache_dir()
33            .unwrap_or_else(|| PathBuf::from("."))
34            .join("manx")
35            .join("models");
36
37        std::fs::create_dir_all(&cache_dir)?;
38        let metadata_file = cache_dir.join("metadata.json");
39
40        let models = if metadata_file.exists() {
41            let content = std::fs::read_to_string(&metadata_file)?;
42            serde_json::from_str(&content).unwrap_or_default()
43        } else {
44            HashMap::new()
45        };
46
47        Ok(Self {
48            metadata_file,
49            models,
50        })
51    }
52
53    /// Get metadata for a specific model
54    pub fn get_model(&self, model_name: &str) -> Option<&ModelMetadata> {
55        self.models.get(model_name)
56    }
57
58    /// Add or update model metadata
59    pub fn add_model(&mut self, metadata: ModelMetadata) -> Result<()> {
60        self.models.insert(metadata.model_name.clone(), metadata);
61        self.save()
62    }
63
64    /// Remove model metadata
65    pub fn remove_model(&mut self, model_name: &str) -> Result<()> {
66        self.models.remove(model_name);
67        self.save()
68    }
69
70    /// List all installed models
71    pub fn list_models(&self) -> Vec<&ModelMetadata> {
72        self.models.values().collect()
73    }
74
75    /// Update last used timestamp for a model
76    #[allow(dead_code)] // Used in ONNX provider when available
77    pub fn mark_used(&mut self, model_name: &str) -> Result<()> {
78        if let Some(model) = self.models.get_mut(model_name) {
79            model.last_used = Some(Utc::now());
80            self.save()?;
81        }
82        Ok(())
83    }
84
85    /// Get the cache directory for models
86    pub fn get_models_dir() -> PathBuf {
87        dirs::cache_dir()
88            .unwrap_or_else(|| PathBuf::from("."))
89            .join("manx")
90            .join("models")
91    }
92
93    /// Save metadata to disk
94    fn save(&self) -> Result<()> {
95        let content = serde_json::to_string_pretty(&self.models)?;
96        std::fs::write(&self.metadata_file, content)?;
97        Ok(())
98    }
99}