Lumotia/crates/llm/src/lib.rs

use std::num::NonZeroU32;
use std::path::Path;
use std::sync::{Arc, Mutex};

use encoding_rs::UTF_8;
use llama_cpp_2::context::params::LlamaContextParams;
use llama_cpp_2::llama_backend::LlamaBackend;
use llama_cpp_2::llama_batch::LlamaBatch;
use llama_cpp_2::model::params::LlamaModelParams;
use llama_cpp_2::model::{AddBos, LlamaChatMessage, LlamaChatTemplate, LlamaModel};
use llama_cpp_2::sampling::LlamaSampler;
use serde::{Deserialize, Serialize};

pub mod grammars;
pub mod model_manager;
pub mod prompts;

pub use model_manager::{recommend_tier, LlmModelId, LlmModelInfo};

const DEFAULT_CONTEXT_TOKENS: u32 = 4096;
const MAX_CONTEXT_TOKENS: u32 = 8192;
const CONTEXT_RESERVE_TOKENS: u32 = 64;
const GENERATION_SEED: u32 = 0;

#[derive(Debug, thiserror::Error)]
pub enum EngineError {
    #[error("LLM not loaded. Download an AI model in Settings.")]
    NotLoaded,
    #[error("LLM load failed: {0}")]
    LoadFailed(String),
    #[error(
        "prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens"
    )]
    PromptTooLong {
        prompt_tokens: usize,
        max_tokens: u32,
        available_prompt_tokens: u32,
        context_window: u32,
    },
    #[error("inference failed: {0}")]
    Inference(String),
    #[error("model output not valid JSON: {0}")]
    InvalidJson(String),
}

#[derive(Debug, Clone)]
pub struct GenerationConfig {
    pub max_tokens: u32,
    pub temperature: f32,
    pub stop_sequences: Vec<String>,
    pub grammar: Option<String>,
}

impl Default for GenerationConfig {
    fn default() -> Self {
        Self {
            max_tokens: 1024,
            temperature: 0.0,
            stop_sequences: Vec::new(),
            grammar: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LoadedModelState {
    pub model_id: String,
    pub model_path: String,
    pub use_gpu: bool,
}

#[derive(Default)]
struct LlmState {
    backend: Option<Arc<LlamaBackend>>,
    model: Option<Arc<LlamaModel>>,
    loaded: Option<LoadedModelState>,
}

#[derive(Clone, Default)]
pub struct LlmEngine {
    inner: Arc<Mutex<LlmState>>,
}

impl LlmEngine {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn load(&self, model_path: &Path) -> Result<(), EngineError> {
        self.load_model(LlmModelId::default_tier(), model_path, true)
    }

    pub fn load_model(
        &self,
        model_id: LlmModelId,
        model_path: &Path,
        use_gpu: bool,
    ) -> Result<(), EngineError> {
        let mut guard = self.inner.lock().unwrap();

        if let Some(loaded) = &guard.loaded {
            if loaded.model_id == model_id.as_str()
                && loaded.model_path == model_path.display().to_string()
                && loaded.use_gpu == use_gpu
            {
                return Ok(());
            }
        }

        let backend = match guard.backend.clone() {
            Some(existing) => existing,
            None => Arc::new(
                LlamaBackend::init()
                    .map_err(|e| EngineError::LoadFailed(format!("backend init: {e}")))?,
            ),
        };

        let gpu_layers = if use_gpu { u32::MAX } else { 0 };
        let params = LlamaModelParams::default().with_n_gpu_layers(gpu_layers);
        let model = LlamaModel::load_from_file(&backend, model_path, &params)
            .map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))?;

        guard.backend = Some(backend);
        guard.model = Some(Arc::new(model));
        guard.loaded = Some(LoadedModelState {
            model_id: model_id.as_str().to_string(),
            model_path: model_path.display().to_string(),
            use_gpu,
        });
        Ok(())
    }

    pub fn unload(&self) -> Result<(), EngineError> {
        let mut guard = self.inner.lock().unwrap();
        guard.model = None;
        guard.backend = None;
        guard.loaded = None;
        Ok(())
    }

    pub fn is_loaded(&self) -> bool {
        self.inner.lock().unwrap().model.is_some()
    }

    pub fn loaded_model(&self) -> Option<LoadedModelState> {
        self.inner.lock().unwrap().loaded.clone()
    }

    pub fn loaded_model_id(&self) -> Option<String> {
        self.loaded_model().map(|loaded| loaded.model_id)
    }

    pub fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String, EngineError> {
        let (backend, model) = self.loaded_handles()?;
        let prompt_tokens = model
            .str_to_token(prompt, AddBos::Never)
            .map_err(|e| EngineError::Inference(format!("tokenize: {e}")))?;
        if prompt_tokens.is_empty() {
            return Ok(String::new());
        }

        let n_ctx = preflight_context_window(prompt_tokens.len(), config.max_tokens)?;
        let thread_count = i32::try_from(num_cpus::get().max(1)).unwrap_or(4);
        let ctx_params = LlamaContextParams::default()
            .with_n_ctx(Some(
                NonZeroU32::new(n_ctx).expect("n_ctx must be non-zero"),
            ))
            .with_n_batch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
            .with_n_ubatch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
            .with_n_threads(thread_count)
            .with_n_threads_batch(thread_count);
        let mut ctx = model
            .new_context(&backend, ctx_params)
            .map_err(|e| EngineError::Inference(format!("context: {e}")))?;

        let mut batch = LlamaBatch::new(prompt_tokens.len().max(1), 1);
        for (index, token) in prompt_tokens.iter().enumerate() {
            batch
                .add(*token, index as i32, &[0], index + 1 == prompt_tokens.len())
                .map_err(|e| EngineError::Inference(format!("batch add: {e}")))?;
        }
        ctx.decode(&mut batch)
            .map_err(|e| EngineError::Inference(format!("prefill decode: {e}")))?;

        let mut sampler = self.build_sampler(&model, config)?;
        let mut decoder = UTF_8.new_decoder();
        let mut generated = String::new();
        let mut cursor = prompt_tokens.len() as i32;

        for _ in 0..config.max_tokens {
            let next = sampler.sample(&ctx, batch.n_tokens() - 1);
            if model.is_eog_token(next) || next == model.token_eos() {
                break;
            }

            let piece = model
                .token_to_piece(next, &mut decoder, true, None)
                .map_err(|e| EngineError::Inference(format!("detokenize: {e}")))?;
            generated.push_str(&piece);
            sampler.accept(next);

            if let Some(stop_index) = first_stop_index(&generated, &config.stop_sequences) {
                generated.truncate(stop_index);
                break;
            }

            batch.clear();
            batch
                .add(next, cursor, &[0], true)
                .map_err(|e| EngineError::Inference(format!("sample batch: {e}")))?;
            cursor += 1;
            ctx.decode(&mut batch)
                .map_err(|e| EngineError::Inference(format!("sample decode: {e}")))?;
        }

        Ok(generated.trim().to_string())
    }

    pub fn cleanup_text(
        &self,
        system_prompt: &str,
        transcript: &str,
    ) -> Result<String, EngineError> {
        if transcript.trim().is_empty() {
            return Ok(String::new());
        }
        let model = self.loaded_model_arc()?;
        let prompt =
            render_chat_prompt(&model, &[("system", system_prompt), ("user", transcript)])?;
        self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 1024,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: None,
            },
        )
    }

    pub fn decompose_task(&self, task_text: &str) -> Result<Vec<String>, EngineError> {
        let model = self.loaded_model_arc()?;
        let prompt = render_chat_prompt(
            &model,
            &[
                ("system", prompts::DECOMPOSE_TASK_SYSTEM),
                ("user", &format!("Task: {task_text}")),
            ],
        )?;
        let raw = self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 512,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: Some(grammars::TASK_ARRAY_GRAMMAR.to_string()),
            },
        )?;
        parse_string_array(&raw)
    }

    pub fn extract_tasks(&self, transcript: &str) -> Result<Vec<String>, EngineError> {
        if transcript.trim().is_empty() {
            return Ok(Vec::new());
        }

        let model = self.loaded_model_arc()?;
        let prompt = render_chat_prompt(
            &model,
            &[
                ("system", prompts::EXTRACT_TASKS_SYSTEM),
                ("user", &format!("Transcript:\n{transcript}")),
            ],
        )?;
        let raw = self.generate(
            &prompt,
            &GenerationConfig {
                max_tokens: 768,
                temperature: 0.0,
                stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
                grammar: Some(grammars::OPTIONAL_TASK_ARRAY_GRAMMAR.to_string()),
            },
        )?;
        parse_string_array(&raw)
    }

    fn loaded_handles(&self) -> Result<(Arc<LlamaBackend>, Arc<LlamaModel>), EngineError> {
        let guard = self.inner.lock().unwrap();
        let backend = guard.backend.clone().ok_or(EngineError::NotLoaded)?;
        let model = guard.model.clone().ok_or(EngineError::NotLoaded)?;
        Ok((backend, model))
    }

    fn loaded_model_arc(&self) -> Result<Arc<LlamaModel>, EngineError> {
        self.loaded_handles().map(|(_, model)| model)
    }

    fn build_sampler(
        &self,
        model: &LlamaModel,
        config: &GenerationConfig,
    ) -> Result<LlamaSampler, EngineError> {
        let mut samplers = Vec::new();

        if let Some(grammar) = &config.grammar {
            samplers.push(
                LlamaSampler::grammar(model, grammar, "root")
                    .map_err(|e| EngineError::Inference(format!("grammar: {e}")))?,
            );
        }

        if config.temperature <= f32::EPSILON {
            samplers.push(LlamaSampler::greedy());
        } else {
            samplers.push(LlamaSampler::temp(config.temperature));
            samplers.push(LlamaSampler::dist(GENERATION_SEED));
        }

        Ok(if samplers.len() == 1 {
            samplers.remove(0)
        } else {
            LlamaSampler::chain_simple(samplers)
        })
    }
}

fn context_window_size(prompt_tokens: usize, max_tokens: u32) -> u32 {
    let required = prompt_tokens
        .saturating_add(max_tokens as usize)
        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
    DEFAULT_CONTEXT_TOKENS.max(required.min(MAX_CONTEXT_TOKENS as usize) as u32)
}

fn preflight_context_window(prompt_tokens: usize, max_tokens: u32) -> Result<u32, EngineError> {
    let required = prompt_tokens
        .saturating_add(max_tokens as usize)
        .saturating_add(CONTEXT_RESERVE_TOKENS as usize);
    if required > MAX_CONTEXT_TOKENS as usize {
        let available_prompt_tokens =
            MAX_CONTEXT_TOKENS.saturating_sub(max_tokens.saturating_add(CONTEXT_RESERVE_TOKENS));
        return Err(EngineError::PromptTooLong {
            prompt_tokens,
            max_tokens,
            available_prompt_tokens,
            context_window: MAX_CONTEXT_TOKENS,
        });
    }

    Ok(context_window_size(prompt_tokens, max_tokens))
}

fn first_stop_index(text: &str, stop_sequences: &[String]) -> Option<usize> {
    stop_sequences
        .iter()
        .filter(|stop| !stop.is_empty())
        .filter_map(|stop| text.find(stop))
        .min()
}

fn render_chat_prompt(
    model: &LlamaModel,
    messages: &[(&str, &str)],
) -> Result<String, EngineError> {
    let chat_messages = messages
        .iter()
        .map(|(role, content)| {
            LlamaChatMessage::new((*role).to_string(), (*content).to_string())
                .map_err(|e| EngineError::Inference(format!("chat message: {e}")))
        })
        .collect::<Result<Vec<_>, _>>()?;

    match model.chat_template(None) {
        Ok(template) => model
            .apply_chat_template(&template, &chat_messages, true)
            .map_err(|e| EngineError::Inference(format!("chat template apply: {e}"))),
        Err(err) => {
            tracing::warn!("model chat template unavailable, falling back to ChatML: {err}");
            let template = LlamaChatTemplate::new("chatml")
                .map_err(|e| EngineError::Inference(format!("chatml template: {e}")))?;
            model
                .apply_chat_template(&template, &chat_messages, true)
                .map_err(|e| EngineError::Inference(format!("chatml template apply: {e}")))
        }
    }
}

fn parse_string_array(raw: &str) -> Result<Vec<String>, EngineError> {
    let parsed = serde_json::from_str::<Vec<String>>(raw.trim())
        .map_err(|e| EngineError::InvalidJson(format!("{e} in: {raw:?}")))?;

    let mut seen = std::collections::HashSet::new();
    let normalized = parsed
        .into_iter()
        .map(|item| item.trim().to_string())
        .filter(|item| !item.is_empty())
        .filter(|item| seen.insert(item.to_lowercase()))
        .collect();

    Ok(normalized)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn generate_fails_when_not_loaded() {
        let engine = LlmEngine::new();
        let err = engine
            .generate("hello", &GenerationConfig::default())
            .unwrap_err();
        assert!(matches!(err, EngineError::NotLoaded));
    }

    #[test]
    fn decompose_returns_error_when_not_loaded() {
        let engine = LlmEngine::new();
        assert!(!engine.is_loaded());
        let result = engine.decompose_task("Write a blog post");
        assert!(matches!(result, Err(EngineError::NotLoaded)));
    }

    #[test]
    fn default_creates_unloaded_engine() {
        let engine = LlmEngine::default();
        assert!(!engine.is_loaded());
    }

    #[test]
    fn engine_is_clone_and_shares_state() {
        let engine = LlmEngine::new();
        let clone = engine.clone();
        assert!(!clone.is_loaded());
    }

    #[test]
    fn parse_string_array_trims_and_dedupes() {
        let parsed = parse_string_array(r#"[" Buy milk ", "buy milk", "Call plumber"]"#).unwrap();
        assert_eq!(parsed, vec!["Buy milk", "Call plumber"]);
    }

    #[test]
    fn first_stop_index_finds_earliest_match() {
        let text = "hello<|im_end|>trailing";
        let index = first_stop_index(text, &["<|im_end|>".into(), "zzz".into()]);
        assert_eq!(index, Some(5));
    }

    #[test]
    fn prompt_preflight_rejects_oversized_prompt_tokens() {
        let err = preflight_context_window(7_105, 1_024).unwrap_err();
        assert!(matches!(
            err,
            EngineError::PromptTooLong {
                prompt_tokens: 7_105,
                max_tokens: 1_024,
                available_prompt_tokens: 7_104,
                context_window: MAX_CONTEXT_TOKENS,
            }
        ));
    }

    #[test]
    fn prompt_preflight_keeps_prompts_within_budget() {
        let n_ctx = preflight_context_window(7_104, 1_024).unwrap();
        assert_eq!(n_ctx, MAX_CONTEXT_TOKENS);
    }
}