Lumotia/src-tauri/src/commands/llm.rs

use tauri::{Emitter, State};

use crate::commands::power::PowerAssertion;
use crate::AppState;
use kon_ai_formatting::llm_cleanup_text;
use kon_core::hardware;
use kon_llm::model_manager::{self, model_info};
use kon_llm::LlmModelId;

#[derive(Debug, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LlmModelStatusDto {
    pub id: String,
    pub display_name: String,
    pub downloaded: bool,
    pub loaded: bool,
    pub size_bytes: u64,
    pub description: String,
    pub minimum_ram_bytes: u64,
    pub recommended_vram_bytes: Option<u64>,
}

fn parse_model_id(model_id: String) -> Result<LlmModelId, String> {
    model_id.parse()
}

#[tauri::command]
pub fn recommend_llm_tier() -> Result<String, String> {
    let profile = hardware::probe_system();
    let ram_bytes = profile.ram.0.saturating_mul(1024 * 1024);
    let vram_bytes = profile
        .gpu
        .map(|gpu| gpu.vram.0.saturating_mul(1024 * 1024));
    Ok(model_manager::recommend_tier(ram_bytes, vram_bytes)
        .as_str()
        .to_string())
}

#[tauri::command]
pub fn check_llm_model(
    state: State<'_, AppState>,
    model_id: String,
) -> Result<LlmModelStatusDto, String> {
    let id = parse_model_id(model_id)?;
    let info = model_info(id);
    let loaded_model_id = state.llm_engine.loaded_model_id();

    Ok(LlmModelStatusDto {
        id: info.id,
        display_name: info.display_name.to_string(),
        downloaded: model_manager::is_downloaded(id),
        loaded: loaded_model_id.as_deref() == Some(id.as_str()),
        size_bytes: info.size_bytes,
        description: info.description.to_string(),
        minimum_ram_bytes: info.minimum_ram_bytes,
        recommended_vram_bytes: info.recommended_vram_bytes,
    })
}

#[tauri::command]
pub async fn download_llm_model(app: tauri::AppHandle, model_id: String) -> Result<(), String> {
    let id = parse_model_id(model_id)?;
    let app_clone = app.clone();
    model_manager::download_model(id, move |done, total| {
        let percent = if total > 0 {
            ((done as f64 / total as f64) * 100.0).round() as u8
        } else {
            0
        };
        let _ = app_clone.emit(
            "kon:llm-download-progress",
            serde_json::json!({
                "modelId": id.as_str(),
                "done": done,
                "total": total,
                "percent": percent,
            }),
        );
    })
    .await
    .map_err(|e| e.to_string())
}

#[tauri::command]
pub async fn load_llm_model(
    state: State<'_, AppState>,
    model_id: String,
    use_gpu: Option<bool>,
) -> Result<(), String> {
    let id = parse_model_id(model_id)?;
    let path = model_manager::model_path(id);
    if !path.exists() {
        return Err("Model not downloaded — call download_llm_model first".to_string());
    }

    let engine = state.llm_engine.clone();
    let use_gpu = use_gpu.unwrap_or(true);
    tokio::task::spawn_blocking(move || engine.load_model(id, &path, use_gpu))
        .await
        .map_err(|e| e.to_string())?
        .map_err(|e| e.to_string())
}

#[tauri::command]
pub fn unload_llm_model(state: State<'_, AppState>) -> Result<(), String> {
    state.llm_engine.unload().map_err(|e| e.to_string())
}

#[tauri::command]
pub fn delete_llm_model(state: State<'_, AppState>, model_id: String) -> Result<(), String> {
    let id = parse_model_id(model_id)?;
    if state.llm_engine.loaded_model_id().as_deref() == Some(id.as_str()) {
        state.llm_engine.unload().map_err(|e| e.to_string())?;
    }
    model_manager::delete_model(id).map_err(|e| e.to_string())
}

#[tauri::command]
pub fn get_llm_status(state: State<'_, AppState>) -> Result<bool, String> {
    Ok(state.llm_engine.is_loaded())
}

#[tauri::command]
pub async fn cleanup_transcript_text_cmd(
    state: State<'_, AppState>,
    transcript: String,
    profile_id: Option<String>,
) -> Result<String, String> {
    let resolved_profile_id =
        profile_id.unwrap_or_else(|| kon_storage::DEFAULT_PROFILE_ID.to_string());
    let profile_terms: Vec<String> =
        kon_storage::database::list_profile_terms(&state.db, &resolved_profile_id)
            .await
            .map_err(|e| e.to_string())?
            .into_iter()
            .map(|term| term.term)
            .collect();

    let engine = state.llm_engine.clone();
    tokio::task::spawn_blocking(move || {
        // macOS: pin a power assertion for the duration of the LLM
        // generation so App Nap can't decide to throttle us mid-token.
        // No-op on every other OS. Item #9.
        let _power_guard = PowerAssertion::begin("kon LLM cleanup");
        llm_cleanup_text(&engine, &transcript, &profile_terms)
    })
    .await
    .map_err(|e| e.to_string())?
    .map_err(|e| e.to_string())
}