Use .puma as the root folder

Cargo.toml

@@ -16,3 +16,4 @@ serde_derive = "1.0"
 env_logger = "0.11.6"
 log = "0.4.26"
 indicatif = "0.17.11"
+dirs = "6.0.0"

src/cli/cmds.rs

@@ -2,6 +2,7 @@ use clap::{Parser, Subcommand};
 use prettytable::{format, row, Table};
 
 use crate::downloader::ollama::OllamaDownloader;
+use crate::util::file;
 
 #[derive(Parser)]
 #[command(name = "PUMA")]
@@ -105,9 +106,9 @@ pub async fn run(cli: Cli) {
             }
             Provider::Ollama => {
                 let d = OllamaDownloader::new(&args.model);
-                d.download_model("/Users/kerthcet/Workspaces/InftyAI/puma/tmp")
-                    .await
-                    .unwrap();
+                let model_path = file::root_home().join(file::model_folder_name(&args.model));
+                file::create_folder_if_not_exists(&model_path).unwrap();
+                d.download_model(&model_path).await.unwrap();
             }
             Provider::Modelscope => {
                 println!("Downloading model from Modelscope...");

src/downloader/ollama.rs

@@ -1,3 +1,4 @@
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use indicatif::{MultiProgress, ProgressStyle};
@@ -50,7 +51,7 @@ impl OllamaDownloader {
         }
     }
 
-    pub async fn download_model(&self, path: &str) -> Result<(), DownloadError> {
+    pub async fn download_model(&self, path: &PathBuf) -> Result<(), DownloadError> {
         info!(
             "Downloading model {} from ollama provider...",
             self.model_name
@@ -89,7 +90,7 @@ impl OllamaDownloader {
             let semaphore: Arc<Semaphore> = Arc::clone(&semaphore);
             let arc_m = Arc::clone(&arc_m);
 
-            let full_path = format!("{}/{}", path, layer.path());
+            let full_path = path.join(layer.path());
             let size = layer.size;
             let sty = sty.clone();
 
@@ -98,19 +99,18 @@ impl OllamaDownloader {
                 // TODO: return the error.
                 let _ = request::download_file(
                     client,
-                    layer_url.clone(),
+                    layer_url,
                     size,
                     layer.path().to_string(),
-                    full_path,
+                    &full_path,
                     arc_m,
                     sty.clone(),
                 )
                 .await
                 .map_err(|e| {
                     error!(
-                        "Failed to download file {} from {}: {}",
+                        "Failed to download file {}: {}",
                         layer.path(),
-                        layer_url,
                         e.to_string()
                     );
                 });

src/main.rs

@@ -7,10 +7,14 @@ use env_logger;
 use tokio::runtime::Builder;
 
 use crate::cli::cmds::{run, Cli};
+use crate::util::file;
 
 fn main() {
     env_logger::init();
 
+    // create the root folder.
+    file::create_folder_if_not_exists(&file::root_home()).unwrap();
+
     let runtime = Builder::new_multi_thread()
         .worker_threads(4)
         .enable_all()

src/util/file.rs

@@ -0,0 +1,18 @@
+use std::fs;
+use std::path::PathBuf;
+
+use dirs::home_dir;
+
+pub fn create_folder_if_not_exists(folder_path: &PathBuf) -> std::io::Result<()> {
+    fs::create_dir_all(folder_path)?;
+    Ok(())
+}
+
+pub fn model_folder_name(model_name: &str) -> String {
+    model_name.replace(":", "--").replace("/", "--")
+}
+
+pub fn root_home() -> PathBuf {
+    let home = home_dir().expect("Failed to get home directory");
+    home.join(".puma")
+}

src/util/mod.rs

@@ -1 +1,2 @@
+pub mod file;
 pub mod request;

src/util/request.rs

@@ -3,6 +3,7 @@ use std::error::Error;
 use std::fs::File;
 use std::io;
 use std::os::unix::fs::FileExt;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
@@ -20,11 +21,15 @@ pub async fn download_file(
     url: String,
     content_length: u64,
     filename: String,
-    output_path: String,
+    output_path: &PathBuf,
     m: Arc<MultiProgress>,
     sty: ProgressStyle,
 ) -> Result<(), Box<dyn Error>> {
-    debug!("Start to download file {} to {}", filename, output_path);
+    debug!(
+        "Start to download file {} to {}",
+        filename,
+        output_path.display()
+    );
 
     let mut tasks = Vec::new();
     let mut start = 0;