refactor: extract registry to its sub dir.

crates/tabby-download/Cargo.toml

@@ -16,3 +16,4 @@ tokio-retry = "0.3.0"
 urlencoding = "2.1.3"
 serde_json = { workspace = true }
 cached = { version = "0.46.0", features = ["async", "proc_macro"] }
+async-trait = { workspace = true }

crates/tabby-download/src/lib.rs

@@ -7,7 +7,7 @@ use anyhow::{anyhow, Result};
 use cache_info::CacheInfo;
 use futures_util::StreamExt;
 use indicatif::{ProgressBar, ProgressStyle};
-use registry::HuggingFaceRegistry;
+use registry::{create_registry, Registry};
 use tabby_common::path::ModelDir;
 use tokio_retry::{
     strategy::{jitter, ExponentialBackoff},
@@ -17,15 +17,15 @@ use tokio_retry::{
 pub struct Downloader {
     model_id: String,
     prefer_local_file: bool,
-    registry: HuggingFaceRegistry,
+    registry: Box<dyn Registry>,
 }
 
 impl Downloader {
     pub fn new(model_id: &str, prefer_local_file: bool) -> Self {
         Self {
             model_id: model_id.to_owned(),
             prefer_local_file,
-            registry: HuggingFaceRegistry::default(),
+            registry: create_registry(),
         }
     }
 
@@ -62,7 +62,7 @@ impl Downloader {
         let mut cache_info = CacheInfo::from(&self.model_id).await;
         for (path, required) in files {
             download_model_file(
-                &self.registry,
+                self.registry.as_ref(),
                 &mut cache_info,
                 &self.model_id,
                 path,
@@ -76,7 +76,7 @@ impl Downloader {
 }
 
 async fn download_model_file(
-    registry: &HuggingFaceRegistry,
+    registry: &dyn Registry,
     cache_info: &mut CacheInfo,
     model_id: &str,
     path: &str,
@@ -122,7 +122,7 @@ async fn download_model_file(
 }
 
 async fn download_file(
-    registry: &HuggingFaceRegistry,
+    registry: &dyn Registry,
     url: &str,
     path: &str,
     local_cache_key: Option<&str>,

crates/tabby-download/src/registry/huggingface.rs

@@ -0,0 +1,24 @@
+use anyhow::{anyhow, Result};
+use async_trait::async_trait;
+
+use crate::Registry;
+
+#[derive(Default)]
+pub struct HuggingFaceRegistry {}
+
+#[async_trait]
+impl Registry for HuggingFaceRegistry {
+    fn build_url(&self, model_id: &str, path: &str) -> String {
+        format!("https://huggingface.co/{}/resolve/main/{}", model_id, path)
+    }
+
+    async fn build_cache_key(&self, url: &str) -> Result<String> {
+        let res = reqwest::get(url).await?;
+        let cache_key = res
+            .headers()
+            .get("etag")
+            .ok_or(anyhow!("etag key missing"))?
+            .to_str()?;
+        Ok(cache_key.to_owned())
+    }
+}

crates/tabby-download/src/registry/mod.rs

@@ -0,0 +1,25 @@
+mod huggingface;
+mod modelscope;
+
+use anyhow::Result;
+use async_trait::async_trait;
+use huggingface::HuggingFaceRegistry;
+
+use self::modelscope::ModelScopeRegistry;
+
+#[async_trait]
+pub trait Registry {
+    fn build_url(&self, model_id: &str, path: &str) -> String;
+    async fn build_cache_key(&self, url: &str) -> Result<String>;
+}
+
+pub fn create_registry() -> Box<dyn Registry> {
+    let registry = std::env::var("TABBY_REGISTRY").unwrap_or("huggingface".to_owned());
+    if registry == "huggingface" {
+        Box::<HuggingFaceRegistry>::default()
+    } else if registry == "modelscope" {
+        Box::<ModelScopeRegistry>::default()
+    } else {
+        panic!("Unsupported registry {}", registry);
+    }
+}

crates/tabby-download/src/registry.rs → crates/tabby-download/src/registry/modelscope.rs

@@ -1,45 +1,30 @@
 use std::collections::HashMap;
 
 use anyhow::{anyhow, Result};
+use async_trait::async_trait;
 use cached::proc_macro::cached;
 use reqwest::Url;
 use serde::Deserialize;
 
-#[derive(Default)]
-pub struct HuggingFaceRegistry {}
-
-impl HuggingFaceRegistry {
-    pub fn build_url(&self, model_id: &str, path: &str) -> String {
-        format!("https://huggingface.co/{}/resolve/main/{}", model_id, path)
-    }
-
-    pub async fn build_cache_key(&self, url: &str) -> Result<String> {
-        let res = reqwest::get(url).await?;
-        let cache_key = res
-            .headers()
-            .get("etag")
-            .ok_or(anyhow!("etag key missing"))?
-            .to_str()?;
-        Ok(cache_key.to_owned())
-    }
-}
+use crate::Registry;
 
 #[derive(Default)]
-pub struct ModelScopeRegistry {
-}
+pub struct ModelScopeRegistry {}
 
-impl ModelScopeRegistry {
-    pub fn build_url(&self, model_id: &str, path: &str) -> String {
+#[async_trait]
+impl Registry for ModelScopeRegistry {
+    fn build_url(&self, model_id: &str, path: &str) -> String {
         format!(
             "https://modelscope.cn/api/v1/models/{}/repo?FilePath={}",
             model_id,
             urlencoding::encode(path)
         )
     }
 
-    pub async fn build_cache_key(&self, url: &str) -> Result<String> {
+    async fn build_cache_key(&self, url: &str) -> Result<String> {
         let url = Url::parse(url)?;
-        let model_id = url.path()
+        let model_id = url
+            .path()
             .strip_prefix("/api/v1/models/")
             .ok_or(anyhow!("Invalid url"))?
             .strip_suffix("/repo")
@@ -49,13 +34,10 @@ impl ModelScopeRegistry {
         let path = query
             .get("FilePath")
             .ok_or(anyhow!("Failed to extract FilePath"))?;
-        self.get_cache_key(model_id, path).await
-    }
 
-    async fn get_cache_key(&self, model_id: &str, path: &str) -> Result<String> {
         let revision_map = fetch_revision_map(model_id.to_owned()).await?;
         for x in revision_map.data.files {
-            if x.path == path {
+            if x.path == *path {
                 return Ok(x.sha256);
             }
         }