Skip to main content

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 total size if all file sizes are known.
28    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    /// Get the number of files.
34    pub const fn file_count(&self) -> usize {
35        self.files.len()
36    }
37}
38
39/// A single resolved file.
40#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
41pub struct ResolvedFile {
42    /// Path within the repository.
43    pub path: String,
44    /// Size in bytes (if available from API).
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub size: Option<u64>,
47    /// Git LFS OID (SHA256 hash from `HuggingFace` tree API).
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub oid: Option<String>,
50}
51
52impl ResolvedFile {
53    /// Create a new resolved file.
54    pub fn new(path: impl Into<String>) -> Self {
55        Self {
56            path: path.into(),
57            size: None,
58            oid: None,
59        }
60    }
61
62    /// Create a new resolved file with size.
63    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    /// Create a new resolved file with size and OID.
72    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// ============================================================================
82// Resolver Trait
83// ============================================================================
84
85/// Trait for resolving quantization-specific files from a model repository.
86///
87/// Implementations handle the specifics of querying APIs (`HuggingFace`, etc.)
88/// to find GGUF files matching a requested quantization.
89///
90/// # Usage
91///
92/// ```ignore
93/// let resolver: Arc<dyn QuantizationResolver> = /* ... */;
94/// let resolution = resolver.resolve("unsloth/Llama-3-GGUF", Quantization::Q4KM).await?;
95/// println!("Found {} files", resolution.file_count());
96/// ```
97#[async_trait]
98pub trait QuantizationResolver: Send + Sync {
99    /// Resolve files for a specific quantization.
100    ///
101    /// Returns a `Resolution` containing the list of files to download
102    /// and metadata about the resolution.
103    async fn resolve(
104        &self,
105        repo_id: &str,
106        quantization: Quantization,
107    ) -> Result<Resolution, DownloadError>;
108
109    /// List all available quantizations in a repository.
110    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}