gglib_core/ports/
download.rs

1//! Download port definitions (trait abstractions).
2//!
3//! This module contains trait definitions for download-related operations
4//! that abstract away infrastructure concerns.
5
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8
9use crate::download::{DownloadError, Quantization};
10
11// ============================================================================
12// Resolution Types
13// ============================================================================
14
15/// Result of resolving files for a quantization.
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Resolution {
18    /// The resolved quantization type.
19    pub quantization: Quantization,
20    /// List of files to download (sorted for sharded files).
21    pub files: Vec<ResolvedFile>,
22    /// Whether this is a sharded (multi-part) download.
23    pub is_sharded: bool,
24}
25
26impl Resolution {
27    /// Get filenames as a simple list.
28    pub fn filenames(&self) -> Vec<String> {
29        self.files.iter().map(|f| f.path.clone()).collect()
30    }
31
32    /// Get total size if all file sizes are known.
33    pub fn total_size(&self) -> Option<u64> {
34        let sizes: Option<Vec<u64>> = self.files.iter().map(|f| f.size).collect();
35        sizes.map(|s| s.iter().sum())
36    }
37
38    /// Get the first file path (used for database registration of sharded models).
39    pub fn first_file(&self) -> Option<&str> {
40        self.files.first().map(|f| f.path.as_str())
41    }
42
43    /// Get the number of files.
44    pub const fn file_count(&self) -> usize {
45        self.files.len()
46    }
47}
48
49/// A single resolved file.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51pub struct ResolvedFile {
52    /// Path within the repository.
53    pub path: String,
54    /// Size in bytes (if available from API).
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub size: Option<u64>,
57}
58
59impl ResolvedFile {
60    /// Create a new resolved file.
61    pub fn new(path: impl Into<String>) -> Self {
62        Self {
63            path: path.into(),
64            size: None,
65        }
66    }
67
68    /// Create a new resolved file with size.
69    pub fn with_size(path: impl Into<String>, size: u64) -> Self {
70        Self {
71            path: path.into(),
72            size: Some(size),
73        }
74    }
75}
76
77// ============================================================================
78// Resolver Trait
79// ============================================================================
80
81/// Trait for resolving quantization-specific files from a model repository.
82///
83/// Implementations handle the specifics of querying APIs (`HuggingFace`, etc.)
84/// to find GGUF files matching a requested quantization.
85///
86/// # Usage
87///
88/// ```ignore
89/// let resolver: Arc<dyn QuantizationResolver> = /* ... */;
90/// let resolution = resolver.resolve("unsloth/Llama-3-GGUF", Quantization::Q4KM).await?;
91/// println!("Found {} files", resolution.file_count());
92/// ```
93#[async_trait]
94pub trait QuantizationResolver: Send + Sync {
95    /// Resolve files for a specific quantization.
96    ///
97    /// Returns a `Resolution` containing the list of files to download
98    /// and metadata about the resolution.
99    async fn resolve(
100        &self,
101        repo_id: &str,
102        quantization: Quantization,
103    ) -> Result<Resolution, DownloadError>;
104
105    /// List all available quantizations in a repository.
106    async fn list_available(&self, repo_id: &str) -> Result<Vec<Quantization>, DownloadError>;
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_resolution_methods() {
115        let resolution = Resolution {
116            quantization: Quantization::Q4KM,
117            files: vec![
118                ResolvedFile::with_size("model.gguf", 1000),
119                ResolvedFile::with_size("model-00001-of-00002.gguf", 500),
120            ],
121            is_sharded: true,
122        };
123
124        assert_eq!(resolution.file_count(), 2);
125        assert_eq!(resolution.total_size(), Some(1500));
126        assert_eq!(resolution.first_file(), Some("model.gguf"));
127        assert_eq!(
128            resolution.filenames(),
129            vec!["model.gguf", "model-00001-of-00002.gguf"]
130        );
131    }
132
133    #[test]
134    fn test_resolved_file_creation() {
135        let file = ResolvedFile::new("test.gguf");
136        assert_eq!(file.path, "test.gguf");
137        assert_eq!(file.size, None);
138
139        let file_with_size = ResolvedFile::with_size("test.gguf", 1024);
140        assert_eq!(file_with_size.size, Some(1024));
141    }
142}
gglib_core/ports/download.rs

gglib_core/ports/
download.rs
