znote 1.2.0

A lightweight desktop editor for Typst with live preview, built with Tauri
use std::io::Read;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};

use tauri::ipc::Channel;
use tauri::{AppHandle, Manager};

use zsuite_finder::models::{Model, Piece, Progress};

use crate::knowledge;

static STOP_NAMES: AtomicBool = AtomicBool::new(false);
static STOP_SEARCH: AtomicBool = AtomicBool::new(false);

fn stop(id: &str) -> &'static AtomicBool {
    if id == "names" {
        &STOP_NAMES
    } else {
        &STOP_SEARCH
    }
}

pub const NAME_FINDER: Model = Model {
    id: "names",
    folder: "name-finder",
    repository: "knowledgator/gliner-pii-edge-v1.0",
    revision: "9b7f39b0a2da971a5beea78d35f1539d4009c891",
    pieces: &[
        Piece {
            remote: "onnx/model_quint8.onnx",
            local: "model.onnx",
            sha256: "988acb03456b26e2d9f2521016d820310c2ed64deb4a846297d3289f0c2eb7e4",
            size: 45_820_894,
        },
        Piece {
            remote: "tokenizer.json",
            local: "tokenizer.json",
            sha256: "84b3a9b18f04a0ccd03b72d9f871b7e0bec40fd7021ef50bc30a7c3693c11205",
            size: 3_583_593,
        },
    ],
};

pub const SMART_SEARCH: Model = Model {
    id: "search",
    folder: "smart-search",
    repository: "minishlab/potion-base-8M",
    revision: "bf8b056651a2c21b8d2565580b8569da283cab23",
    pieces: &[
        Piece {
            remote: "model.safetensors",
            local: "model.safetensors",
            sha256: "f65d0f325faadc1e121c319e2faa41170d3fa07d8c89abd48ca5358d9a223de2",
            size: 30_236_760,
        },
        Piece {
            remote: "tokenizer.json",
            local: "tokenizer.json",
            sha256: "e67e803f624fb4d67dea1c730d06e1067e1b14d830e2c2202569e3ef0f70bb50",
            size: 683_666,
        },
        Piece {
            remote: "config.json",
            local: "config.json",
            sha256: "2a6ac0e9aaa356a68a5688070db78fc3a464fefe85d2f06a1905ce3718687553",
            size: 202,
        },
    ],
};

pub fn folder(app: &AppHandle) -> Result<PathBuf, String> {
    Ok(app
        .path()
        .app_local_data_dir()
        .map_err(|error| error.to_string())?
        .join("models"))
}

fn web(address: &str) -> Result<Box<dyn Read>, String> {
    let response = zsuite_agent::connection()?
        .build()
        .get(address)
        .call()
        .map_err(zsuite_agent::refused)?;
    Ok(Box::new(response.into_reader()))
}

#[tauri::command]
pub async fn models_download(
    app: AppHandle,
    model: String,
    progress: Channel<Progress>,
) -> Result<(), String> {
    crate::blocking(move || models_download_now(app, model, progress)).await
}

fn models_download_now(
    app: AppHandle,
    model: String,
    progress: Channel<Progress>,
) -> Result<(), String> {
    let chosen = match model.as_str() {
        "names" => &NAME_FINDER,
        "search" => &SMART_SEARCH,
        _ => return Err(format!("“{model}” is not a model znote knows.")),
    };
    stop(chosen.id).store(false, Ordering::Relaxed);
    chosen.fetch(&folder(&app)?, &web, stop(chosen.id), &|step| {
        let _ = progress.send(step);
    })?;
    knowledge::models_ready(&app, chosen.id);
    Ok(())
}

#[tauri::command]
pub fn models_cancel(model: String) {
    if matches!(model.as_str(), "names" | "search") {
        stop(&model).store(true, Ordering::Relaxed);
    }
}