feat: support downloading multiple model files

crates/aim-downloader/src/https.rs

@@ -2,20 +2,36 @@ use std::{cmp::min, io::Error};
 
 use futures_util::StreamExt;
 use regex::Regex;
-use reqwest::Client;
+use reqwest::{header::HeaderMap, Client};
 use tokio_util::io::ReaderStream;
 
 use crate::{
     address::ParsedAddress,
     bar::WrappedBar,
     consts::*,
-    error::{DownloadError, ValidateError},
+    error::{DownloadError, HTTPHeaderError, ValidateError},
     hash::HashChecker,
     io,
 };
 
 pub struct HTTPSHandler;
 impl HTTPSHandler {
+    pub async fn head(input: &str) -> Result<HeaderMap, HTTPHeaderError> {
+        let parsed_address = ParsedAddress::parse_address(input, true);
+        let res = Client::new()
+            .head(input)
+            .header(
+                reqwest::header::USER_AGENT,
+                reqwest::header::HeaderValue::from_static(CLIENT_ID),
+            )
+            .basic_auth(parsed_address.username, Some(parsed_address.password))
+            .send()
+            .await
+            .map_err(|_| format!("Failed to HEAD from {}", &input))
+            .unwrap();
+        Ok(res.headers().clone())
+    }
+
     pub async fn get(
         input: &str,
         output: &str,
@@ -275,3 +291,10 @@ async fn get_links_works_when_typical() {
 
     assert_eq!(result[0], expected);
 }
+
+#[ignore]
+#[tokio::test]
+async fn head_works() {
+    let result = HTTPSHandler::head("https://github.com/XAMPPRocky/tokei/releases/download/v12.0.4/tokei-x86_64-unknown-linux-gnu.tar.gz").await;
+    assert!(result.is_ok());
+}

crates/llama-cpp-server/src/lib.rs

@@ -10,7 +10,7 @@ use serde::Deserialize;
 use supervisor::LlamaCppSupervisor;
 use tabby_common::{
     config::{HttpModelConfigBuilder, LocalModelConfig, ModelConfig},
-    registry::{parse_model_id, ModelRegistry, GGML_MODEL_RELATIVE_PATH},
+    registry::{parse_model_id, ModelRegistry},
 };
 use tabby_inference::{ChatCompletionStream, CompletionOptions, CompletionStream, Embedding};
 
@@ -277,14 +277,9 @@ pub async fn create_embedding(config: &ModelConfig) -> Arc<dyn Embedding> {
 }
 
 async fn resolve_model_path(model_id: &str) -> String {
-    let path = PathBuf::from(model_id);
-    let path = if path.exists() {
-        path.join(GGML_MODEL_RELATIVE_PATH.as_str())
-    } else {
-        let (registry, name) = parse_model_id(model_id);
-        let registry = ModelRegistry::new(registry).await;
-        registry.get_model_path(name)
-    };
+    let (registry, name) = parse_model_id(model_id);
+    let registry = ModelRegistry::new(registry).await;
+    let path = registry.get_model_entry_path(name);
     path.display().to_string()
 }
 

crates/tabby-common/src/env.rs

@@ -0,0 +1,12 @@
+pub fn get_download_host() -> String {
+    std::env::var("TABBY_DOWNLOAD_HOST").unwrap_or_else(|_| "huggingface.co".to_string())
+}
+
+pub fn get_huggingface_mirror_host() -> Option<String> {
+    std::env::var("TABBY_HUGGINGFACE_HOST_OVERRIDE").ok()
+}
+
+// for debug only
+pub fn use_local_model_json() -> bool {
+    std::env::var("TABBY_USE_LOCAL_MODEL_JSON").is_ok()
+}

crates/tabby-common/src/lib.rs

@@ -4,6 +4,7 @@ pub mod api;
 pub mod axum;
 pub mod config;
 pub mod constants;
+pub mod env;
 pub mod index;
 pub mod languages;
 pub mod path;

crates/tabby-common/src/registry.rs

@@ -1,10 +1,14 @@
 use std::{fs, path::PathBuf};
 
 use anyhow::{Context, Result};
-use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 
-use crate::path::models_dir;
+use crate::{env::use_local_model_json, path::models_dir};
+
+// default_entrypoint is legacy entrypoint for single model file
+fn default_entrypoint() -> String {
+    "model.gguf".to_string()
+}
 
 #[derive(Serialize, Deserialize)]
 pub struct ModelInfo {
@@ -16,6 +20,10 @@ pub struct ModelInfo {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub urls: Option<Vec<String>>,
     pub sha256: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub urls_sha256: Option<Vec<String>>,
+    #[serde(default = "default_entrypoint")]
+    pub entrypoint: String,
 }
 
 fn models_json_file(registry: &str) -> PathBuf {
@@ -54,30 +62,71 @@ pub struct ModelRegistry {
     pub models: Vec<ModelInfo>,
 }
 
+// model registry tree structure
+
+// root: ~/.tabby/models/TABBYML
+
+// fn get_model_root_dir(model_name) -> {root}/{model_name}
+
+// fn get_model_dir(model_name) -> {root}/{model_name}/ggml
+
+// fn get_model_path(model_name)
+// for single model file
+// -> {root}/{model_name}/ggml/model.gguf
+// for multiple model files
+// -> {root}/{model_name}/ggml/{entrypoint}
+
 impl ModelRegistry {
     pub async fn new(registry: &str) -> Self {
-        Self {
-            name: registry.to_owned(),
-            models: load_remote_registry(registry).await.unwrap_or_else(|err| {
-                load_local_registry(registry).unwrap_or_else(|_| {
-                    panic!(
-                        "Failed to fetch model organization <{}>: {:?}",
-                        registry, err
-                    )
-                })
-            }),
+        if use_local_model_json() {
+            Self {
+                name: registry.to_owned(),
+                models: load_local_registry(registry).unwrap_or_else(|_| {
+                    panic!("Failed to fetch model organization <{}>", registry)
+                }),
+            }
+        } else {
+            Self {
+                name: registry.to_owned(),
+                models: load_remote_registry(registry).await.unwrap_or_else(|err| {
+                    load_local_registry(registry).unwrap_or_else(|_| {
+                        panic!(
+                            "Failed to fetch model organization <{}>: {:?}",
+                            registry, err
+                        )
+                    })
+                }),
+            }
         }
     }
 
-    fn get_model_dir(&self, name: &str) -> PathBuf {
+    // get_model_store_dir returns {root}/{name}/ggml, e.g.. ~/.tabby/models/TABBYML/StarCoder-1B/ggml
+    pub fn get_model_store_dir(&self, name: &str) -> PathBuf {
+        models_dir().join(&self.name).join(name).join("ggml")
+    }
+
+    // get_model_dir returns {root}/{name}, e.g. ~/.tabby/models/TABBYML/StarCoder-1B
+    pub fn get_model_dir(&self, name: &str) -> PathBuf {
         models_dir().join(&self.name).join(name)
     }
 
+    // get_legacy_model_path returns {root}/{name}/q8_0.v2.gguf, e.g. ~/.tabby/models/TABBYML/StarCoder-1B/q8_0.v2.gguf
+    fn get_legacy_model_path(&self, name: &str) -> PathBuf {
+        self.get_model_store_dir(name).join("q8_0.v2.gguf")
+    }
+
+    // get_model_path returns the entrypoint of the model,
+    // for single model file, it returns {root}/{name}/ggml/model.gguf
+    // for multiple model files, it returns {root}/{name}/ggml/{entrypoint}
+    pub fn get_model_entry_path(&self, name: &str) -> PathBuf {
+        let model_info = self.get_model_info(name);
+        self.get_model_store_dir(name)
+            .join(model_info.entrypoint.clone())
+    }
+
     pub fn migrate_model_path(&self, name: &str) -> Result<(), std::io::Error> {
-        let model_path = self.get_model_path(name);
-        let old_model_path = self
-            .get_model_dir(name)
-            .join(LEGACY_GGML_MODEL_RELATIVE_PATH.as_str());
+        let model_path = self.get_model_entry_path(name);
+        let old_model_path = self.get_legacy_model_path(name);
 
         if !model_path.exists() && old_model_path.exists() {
             std::fs::rename(&old_model_path, &model_path)?;
@@ -89,11 +138,6 @@ impl ModelRegistry {
         Ok(())
     }
 
-    pub fn get_model_path(&self, name: &str) -> PathBuf {
-        self.get_model_dir(name)
-            .join(GGML_MODEL_RELATIVE_PATH.as_str())
-    }
-
     pub fn save_model_info(&self, name: &str) {
         let model_info = self.get_model_info(name);
         let path = self.get_model_dir(name).join("tabby.json");
@@ -120,18 +164,11 @@ pub fn parse_model_id(model_id: &str) -> (&str, &str) {
     }
 }
 
-lazy_static! {
-    pub static ref LEGACY_GGML_MODEL_RELATIVE_PATH: String =
-        format!("ggml{}q8_0.v2.gguf", std::path::MAIN_SEPARATOR_STR);
-    pub static ref GGML_MODEL_RELATIVE_PATH: String =
-        format!("ggml{}model.gguf", std::path::MAIN_SEPARATOR_STR);
-}
-
 #[cfg(test)]
 mod tests {
     use temp_testdir::TempDir;
 
-    use super::{ModelRegistry, *};
+    use super::ModelRegistry;
     use crate::path::set_tabby_root;
 
     #[tokio::test]
@@ -140,9 +177,9 @@ mod tests {
         set_tabby_root(root.to_path_buf());
 
         let registry = ModelRegistry::new("TabbyML").await;
-        let dir = registry.get_model_dir("StarCoder-1B");
+        let name = "StarCoder-1B";
 
-        let old_model_path = dir.join(LEGACY_GGML_MODEL_RELATIVE_PATH.as_str());
+        let old_model_path = registry.get_legacy_model_path(name);
         tokio::fs::create_dir_all(old_model_path.parent().unwrap())
             .await
             .unwrap();
@@ -154,7 +191,7 @@ mod tests {
             .unwrap();
 
         registry.migrate_model_path("StarCoder-1B").unwrap();
-        assert!(registry.get_model_path("StarCoder-1B").exists());
+        assert!(registry.get_model_entry_path("StarCoder-1B").exists());
         assert!(old_model_path.exists());
     }
 }

crates/tabby-download/Cargo.toml

@@ -9,3 +9,7 @@ tabby-common = { path = "../tabby-common" }
 anyhow = { workspace = true }
 tracing = { workspace = true }
 tokio-retry = "0.3.0"
+futures.workspace = true
+reqwest = { workspace = true }
+tokio = {workspace=true}
+regex = {workspace=true}