use crate::readiness::{LocalFileStatus, ModelReadinessReport};
use std::path::PathBuf;
pub const DEFAULT_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
pub const MAX_MODEL_DOWNLOAD_CONCURRENCY: usize = 2;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DownloadAction {
Skip,
Download,
Replace,
Retry,
}
impl DownloadAction {
pub fn requires_download(&self) -> bool {
!matches!(self, DownloadAction::Skip)
}
}
#[derive(Debug, Clone)]
pub struct ModelFilePlan {
pub logical_name: &'static str,
pub relative_path: &'static str,
pub remote_url: &'static str,
pub local_status: LocalFileStatus,
pub action: DownloadAction,
pub temp_path_suffix: &'static str,
}
impl ModelFilePlan {
pub fn temp_path(&self, model_dir: &std::path::Path) -> PathBuf {
model_dir.join(format!("{}.part", self.relative_path))
}
pub fn final_path(&self, model_dir: &std::path::Path) -> PathBuf {
model_dir.join(self.relative_path)
}
}
#[derive(Debug, Clone)]
pub struct DownloadPlan {
pub max_concurrent: usize,
pub files: Vec<ModelFilePlan>,
}
impl DownloadPlan {
pub fn files_to_download(&self) -> Vec<&ModelFilePlan> {
self.files
.iter()
.filter(|f| f.action.requires_download())
.collect()
}
pub fn files_to_skip(&self) -> Vec<&ModelFilePlan> {
self.files
.iter()
.filter(|f| f.action == DownloadAction::Skip)
.collect()
}
pub fn has_work(&self) -> bool {
self.files.iter().any(|f| f.action.requires_download())
}
}
const TOKENIZER_URL: &str =
"https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/tokenizer.json";
const MODEL_URL: &str =
"https://huggingface.co/intfloat/multilingual-e5-small/resolve/main/onnx/model.onnx";
pub fn build_download_plan(report: &ModelReadinessReport) -> DownloadPlan {
let urls: std::collections::HashMap<&str, &str> = [
("tokenizer.json", TOKENIZER_URL),
("onnx/model.onnx", MODEL_URL),
]
.into_iter()
.collect();
let files = report
.files
.iter()
.map(|f| {
let action = plan_action(&f.status);
ModelFilePlan {
logical_name: f.logical_name,
relative_path: f.relative_path,
remote_url: urls.get(f.relative_path).copied().unwrap_or(""),
local_status: f.status.clone(),
action,
temp_path_suffix: ".part",
}
})
.collect();
DownloadPlan {
max_concurrent: DEFAULT_MODEL_DOWNLOAD_CONCURRENCY,
files,
}
}
fn plan_action(status: &LocalFileStatus) -> DownloadAction {
match status {
LocalFileStatus::Ready => DownloadAction::Skip,
LocalFileStatus::Missing => DownloadAction::Download,
LocalFileStatus::Partial => DownloadAction::Retry,
LocalFileStatus::Invalid => DownloadAction::Replace,
LocalFileStatus::CannotCheck => DownloadAction::Skip, }
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FileDownloadStatus {
Pending,
Checking,
Downloading,
Validating,
Complete,
Failed,
Skipped,
}
#[derive(Debug, Clone)]
pub struct FileDownloadProgress {
pub relative_path: &'static str,
pub downloaded_bytes: u64,
pub total_bytes: Option<u64>,
pub status: FileDownloadStatus,
}
#[derive(Debug, Clone)]
pub enum OverallDownloadProgress {
Known {
downloaded_bytes: u64,
total_bytes: u64,
},
Step {
completed_files: usize,
total_files: usize,
},
Indeterminate,
}
#[derive(Debug, Clone)]
pub enum FriendlyDownloadProblem {
NetworkUnavailable,
ServerBusy,
NotEnoughSpace,
CannotWriteFiles,
CannotCheckFiles,
ValidationFailed,
Unexpected,
}
impl FriendlyDownloadProblem {
pub fn user_message(&self) -> &'static str {
match self {
FriendlyDownloadProblem::NetworkUnavailable => {
"Download did not finish. Please check your connection and try again."
}
FriendlyDownloadProblem::ServerBusy => {
"The download is taking longer than expected. Please try again later."
}
FriendlyDownloadProblem::NotEnoughSpace => {
"More space is needed to finish the download."
}
FriendlyDownloadProblem::CannotWriteFiles => {
"orbok could not save the search helper files here. Please choose another location or check folder permissions."
}
FriendlyDownloadProblem::CannotCheckFiles => {
"orbok could not check the search helper files. Please choose the folder again."
}
FriendlyDownloadProblem::ValidationFailed => {
"Some downloaded files could not be used. orbok can download them again."
}
FriendlyDownloadProblem::Unexpected => "Download did not finish. Please try again.",
}
}
}