Initial ForgeLM import

tauri-app/src-tauri/src/model_files.rs

@@ -0,0 +1,188 @@
+use std::path::Path;
+
+use tauri::AppHandle;
+use tauri_plugin_dialog::DialogExt;
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq)]
+pub struct ModelFile {
+    pub name: String,
+    pub path: String,
+}
+
+#[derive(serde::Serialize)]
+pub struct PickDirResult {
+    pub dir: String,
+    pub models: Vec<ModelFile>,
+}
+
+fn model_file_from_path(path: &Path) -> ModelFile {
+    ModelFile {
+        name: path
+            .file_name()
+            .unwrap_or_default()
+            .to_string_lossy()
+            .to_string(),
+        path: path.to_string_lossy().to_string(),
+    }
+}
+
+fn scan_model_dir(path: &Path) -> Result<Vec<ModelFile>, String> {
+    if !path.exists() {
+        return Err(format!("Directory not found: {}", path.display()));
+    }
+    if !path.is_dir() {
+        return Err(format!("Not a directory: {}", path.display()));
+    }
+
+    let entries = std::fs::read_dir(path).map_err(|e| e.to_string())?;
+    let mut models = Vec::new();
+
+    for entry in entries.flatten() {
+        let path = entry.path();
+        if path.is_file()
+            && path
+            .extension()
+            .map(|ext| ext.eq_ignore_ascii_case("gguf"))
+            .unwrap_or(false)
+        {
+            models.push(model_file_from_path(&path));
+        }
+    }
+
+    models.sort_by(|a, b| a.name.cmp(&b.name));
+    Ok(models)
+}
+
+fn import_model_file(src: &Path, dest_dir: &Path) -> Result<ModelFile, String> {
+    if !src.exists() {
+        return Err(format!("Source file not found: {}", src.display()));
+    }
+    if !src.is_file() {
+        return Err("Source is not a file".into());
+    }
+    if !src
+        .extension()
+        .unwrap_or_default()
+        .eq_ignore_ascii_case("gguf")
+    {
+        return Err("Only .gguf files are supported".into());
+    }
+
+    if !dest_dir.exists() {
+        std::fs::create_dir_all(dest_dir)
+            .map_err(|e| format!("Failed to create models directory: {e}"))?;
+    }
+
+    let file_name = src.file_name().unwrap_or_default();
+    let dest_path = dest_dir.join(file_name);
+
+    if src == dest_path {
+        return Ok(model_file_from_path(&dest_path));
+    }
+
+    std::fs::copy(src, &dest_path).map_err(|e| format!("Failed to copy model: {e}"))?;
+    Ok(model_file_from_path(&dest_path))
+}
+
+#[tauri::command]
+pub fn scan_models(dir: String) -> Result<Vec<ModelFile>, String> {
+    scan_model_dir(Path::new(&dir))
+}
+
+#[tauri::command]
+pub async fn pick_models_dir(app: AppHandle) -> Result<PickDirResult, String> {
+    let folder = app
+        .dialog()
+        .file()
+        .set_title("Select models folder")
+        .blocking_pick_folder();
+
+    match folder {
+        Some(path) => {
+            let dir = path.to_string();
+            let models = scan_model_dir(Path::new(&dir))?;
+            Ok(PickDirResult { dir, models })
+        }
+        None => Err("cancelled".into()),
+    }
+}
+
+#[tauri::command]
+pub async fn pick_model_files(app: AppHandle) -> Result<Vec<ModelFile>, String> {
+    let files = app
+        .dialog()
+        .file()
+        .set_title("Select model files")
+        .add_filter("GGUF models", &["gguf"])
+        .blocking_pick_files();
+
+    match files {
+        Some(paths) => Ok(paths.iter().map(|path| model_file_from_path(Path::new(&path.to_string()))).collect()),
+        None => Err("cancelled".into()),
+    }
+}
+
+#[tauri::command]
+pub fn import_model(src_path: String, dest_dir: String) -> Result<ModelFile, String> {
+    import_model_file(Path::new(&src_path), Path::new(&dest_dir))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{import_model_file, scan_model_dir};
+    use std::path::PathBuf;
+    use std::time::{SystemTime, UNIX_EPOCH};
+
+    fn temp_path(name: &str) -> PathBuf {
+        let unique = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .expect("system time")
+            .as_nanos();
+        std::env::temp_dir().join(format!("forgelm-{name}-{unique}"))
+    }
+
+    #[test]
+    fn scan_model_dir_filters_and_sorts_gguf_files() {
+        let dir = temp_path("scan-models");
+        std::fs::create_dir_all(&dir).expect("create temp dir");
+        std::fs::write(dir.join("zeta.gguf"), b"demo").expect("write zeta");
+        std::fs::write(dir.join("alpha.gguf"), b"demo").expect("write alpha");
+        std::fs::write(dir.join("notes.txt"), b"ignore").expect("write non-model file");
+
+        let models = scan_model_dir(&dir).expect("scan model dir");
+        let names: Vec<_> = models.into_iter().map(|model| model.name).collect();
+
+        assert_eq!(names, vec!["alpha.gguf", "zeta.gguf"]);
+
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn scan_model_dir_ignores_directories_named_like_models() {
+        let dir = temp_path("scan-model-dirs");
+        std::fs::create_dir_all(dir.join("folder.gguf")).expect("create fake model dir");
+        std::fs::write(dir.join("real.gguf"), b"demo").expect("write model file");
+
+        let models = scan_model_dir(&dir).expect("scan model dir");
+        let names: Vec<_> = models.into_iter().map(|model| model.name).collect();
+
+        assert_eq!(names, vec!["real.gguf"]);
+
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+
+    #[test]
+    fn import_model_file_returns_existing_model_when_source_matches_destination() {
+        let dir = temp_path("import-model");
+        std::fs::create_dir_all(&dir).expect("create temp dir");
+        let model_path = dir.join("demo.gguf");
+        std::fs::write(&model_path, b"demo").expect("write demo model");
+
+        let model = import_model_file(&model_path, &dir).expect("import model file");
+
+        assert_eq!(model.name, "demo.gguf");
+        assert_eq!(PathBuf::from(model.path), model_path);
+
+        let _ = std::fs::remove_dir_all(&dir);
+    }
+}