Skip to main content

orbok_models/
download_plan.rs

1//! Download plan: maps readiness → actions (RFC-043 §10).
2//!
3//! A `DownloadPlan` tells the download worker exactly which files to
4//! fetch, which to skip, and what concurrency limit to apply. It is
5//! always derived from a fresh `ModelReadinessReport`; nothing is
6//! downloaded without a plan.
7
8use crate::readiness::{LocalFileStatus, ModelReadinessReport};
9use std::path::PathBuf;
10
11// ── Constants ─────────────────────────────────────────────────────────
12
13/// Maximum concurrent file downloads (RFC-043 §11.1).
14pub const DEFAULT_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
15pub const MAX_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
16
17// ── Download action ───────────────────────────────────────────────────
18
19/// What to do with one model file (RFC-043 §10.3).
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum DownloadAction {
22    /// File is already valid — skip.
23    Skip,
24    /// File is missing — download fresh.
25    Download,
26    /// File is invalid — replace (delete and re-download).
27    Replace,
28    /// File is partial — restart the download.
29    Retry,
30}
31
32impl DownloadAction {
33    /// Whether this action requires network activity.
34    pub fn requires_download(&self) -> bool {
35        !matches!(self, DownloadAction::Skip)
36    }
37}
38
39// ── Per-file plan ─────────────────────────────────────────────────────
40
41/// Plan for one required model file (RFC-043 §10.2).
42#[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    /// Temporary path used during download (RFC-043 §9.1).
50    pub temp_path_suffix: &'static str,
51}
52
53impl ModelFilePlan {
54    /// Full path for the temporary `.part` file.
55    pub fn temp_path(&self, model_dir: &std::path::Path) -> PathBuf {
56        model_dir.join(format!("{}.part", self.relative_path))
57    }
58
59    /// Full final path for the completed validated file.
60    pub fn final_path(&self, model_dir: &std::path::Path) -> PathBuf {
61        model_dir.join(self.relative_path)
62    }
63}
64
65// ── Download plan ─────────────────────────────────────────────────────
66
67/// Complete download plan for a model (RFC-043 §10.1).
68#[derive(Debug, Clone)]
69pub struct DownloadPlan {
70    /// Maximum concurrent file downloads (RFC-043 §11).
71    pub max_concurrent: usize,
72    pub files: Vec<ModelFilePlan>,
73}
74
75impl DownloadPlan {
76    /// Files that actually require download work.
77    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    /// Files that will be skipped.
85    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    /// Whether anything actually needs downloading.
93    pub fn has_work(&self) -> bool {
94        self.files.iter().any(|f| f.action.requires_download())
95    }
96}
97
98// ── File URLs (current recommended model) ────────────────────────────
99
100const 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
105// ── Plan builder ──────────────────────────────────────────────────────
106
107/// Build a download plan from a fresh readiness report (RFC-043 §10.4).
108///
109/// Called after every `check_model_readiness`. The plan is always
110/// derived from the current on-disk state; nothing is assumed from
111/// a previous session.
112pub 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, // surface via friendly problem
149    }
150}
151
152// ── Download progress types ───────────────────────────────────────────
153
154/// Per-file download status (RFC-043 §13.1).
155#[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/// Per-file download progress (RFC-043 §13.1).
167#[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/// Overall combined progress (RFC-043 §13.2).
176#[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/// Friendly download problem — safe to show to users (RFC-043 §16.4).
190#[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    /// User-facing copy (RFC-043 §20, avoiding technical terms).
203    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}