gglib_core/ports/
download.rs1use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8
9use crate::download::{DownloadError, Quantization};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Resolution {
18 pub quantization: Quantization,
20 pub files: Vec<ResolvedFile>,
22 pub is_sharded: bool,
24}
25
26impl Resolution {
27 pub fn total_size(&self) -> Option<u64> {
29 let sizes: Option<Vec<u64>> = self.files.iter().map(|f| f.size).collect();
30 sizes.map(|s| s.iter().sum())
31 }
32
33 pub const fn file_count(&self) -> usize {
35 self.files.len()
36 }
37}
38
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
41pub struct ResolvedFile {
42 pub path: String,
44 #[serde(skip_serializing_if = "Option::is_none")]
46 pub size: Option<u64>,
47 #[serde(skip_serializing_if = "Option::is_none")]
49 pub oid: Option<String>,
50}
51
52impl ResolvedFile {
53 pub fn new(path: impl Into<String>) -> Self {
55 Self {
56 path: path.into(),
57 size: None,
58 oid: None,
59 }
60 }
61
62 pub fn with_size(path: impl Into<String>, size: u64) -> Self {
64 Self {
65 path: path.into(),
66 size: Some(size),
67 oid: None,
68 }
69 }
70
71 pub fn with_size_and_oid(path: impl Into<String>, size: u64, oid: Option<String>) -> Self {
73 Self {
74 path: path.into(),
75 size: Some(size),
76 oid,
77 }
78 }
79}
80
81#[async_trait]
98pub trait QuantizationResolver: Send + Sync {
99 async fn resolve(
104 &self,
105 repo_id: &str,
106 quantization: Quantization,
107 ) -> Result<Resolution, DownloadError>;
108
109 async fn list_available(&self, repo_id: &str) -> Result<Vec<Quantization>, DownloadError>;
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116
117 #[test]
118 fn test_resolution_methods() {
119 let resolution = Resolution {
120 quantization: Quantization::Q4KM,
121 files: vec![
122 ResolvedFile::with_size("model.gguf", 1000),
123 ResolvedFile::with_size("model-00001-of-00002.gguf", 500),
124 ],
125 is_sharded: true,
126 };
127
128 assert_eq!(resolution.file_count(), 2);
129 assert_eq!(resolution.total_size(), Some(1500));
130 }
131
132 #[test]
133 fn test_resolved_file_creation() {
134 let file = ResolvedFile::new("test.gguf");
135 assert_eq!(file.path, "test.gguf");
136 assert_eq!(file.size, None);
137
138 let file_with_size = ResolvedFile::with_size("test.gguf", 1024);
139 assert_eq!(file_with_size.size, Some(1024));
140 }
141}