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, PartialEq, Eq)]
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    /// Git LFS OID (SHA256 hash from `HuggingFace` tree API).
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub oid: Option<String>,
60}
61
62impl ResolvedFile {
63    /// Create a new resolved file.
64    pub fn new(path: impl Into<String>) -> Self {
65        Self {
66            path: path.into(),
67            size: None,
68            oid: None,
69        }
70    }
71
72    /// Create a new resolved file with size.
73    pub fn with_size(path: impl Into<String>, size: u64) -> Self {
74        Self {
75            path: path.into(),
76            size: Some(size),
77            oid: None,
78        }
79    }
80
81    /// Create a new resolved file with size and OID.
82    pub fn with_size_and_oid(path: impl Into<String>, size: u64, oid: Option<String>) -> Self {
83        Self {
84            path: path.into(),
85            size: Some(size),
86            oid,
87        }
88    }
89}
90
91// ============================================================================
92// Resolver Trait
93// ============================================================================
94
95/// Trait for resolving quantization-specific files from a model repository.
96///
97/// Implementations handle the specifics of querying APIs (`HuggingFace`, etc.)
98/// to find GGUF files matching a requested quantization.
99///
100/// # Usage
101///
102/// ```ignore
103/// let resolver: Arc<dyn QuantizationResolver> = /* ... */;
104/// let resolution = resolver.resolve("unsloth/Llama-3-GGUF", Quantization::Q4KM).await?;
105/// println!("Found {} files", resolution.file_count());
106/// ```
107#[async_trait]
108pub trait QuantizationResolver: Send + Sync {
109    /// Resolve files for a specific quantization.
110    ///
111    /// Returns a `Resolution` containing the list of files to download
112    /// and metadata about the resolution.
113    async fn resolve(
114        &self,
115        repo_id: &str,
116        quantization: Quantization,
117    ) -> Result<Resolution, DownloadError>;
118
119    /// List all available quantizations in a repository.
120    async fn list_available(&self, repo_id: &str) -> Result<Vec<Quantization>, DownloadError>;
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    #[test]
128    fn test_resolution_methods() {
129        let resolution = Resolution {
130            quantization: Quantization::Q4KM,
131            files: vec![
132                ResolvedFile::with_size("model.gguf", 1000),
133                ResolvedFile::with_size("model-00001-of-00002.gguf", 500),
134            ],
135            is_sharded: true,
136        };
137
138        assert_eq!(resolution.file_count(), 2);
139        assert_eq!(resolution.total_size(), Some(1500));
140        assert_eq!(resolution.first_file(), Some("model.gguf"));
141        assert_eq!(
142            resolution.filenames(),
143            vec!["model.gguf", "model-00001-of-00002.gguf"]
144        );
145    }
146
147    #[test]
148    fn test_resolved_file_creation() {
149        let file = ResolvedFile::new("test.gguf");
150        assert_eq!(file.path, "test.gguf");
151        assert_eq!(file.size, None);
152
153        let file_with_size = ResolvedFile::with_size("test.gguf", 1024);
154        assert_eq!(file_with_size.size, Some(1024));
155    }
156}
gglib_core/ports/download.rs

gglib_core/ports/
download.rs
