ricecoder_local_models/
models.rs

1//! Data models for local model management
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6/// Information about a local model
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
8pub struct LocalModel {
9    /// Model name/ID (e.g., "mistral:latest")
10    pub name: String,
11
12    /// Model size in bytes
13    pub size: u64,
14
15    /// Model digest/hash
16    pub digest: String,
17
18    /// When the model was last modified
19    pub modified_at: DateTime<Utc>,
20
21    /// Model metadata
22    pub metadata: ModelMetadata,
23}
24
25/// Metadata about a model
26#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
27pub struct ModelMetadata {
28    /// Model format (e.g., "gguf")
29    pub format: String,
30
31    /// Model family (e.g., "llama", "mistral")
32    pub family: String,
33
34    /// Parameter size (e.g., "7B", "13B")
35    pub parameter_size: String,
36
37    /// Quantization level (e.g., "Q4_0", "Q5_K_M")
38    pub quantization_level: String,
39}
40
41/// Progress information for model pull operations
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct PullProgress {
44    /// Model name being pulled
45    pub model: String,
46
47    /// Current status message
48    pub status: String,
49
50    /// Model digest
51    pub digest: String,
52
53    /// Total bytes to download
54    pub total: u64,
55
56    /// Bytes downloaded so far
57    pub completed: u64,
58}
59
60impl PullProgress {
61    /// Get the progress percentage (0-100)
62    pub fn percentage(&self) -> f64 {
63        if self.total == 0 {
64            0.0
65        } else {
66            (self.completed as f64 / self.total as f64) * 100.0
67        }
68    }
69
70    /// Check if pull is complete
71    pub fn is_complete(&self) -> bool {
72        self.completed >= self.total && self.total > 0
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn test_pull_progress_percentage() {
82        let progress = PullProgress {
83            model: "mistral".to_string(),
84            status: "downloading".to_string(),
85            digest: "abc123".to_string(),
86            total: 1000,
87            completed: 500,
88        };
89
90        assert_eq!(progress.percentage(), 50.0);
91    }
92
93    #[test]
94    fn test_pull_progress_complete() {
95        let progress = PullProgress {
96            model: "mistral".to_string(),
97            status: "complete".to_string(),
98            digest: "abc123".to_string(),
99            total: 1000,
100            completed: 1000,
101        };
102
103        assert!(progress.is_complete());
104    }
105
106    #[test]
107    fn test_pull_progress_not_complete() {
108        let progress = PullProgress {
109            model: "mistral".to_string(),
110            status: "downloading".to_string(),
111            digest: "abc123".to_string(),
112            total: 1000,
113            completed: 500,
114        };
115
116        assert!(!progress.is_complete());
117    }
118
119    #[test]
120    fn test_pull_progress_zero_total() {
121        let progress = PullProgress {
122            model: "mistral".to_string(),
123            status: "downloading".to_string(),
124            digest: "abc123".to_string(),
125            total: 0,
126            completed: 0,
127        };
128
129        assert_eq!(progress.percentage(), 0.0);
130        assert!(!progress.is_complete());
131    }
132}