orbok_models/
download_plan.rs1use crate::readiness::{LocalFileStatus, ModelReadinessReport};
9use std::path::PathBuf;
10
11pub const DEFAULT_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
15pub const MAX_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum DownloadAction {
22 Skip,
24 Download,
26 Replace,
28 Retry,
30}
31
32impl DownloadAction {
33 pub fn requires_download(&self) -> bool {
35 !matches!(self, DownloadAction::Skip)
36 }
37}
38
39#[derive(Debug, Clone)]
43pub struct ModelFilePlan {
44 pub logical_name: &'static str,
45 pub relative_path: &'static str,
46 pub remote_url: &'static str,
47 pub local_status: LocalFileStatus,
48 pub action: DownloadAction,
49 pub temp_path_suffix: &'static str,
51}
52
53impl ModelFilePlan {
54 pub fn temp_path(&self, model_dir: &std::path::Path) -> PathBuf {
56 model_dir.join(format!("{}.part", self.relative_path))
57 }
58
59 pub fn final_path(&self, model_dir: &std::path::Path) -> PathBuf {
61 model_dir.join(self.relative_path)
62 }
63}
64
65#[derive(Debug, Clone)]
69pub struct DownloadPlan {
70 pub max_concurrent: usize,
72 pub files: Vec<ModelFilePlan>,
73}
74
75impl DownloadPlan {
76 pub fn files_to_download(&self) -> Vec<&ModelFilePlan> {
78 self.files
79 .iter()
80 .filter(|f| f.action.requires_download())
81 .collect()
82 }
83
84 pub fn files_to_skip(&self) -> Vec<&ModelFilePlan> {
86 self.files
87 .iter()
88 .filter(|f| f.action == DownloadAction::Skip)
89 .collect()
90 }
91
92 pub fn has_work(&self) -> bool {
94 self.files.iter().any(|f| f.action.requires_download())
95 }
96}
97
98const TOKENIZER_URL: &str =
101 "https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/tokenizer.json";
102const MODEL_URL: &str =
103 "https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/onnx/model.onnx";
104
105pub fn build_download_plan(report: &ModelReadinessReport) -> DownloadPlan {
113 let urls: std::collections::HashMap<&str, &str> = [
114 ("tokenizer.json", TOKENIZER_URL),
115 ("onnx/model.onnx", MODEL_URL),
116 ]
117 .into_iter()
118 .collect();
119
120 let files = report
121 .files
122 .iter()
123 .map(|f| {
124 let action = plan_action(&f.status);
125 ModelFilePlan {
126 logical_name: f.logical_name,
127 relative_path: f.relative_path,
128 remote_url: urls.get(f.relative_path).copied().unwrap_or(""),
129 local_status: f.status.clone(),
130 action,
131 temp_path_suffix: ".part",
132 }
133 })
134 .collect();
135
136 DownloadPlan {
137 max_concurrent: DEFAULT_MODEL_DOWNLOAD_CONCURRENCY,
138 files,
139 }
140}
141
142fn plan_action(status: &LocalFileStatus) -> DownloadAction {
143 match status {
144 LocalFileStatus::Ready => DownloadAction::Skip,
145 LocalFileStatus::Missing => DownloadAction::Download,
146 LocalFileStatus::Partial => DownloadAction::Retry,
147 LocalFileStatus::Invalid => DownloadAction::Replace,
148 LocalFileStatus::CannotCheck => DownloadAction::Skip, }
150}
151
152#[derive(Debug, Clone, PartialEq, Eq)]
156pub enum FileDownloadStatus {
157 Pending,
158 Checking,
159 Downloading,
160 Validating,
161 Complete,
162 Failed,
163 Skipped,
164}
165
166#[derive(Debug, Clone)]
168pub struct FileDownloadProgress {
169 pub relative_path: &'static str,
170 pub downloaded_bytes: u64,
171 pub total_bytes: Option<u64>,
172 pub status: FileDownloadStatus,
173}
174
175#[derive(Debug, Clone)]
177pub enum OverallDownloadProgress {
178 Known {
179 downloaded_bytes: u64,
180 total_bytes: u64,
181 },
182 Step {
183 completed_files: usize,
184 total_files: usize,
185 },
186 Indeterminate,
187}
188
189#[derive(Debug, Clone)]
191pub enum FriendlyDownloadProblem {
192 NetworkUnavailable,
193 ServerBusy,
194 NotEnoughSpace,
195 CannotWriteFiles,
196 CannotCheckFiles,
197 ValidationFailed,
198 Unexpected,
199}
200
201impl FriendlyDownloadProblem {
202 pub fn user_message(&self) -> &'static str {
204 match self {
205 FriendlyDownloadProblem::NetworkUnavailable => {
206 "Download did not finish. Please check your connection and try again."
207 }
208 FriendlyDownloadProblem::ServerBusy => {
209 "The download is taking longer than expected. Please try again later."
210 }
211 FriendlyDownloadProblem::NotEnoughSpace => {
212 "More space is needed to finish the download."
213 }
214 FriendlyDownloadProblem::CannotWriteFiles => {
215 "orbok could not save the search helper files here. Please choose another location or check folder permissions."
216 }
217 FriendlyDownloadProblem::CannotCheckFiles => {
218 "orbok could not check the search helper files. Please choose the folder again."
219 }
220 FriendlyDownloadProblem::ValidationFailed => {
221 "Some downloaded files could not be used. orbok can download them again."
222 }
223 FriendlyDownloadProblem::Unexpected => "Download did not finish. Please try again.",
224 }
225 }
226}