goosedump 0.10.0

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

//! llama.cpp model layer. GGUF models are downloaded to the local cache on
//! first use, then loaded locally for CPU inference.

use std::num::NonZeroU32;
use std::path::{Path, PathBuf};

use anyhow::{Context as _, bail};
use llama_cpp_2::context::params::{LlamaContextParams, LlamaPoolingType};
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, LlamaModel};
use llama_cpp_2::sampling::LlamaSampler;

/// Cache subdirectory and Hugging Face repo of the sentence-embedding model.
pub(crate) const EMBEDDER_NAME: &str = "bge-small-en-v1.5-gguf";
pub(crate) const EMBEDDER_REPO: &str = "CompendiumLabs/bge-small-en-v1.5-gguf";
pub(crate) const EMBEDDER_FILE: &str = "bge-small-en-v1.5-q8_0.gguf";
pub(crate) const EMBEDDER_FILES: [&str; 1] = [EMBEDDER_FILE];
pub(crate) const EMBEDDING_MODEL_ID: &str = "BAAI/bge-small-en-v1.5:q8_0:mean:l2:llama.cpp";

/// Cache subdirectory and Hugging Face source of the Qwen2 judge model.
pub(crate) const TEXTGEN_NAME: &str = "qwen2-0_5b-instruct-llama";
pub(crate) const TEXTGEN_WEIGHTS_REPO: &str = "Qwen/Qwen2-0.5B-Instruct-GGUF";
pub(crate) const TEXTGEN_WEIGHTS_FILE: &str = "qwen2-0_5b-instruct-q4_0.gguf";

/// Cache subdirectory and Hugging Face source of the optional Stage-3 mutator.
pub(crate) const MUTATOR_NAME: &str = "qwen3-0.6b-q4-k-m-gguf";
pub(crate) const MUTATOR_REPO: &str = "unsloth/Qwen3-0.6B-GGUF";
pub(crate) const MUTATOR_FILE: &str = "Qwen3-0.6B-Q4_K_M.gguf";

const EMBEDDING_CONTEXT: u32 = 512;
const TEXTGEN_CONTEXT: u32 = 2048;
const MUTATOR_CONTEXT: u32 = 2048;

/// Sentence-embedding model (`bge-small-en-v1.5`, 384-dim) used for semantic
/// deduplication and persistent-memory recall.
pub struct Embedder {
    model: LlamaModel,
    backend: LlamaBackend,
}

impl Embedder {
    /// Download and load the embedder from the local model cache.
    ///
    /// # Errors
    /// Returns an error if downloading or loading the model fails.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(EMBEDDER_NAME);
        ensure_files(EMBEDDER_REPO, &EMBEDDER_FILES, &dir)?;
        Self::from_cache(&dir)
    }

    fn from_cache(dir: &Path) -> anyhow::Result<Self> {
        let mut backend = LlamaBackend::init().context("initialize llama.cpp")?;
        backend.void_logs();
        let model = LlamaModel::load_from_file(
            &backend,
            dir.join(EMBEDDER_FILE),
            &LlamaModelParams::default(),
        )
        .context("load embedding model")?;
        Ok(Self { model, backend })
    }

    /// Mean-pooled, L2-normalized sentence embeddings, one row per input text.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails.
    pub fn embed(&self, texts: &[String]) -> anyhow::Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let threads = inference_threads()?;
        let params = LlamaContextParams::default()
            .with_n_ctx(NonZeroU32::new(EMBEDDING_CONTEXT))
            .with_n_threads(threads)
            .with_n_threads_batch(threads)
            .with_embeddings(true)
            .with_pooling_type(LlamaPoolingType::Mean);
        let mut ctx = self
            .model
            .new_context(&self.backend, params)
            .context("create embedding context")?;
        let context_size = usize::try_from(ctx.n_ctx()).context("embedding context size")?;
        let mut output = Vec::with_capacity(texts.len());

        for text in texts {
            let mut tokens = self
                .model
                .str_to_token(text, AddBos::Always)
                .with_context(|| format!("tokenize embedding input: {text}"))?;
            let separator = self.model.token_eos();
            if tokens.last() != Some(&separator) {
                tokens.push(separator);
            }
            if tokens.is_empty() {
                bail!("embedding input produced no tokens");
            }
            if tokens.len() > context_size {
                let final_token = tokens[tokens.len() - 1];
                tokens.truncate(context_size);
                tokens[context_size - 1] = final_token;
            }

            let mut batch = LlamaBatch::new(tokens.len(), 1);
            batch.add_sequence(&tokens, 0, false)?;
            ctx.clear_kv_cache();
            ctx.encode(&mut batch).context("encode embedding input")?;
            let embedding = ctx
                .embeddings_seq_ith(0)
                .context("read sequence embedding")?;
            output.push(normalize(embedding)?);
        }
        Ok(output)
    }
}

/// Quantized `Qwen2-0.5B-Instruct` model used by `compact` to judge which
/// deterministic directive candidates remain current.
pub struct TextGen {
    model: LlamaModel,
    backend: LlamaBackend,
}

impl TextGen {
    /// Download and load the text model from the local model cache.
    ///
    /// # Errors
    /// Returns an error if downloading or loading the model fails.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(TEXTGEN_NAME);
        ensure_files(TEXTGEN_WEIGHTS_REPO, &[TEXTGEN_WEIGHTS_FILE], &dir)?;
        Self::from_cache(&dir)
    }

    fn from_cache(dir: &Path) -> anyhow::Result<Self> {
        let mut backend = LlamaBackend::init().context("initialize llama.cpp")?;
        backend.void_logs();
        let model = LlamaModel::load_from_file(
            &backend,
            dir.join(TEXTGEN_WEIGHTS_FILE),
            &LlamaModelParams::default(),
        )
        .context("load text generation model")?;
        Ok(Self { model, backend })
    }

    /// Greedy instruction completion of `system` + `user`, capped at
    /// `max_tokens`.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails.
    pub fn complete(
        &mut self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<String> {
        let prompt = format!(
            "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n"
        );
        complete_model(
            &self.model,
            &self.backend,
            TEXTGEN_CONTEXT,
            &prompt,
            max_tokens,
        )
    }
}

/// Quantized `Qwen3-0.6B-Instruct` model used for optional Stage-3 memory
/// mutation. The prompt disables reasoning so the result is a single merge.
pub struct Mutator {
    model: LlamaModel,
    backend: LlamaBackend,
}

impl Mutator {
    /// Download and load the mutator from the local model cache.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(MUTATOR_NAME);
        ensure_files(MUTATOR_REPO, &[MUTATOR_FILE], &dir)?;
        let mut backend = LlamaBackend::init().context("initialize llama.cpp")?;
        backend.void_logs();
        let model = LlamaModel::load_from_file(
            &backend,
            dir.join(MUTATOR_FILE),
            &LlamaModelParams::default(),
        )
        .context("load mutation model")?;
        Ok(Self { model, backend })
    }

    /// Merge two related memory entries into one concise factual statement.
    pub fn merge(&mut self, left: &str, right: &str) -> anyhow::Result<String> {
        let system = "Merge the two memory entries into one concise factual statement. Preserve concrete names, paths, commands, constraints, and unresolved work. Do not add information. Return only the merged statement.";
        let user = format!("[Memory A]\n{left}\n\n[Memory B]\n{right}");
        let prompt = format!(
            "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
        );
        complete_model(&self.model, &self.backend, MUTATOR_CONTEXT, &prompt, 160)
    }
}

fn complete_model(
    model: &LlamaModel,
    backend: &LlamaBackend,
    context: u32,
    prompt: &str,
    max_tokens: usize,
) -> anyhow::Result<String> {
    let tokens = model
        .str_to_token(prompt, AddBos::Always)
        .context("tokenize text generation prompt")?;
    let context_size = usize::try_from(context).context("text context size")?;
    if tokens.is_empty() {
        bail!("text generation prompt produced no tokens");
    }
    if tokens.len().saturating_add(max_tokens) > context_size {
        bail!(
            "text generation requires {} tokens but context holds {context_size}",
            tokens.len().saturating_add(max_tokens)
        );
    }
    let threads = inference_threads()?;
    let params = LlamaContextParams::default()
        .with_n_ctx(NonZeroU32::new(context))
        .with_n_threads(threads)
        .with_n_threads_batch(threads);
    let mut ctx = model
        .new_context(backend, params)
        .context("create text generation context")?;
    let mut batch = LlamaBatch::new(tokens.len().max(1), 1);
    batch.add_sequence(&tokens, 0, false)?;
    ctx.decode(&mut batch)
        .context("decode text generation prompt")?;
    let mut sampler = LlamaSampler::greedy();
    let mut decoder = encoding_rs::UTF_8.new_decoder();
    let mut output = String::new();
    for (position, generated) in (batch.n_tokens()..).zip(0..max_tokens) {
        let token = sampler.sample(&ctx, batch.n_tokens() - 1);
        sampler.accept(token);
        if model.is_eog_token(token) {
            break;
        }
        let piece = model
            .token_to_piece(token, &mut decoder, true, None)
            .context("decode generated token")?;
        output.push_str(&piece);
        if generated + 1 == max_tokens {
            break;
        }
        batch.clear();
        batch.add(token, position, &[0], true)?;
        ctx.decode(&mut batch).context("decode generated token")?;
    }
    output.reserve(4);
    let (_, _, had_errors) = decoder.decode_to_string(b"", &mut output, true);
    if had_errors {
        bail!("generated text ended with invalid UTF-8");
    }
    Ok(output)
}

fn normalize(input: &[f32]) -> anyhow::Result<Vec<f32>> {
    let magnitude = input
        .iter()
        .fold(0.0_f32, |sum, value| value.mul_add(*value, sum))
        .sqrt();
    if !magnitude.is_finite() || magnitude == 0.0 {
        bail!("embedding has invalid magnitude");
    }
    Ok(input.iter().map(|value| value / magnitude).collect())
}

fn inference_threads() -> anyhow::Result<i32> {
    let threads = std::thread::available_parallelism()
        .context("detect available parallelism")?
        .get();
    i32::try_from(threads).context("thread count exceeds i32")
}

fn ensure_files(repo_id: &str, files: &[&str], dest: &Path) -> anyhow::Result<()> {
    if files.iter().all(|file| dest.join(file).is_file()) {
        return Ok(());
    }
    pull_files(repo_id, files, dest)
}

/// Fetch `files` from the Hugging Face `repo_id` into `dest` (flat layout).
pub(crate) fn pull_files(repo_id: &str, files: &[&str], dest: &Path) -> anyhow::Result<()> {
    std::fs::create_dir_all(dest).with_context(|| format!("create {}", dest.display()))?;
    let api = hf_hub::api::sync::Api::new()?;
    let repo = api.model(repo_id.to_string());
    for file in files {
        let cached = repo
            .get(file)
            .with_context(|| format!("download {repo_id}/{file}"))?;
        let target = dest.join(file);
        std::fs::copy(&cached, &target).with_context(|| format!("write {}", target.display()))?;
    }
    Ok(())
}

/// The per-model cache directory (honors `XDG_CACHE_HOME`).
pub(crate) fn model_cache_dir(name: &str) -> PathBuf {
    dirs::cache_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("goosedump")
        .join("models")
        .join(name)
}