download with progress

gau-nernst · gau-nernst · commit 42068ad94031 · 2025-05-15T12:46:46.000+08:00
diff --git a/extensions/model-extension/package.json b/extensions/model-extension/package.json
@@ -26,6 +26,7 @@
   ],
   "dependencies": {
     "@janhq/core": "../../core/package.tgz",
+    "@tauri-apps/api": "^2.5.0",
     "ky": "^1.7.2",
     "p-queue": "^8.0.1"
   },
diff --git a/extensions/model-extension/src/index.ts b/extensions/model-extension/src/index.ts
@@ -1,3 +1,5 @@
+import { invoke } from '@tauri-apps/api/core';
+import { listen } from '@tauri-apps/api/event';
 import {
   ModelExtension,
   Model,
@@ -9,6 +11,7 @@ import {
   ModelSource,
   extractInferenceParams,
   extractModelLoadParams,
+  events,
 } from '@janhq/core'
 import { scanModelsFolder } from './legacy/model-json'
 import { deleteModelFiles } from './legacy/delete'
@@ -26,6 +29,14 @@ type Data<T> = {
   data: T[]
 }
 
+type DownloadInfo = {
+  task_id: string
+  total_size: number
+  downloaded_size: number
+  download_type: string
+  event_type: string
+}
+
 /**
  * A extension for models
  */
@@ -64,6 +75,29 @@ export default class JanModelExtension extends ModelExtension {
       this.updateCortexConfig({ huggingface_token: huggingfaceToken })
     }
 
+    // listen to tauri events
+    // TODO: move this to core? i.e. forward tauri events to core events
+    listen<DownloadInfo>('download', (event) => {
+      let payload = event.payload
+      let eventName = {
+        Updated: 'onFileDownloadUpdate',
+        Error: 'onFileDownloadError',
+        Success: 'onFileDownloadSuccess',
+        Stopped: 'onFileDownloadStopped',
+        Started: 'onFileDownloadStarted',
+      }[payload.event_type]
+
+      events.emit(eventName, {
+        modelId: payload.task_id,
+        percent: payload.downloaded_size / payload.total_size,
+        size: {
+          transferred: payload.downloaded_size,
+          total: payload.total_size,
+        },
+        downloadType: payload.download_type,
+      })
+    })
+
     // Sync with cortexsohub
     this.fetchCortexsoModels()
   }
@@ -92,6 +126,19 @@ export default class JanModelExtension extends ModelExtension {
    * @returns A Promise that resolves when the model is downloaded.
    */
   async pullModel(model: string, id?: string, name?: string): Promise<void> {
+    if (id == null && name == null) {
+      let [modelName, branch] = model.split(":")
+      return invoke<void>("get_jan_data_folder_path").then((path) => {
+        // cortexso format
+        return invoke<void>("download_hf_repo", {
+          taskId: model,
+          repoId: `cortexso/${modelName}`,
+          branch: branch,
+          saveDir: `${path}/models/cortex.so/${modelName}/${branch}`
+        })
+      }).catch(console.error)
+    }
+
     /**
      * Sending POST to /models/pull/{id} endpoint to pull the model
      */
diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml
@@ -30,7 +30,7 @@ rand = "0.8"
 tauri-plugin-http = { version = "2", features = ["unsafe-headers"] }
 tauri-plugin-store = "2"
 hyper = { version = "0.14", features = ["server"] }
-reqwest = { version = "0.11", features = ["json", "blocking"] }
+reqwest = { version = "0.11", features = ["json", "blocking", "stream"] }
 tokio = { version = "1", features = ["full"] }
 rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "main", features = [
     "client",
@@ -40,6 +40,7 @@ rmcp = { git = "https://github.com/modelcontextprotocol/rust-sdk", branch = "mai
 ] }
 uuid = { version = "1.7", features = ["v4"] }
 env = "1.0.1"
+futures-util = "0.3.31"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 tauri-plugin-updater = "2"
diff --git a/src-tauri/src/core/mod.rs b/src-tauri/src/core/mod.rs
@@ -6,4 +6,3 @@ pub mod setup;
 pub mod state;
 pub mod threads;
 pub mod utils;
-pub mod download;
diff --git a/src-tauri/src/core/utils/download.rs b/src-tauri/src/core/utils/download.rs
@@ -1,20 +1,44 @@
-use std::fs::File;
-use std::io::Write;
+use futures_util::StreamExt;
 use std::path::Path;
 use std::time::Duration;
+use tokio::fs::File;
+use tokio::io::AsyncWriteExt;
 
-pub async fn download(url: &str, save_path: &Path) -> Result<(), Box<dyn std::error::Error>> {
+// this is to emulate the current way of downloading files by Cortex + Jan
+// we can change this later
+#[derive(serde::Serialize, Clone, Debug)]
+pub enum DownloadEventType {
+    Started,
+    Updated,
+    Success,
+    Error,
+    Stopped,
+}
+
+#[derive(serde::Serialize, Clone, Debug)]
+pub struct DownloadEvent {
+    pub task_id: String,
+    pub total_size: u64,
+    pub downloaded_size: u64,
+    pub download_type: String, // TODO: make this an enum as well
+    pub event_type: DownloadEventType,
+}
+
+pub async fn download<F>(
+    url: &str,
+    save_path: &Path,
+    mut callback: Option<F>,
+) -> Result<(), Box<dyn std::error::Error>>
+where
+    // F: FnMut(u64) + Send + 'static,
+    F: FnMut(u64),
+{
     let client = reqwest::Client::builder()
         .http2_keep_alive_timeout(Duration::from_secs(15))
         .build()?;
 
-    let mut resp = client
-        .get(url)
-        .header("User-Agent", "rust-reqwest/huggingface-downloader")
-        .send()
-        .await?;
-
-    // Check if request was successful
+    // NOTE: might want to add User-Agent header
+    let resp = client.get(url).send().await?;
     if !resp.status().is_success() {
         return Err(format!(
             "Failed to download: HTTP status {}, {}",
@@ -23,17 +47,26 @@ pub async fn download(url: &str, save_path: &Path) -> Result<(), Box<dyn std::er
         )
         .into());
     }
+
     // Create parent directories if they don't exist
     if let Some(parent) = save_path.parent() {
         if !parent.exists() {
             std::fs::create_dir_all(parent)?;
         }
     }
-    let mut file = File::create(save_path)?;
+    let mut file = File::create(save_path).await?;
 
-    while let Some(chunk) = resp.chunk().await? {
-        file.write_all(&chunk)?;
-    }
+    // write chunk to file, and call callback if needed (e.g. download progress)
+    let mut stream = resp.bytes_stream();
+    while let Some(chunk) = stream.next().await {
+        let chunk = chunk?;
+        file.write_all(&chunk).await?;
 
+        // NOTE: might want to reduce frequency of callback e.g. every 1MB
+        if let Some(cb) = callback.as_mut() {
+            cb(chunk.len() as u64);
+        }
+    }
+    file.flush().await?;
     Ok(())
 }
diff --git a/src-tauri/src/core/utils/hf.rs b/src-tauri/src/core/utils/hf.rs
@@ -1,26 +1,103 @@
+use crate::core::utils::download::{download, DownloadEvent, DownloadEventType};
 use std::path::Path;
+use std::sync::{Arc, Mutex};
+use tauri::{AppHandle, Emitter};
 
-pub async fn download_file(
+// task_id is only for compatibility with current Jan+Cortex we can remove it later
+// should we disallow custom save_dir? i.e. all HF models will be downloaded to
+// <app_dir>/models/<repo_id>/<branch>?
+// (only cortexso uses branch. normally people don't use it)
+#[tauri::command]
+pub async fn download_hf_repo(
+    app: AppHandle,
+    task_id: &str,
     repo_id: &str,
     branch: &str,
-    file_path: &str,
-    save_path: &Path,
-) -> Result<(), Box<dyn std::error::Error>> {
-    let url = format!(
-        "https://huggingface.co/{}/resolve/{}/{}",
-        repo_id, branch, file_path
-    );
-    super::download::download(&url, save_path).await
+    save_dir: &Path,
+) -> Result<(), String> {
+    // TODO: check if it has been downloaded
+    let files = list_files(repo_id, branch)
+        .await
+        .map_err(|e| format!("Failed to list files {}", e))?;
+
+    // obtain total download size. emit download started event
+    let info = DownloadEvent {
+        task_id: task_id.to_string(),
+        total_size: files.iter().map(|f| f.size).sum(),
+        downloaded_size: 0,
+        download_type: "Model".to_string(),
+        event_type: DownloadEventType::Started,
+    };
+    let info_arc = Arc::new(Mutex::new(info));
+    log::info!("Start download repo_id: {} branch: {}", repo_id, branch);
+    app.emit("download", info_arc.lock().unwrap().clone())
+        .unwrap();
+
+    let download_result = async {
+        // NOTE: currently we are downloading sequentially. we can spawn tokio tasks
+        // to download files in parallel.
+        for file in files {
+            let url = format!(
+                "https://huggingface.co/{}/resolve/{}/{}",
+                repo_id, branch, file.path
+            );
+            let full_path = save_dir.join(&file.path);
+
+            // update download progress. clone app handle and info_arc
+            // to move them into the closure
+            let callback = {
+                let app = app.clone();
+                let info_arc = Arc::clone(&info_arc);
+                move |size| {
+                    let mut info = info_arc.lock().unwrap();
+                    info.event_type = DownloadEventType::Updated;
+                    info.downloaded_size += size;
+                    app.emit("download", info.clone()).unwrap();
+                }
+            };
+            download(&url, &full_path, Some(callback))
+                .await
+                .map_err(|e| format!("Failed to download file {}: {}", file.path, e))?;
+        }
+        Ok(())
+    }
+    .await;
+    match download_result {
+        Ok(_) => {
+            let mut info = info_arc.lock().unwrap();
+            info.event_type = DownloadEventType::Success;
+            log::info!("Finished download repo_id: {} branch: {}", repo_id, branch);
+            app.emit("download", info.clone()).unwrap();
+            Ok(())
+        }
+        Err(e) => {
+            // on failure, remove the directory to restore the original state
+            if save_dir.exists() {
+                // NOTE: we don't check error here
+                let _ = std::fs::remove_dir_all(save_dir);
+            }
+            log::info!("Failed to download repo_id: {} branch: {}", repo_id, branch);
+            // TODO: check what cortex and Jan does on download error
+            // app.emit("download", info.clone()).unwrap();
+            Err(e)
+        }
+    }
 }
 
-pub async fn list_files(
+#[derive(Debug)]
+struct FileInfo {
+    path: String,
+    size: u64,
+}
+
+async fn list_files(
     repo_id: &str,
     branch: &str,
-) -> Result<Vec<String>, Box<dyn std::error::Error>> {
+) -> Result<Vec<FileInfo>, Box<dyn std::error::Error>> {
     let mut files = vec![];
     let client = reqwest::Client::new();
 
-    // do recursion with a stack
+    // DFS with a stack (similar to recursion)
     let mut stack = vec!["".to_string()];
     while let Some(directory) = stack.pop() {
         let url = format!(
@@ -38,14 +115,19 @@ pub async fn list_files(
             .into());
         }
 
+        // this struct is only used for internal deserialization
         #[derive(serde::Deserialize)]
         struct Item {
             r#type: String,
             path: String,
+            size: u64,
         }
         for item in resp.json::<Vec<Item>>().await?.into_iter() {
             match item.r#type.as_str() {
-                "file" => files.push(item.path),
+                "file" => files.push(FileInfo {
+                    path: item.path,
+                    size: item.size,
+                }),
                 "directory" => stack.push(item.path),
                 _ => {}
             }
@@ -59,44 +141,8 @@ pub async fn list_files(
 mod tests {
     use super::*;
 
-    #[tokio::test]
-    async fn test_download_file() {
-        let repo_id = "openai-community/gpt2";
-        let branch = "main";
-        let file_path = "config.json";
-        let save_path = std::path::PathBuf::from("subdir/test_config.json");
-
-        if let Some(parent) = save_path.parent() {
-            if parent.exists() {
-                std::fs::remove_dir_all(parent).unwrap();
-            }
-        }
-
-        let result = download_file(repo_id, branch, file_path, &save_path).await;
-        assert!(result.is_ok(), "{}", result.unwrap_err());
-        assert!(save_path.exists());
-
-        // Read the file and verify its content
-        let file_content = std::fs::read_to_string(&save_path).unwrap();
-        let json_result: Result<serde_json::Value, _> = serde_json::from_str(&file_content);
-        assert!(json_result.is_ok(), "Downloaded file is not valid JSON");
-
-        if let Ok(json) = json_result {
-            assert!(json.is_object(), "JSON root should be an object");
-            assert_eq!(
-                json.get("model_type")
-                    .and_then(|v| v.as_str())
-                    .unwrap_or(""),
-                "gpt2",
-                "model_type should be gpt2"
-            );
-        }
-
-        // Clean up
-        // NOTE: this will not run if there are errors
-        // TODO: use tempfile crate instead
-        std::fs::remove_dir_all(save_path.parent().unwrap()).unwrap();
-    }
+    // TODO: add test for download_hf_repo (need to find a small repo)
+    // TODO: test when repo does not exist
 
     #[tokio::test]
     async fn test_list_files() {
@@ -107,7 +153,7 @@ mod tests {
         assert!(result.is_ok(), "{}", result.unwrap_err());
         let files = result.unwrap();
         assert!(
-            files.iter().any(|f| f == "config.json"),
+            files.iter().any(|f| f.path == "config.json"),
             "config.json should be in the list"
         );
     }
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
@@ -58,7 +58,9 @@ pub fn run() {
             core::threads::delete_message,
             core::threads::get_thread_assistant,
             core::threads::create_thread_assistant,
-            core::threads::modify_thread_assistant
+            core::threads::modify_thread_assistant,
+            // Download
+            core::utils::hf::download_hf_repo,
         ])
         .manage(AppState {
             app_token: Some(generate_app_token()),
