Skip to main content

haagenti_network/
config.rs

1//! Network configuration
2
3use serde::{Deserialize, Serialize};
4use std::time::Duration;
5
6/// CDN endpoint configuration
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct CdnEndpoint {
9    /// Base URL for the endpoint
10    pub url: String,
11    /// Priority (lower = preferred)
12    pub priority: u32,
13    /// Geographic region (for latency-based selection)
14    pub region: Option<String>,
15    /// Whether this endpoint supports range requests
16    pub supports_range: bool,
17    /// Maximum concurrent connections
18    pub max_connections: usize,
19    /// Endpoint-specific headers
20    pub headers: Vec<(String, String)>,
21}
22
23impl CdnEndpoint {
24    /// Create a new CDN endpoint
25    pub fn new(url: impl Into<String>) -> Self {
26        Self {
27            url: url.into(),
28            priority: 0,
29            region: None,
30            supports_range: true,
31            max_connections: 8,
32            headers: Vec::new(),
33        }
34    }
35
36    /// Set priority
37    pub fn with_priority(mut self, priority: u32) -> Self {
38        self.priority = priority;
39        self
40    }
41
42    /// Set region
43    pub fn with_region(mut self, region: impl Into<String>) -> Self {
44        self.region = Some(region.into());
45        self
46    }
47
48    /// Add header
49    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
50        self.headers.push((key.into(), value.into()));
51        self
52    }
53}
54
55/// Retry configuration
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct RetryConfig {
58    /// Maximum retry attempts
59    pub max_retries: u32,
60    /// Initial backoff duration
61    pub initial_backoff: Duration,
62    /// Maximum backoff duration
63    pub max_backoff: Duration,
64    /// Backoff multiplier
65    pub multiplier: f64,
66    /// Add jitter to backoff
67    pub jitter: bool,
68}
69
70impl Default for RetryConfig {
71    fn default() -> Self {
72        Self {
73            max_retries: 3,
74            initial_backoff: Duration::from_millis(100),
75            max_backoff: Duration::from_secs(10),
76            multiplier: 2.0,
77            jitter: true,
78        }
79    }
80}
81
82/// Network loader configuration
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct NetworkConfig {
85    /// CDN endpoints (ordered by preference)
86    pub endpoints: Vec<CdnEndpoint>,
87    /// Request timeout
88    pub timeout: Duration,
89    /// Connect timeout
90    pub connect_timeout: Duration,
91    /// Maximum concurrent downloads
92    pub max_concurrent: usize,
93    /// Chunk size for range requests
94    pub chunk_size: usize,
95    /// Retry configuration
96    pub retry: RetryConfig,
97    /// Enable compression
98    pub compression: bool,
99    /// User agent string
100    pub user_agent: String,
101    /// Cache directory
102    pub cache_dir: Option<std::path::PathBuf>,
103    /// Maximum cache size (bytes)
104    pub max_cache_size: u64,
105    /// Enable bandwidth monitoring
106    pub monitor_bandwidth: bool,
107    /// Minimum acceptable bandwidth (bytes/sec)
108    pub min_bandwidth: u64,
109}
110
111impl Default for NetworkConfig {
112    fn default() -> Self {
113        Self {
114            endpoints: Vec::new(),
115            timeout: Duration::from_secs(30),
116            connect_timeout: Duration::from_secs(10),
117            max_concurrent: 4,
118            chunk_size: 1024 * 1024, // 1MB chunks
119            retry: RetryConfig::default(),
120            compression: true,
121            user_agent: format!("haagenti-network/{}", env!("CARGO_PKG_VERSION")),
122            cache_dir: None,
123            max_cache_size: 10 * 1024 * 1024 * 1024, // 10GB
124            monitor_bandwidth: true,
125            min_bandwidth: 1024 * 1024, // 1MB/s minimum
126        }
127    }
128}
129
130impl NetworkConfig {
131    /// Add a CDN endpoint
132    pub fn with_endpoint(mut self, endpoint: CdnEndpoint) -> Self {
133        self.endpoints.push(endpoint);
134        self
135    }
136
137    /// Set cache directory
138    pub fn with_cache_dir(mut self, path: impl Into<std::path::PathBuf>) -> Self {
139        self.cache_dir = Some(path.into());
140        self
141    }
142
143    /// Set maximum concurrent downloads
144    pub fn with_max_concurrent(mut self, max: usize) -> Self {
145        self.max_concurrent = max;
146        self
147    }
148
149    /// Configure for Hugging Face Hub
150    pub fn huggingface_hub() -> Self {
151        Self::default()
152            .with_endpoint(
153                CdnEndpoint::new("https://huggingface.co")
154                    .with_priority(0)
155                    .with_region("global"),
156            )
157            .with_endpoint(
158                CdnEndpoint::new("https://cdn-lfs.huggingface.co")
159                    .with_priority(1)
160                    .with_region("global"),
161            )
162    }
163
164    /// Configure for Civitai
165    pub fn civitai() -> Self {
166        Self::default()
167            .with_endpoint(CdnEndpoint::new("https://civitai.com/api/download").with_priority(0))
168    }
169
170    /// Configure for custom CDN
171    pub fn custom_cdn(base_url: impl Into<String>) -> Self {
172        Self::default().with_endpoint(CdnEndpoint::new(base_url))
173    }
174}