goosedump 0.10.10

// 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;
use llama_cpp_2::token::LlamaToken;

/// 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 context_size = usize::try_from(EMBEDDING_CONTEXT).context("embedding context size")?;
        let mut chunks = Vec::new();
        for (text_index, text) in texts.iter().enumerate() {
            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");
            }
            chunks.extend(
                embedding_chunks(&tokens, separator, context_size)?
                    .into_iter()
                    .map(|chunk| (text_index, chunk)),
            );
        }
        let threads = inference_threads()?;
        let params = LlamaContextParams::default()
            .with_n_ctx(NonZeroU32::new(EMBEDDING_CONTEXT))
            .with_n_batch(EMBEDDING_CONTEXT)
            .with_n_ubatch(EMBEDDING_CONTEXT)
            .with_n_threads(threads)
            .with_n_threads_batch(threads)
            .with_n_seq_max(EMBEDDING_CONTEXT / 2)
            .with_kv_unified(true)
            .with_embeddings(true)
            .with_pooling_type(LlamaPoolingType::Mean);
        let mut ctx = self
            .model
            .new_context(&self.backend, params)
            .context("create embedding context")?;
        let mut embeddings = (0..texts.len()).map(|_| Vec::new()).collect::<Vec<_>>();
        let mut chunks = chunks.as_slice();
        while !chunks.is_empty() {
            let sequence_count = embedding_batch_len(chunks, context_size);
            let batch_chunks = &chunks[..sequence_count];
            let token_count = batch_chunks.iter().map(|(_, chunk)| chunk.len()).sum();
            let mut batch = LlamaBatch::new(token_count, 1);
            for (sequence, (_, chunk)) in batch_chunks.iter().enumerate() {
                let sequence = i32::try_from(sequence).context("embedding sequence index")?;
                batch.add_sequence(chunk, sequence, true)?;
            }
            ctx.clear_kv_cache();
            ctx.encode(&mut batch).context("encode embedding batch")?;
            for (sequence, (text_index, chunk)) in batch_chunks.iter().enumerate() {
                let sequence = i32::try_from(sequence).context("embedding sequence index")?;
                let embedding = ctx
                    .embeddings_seq_ith(sequence)
                    .context("read sequence embedding")?;
                embeddings[*text_index].push((normalize(embedding)?, chunk.len()));
            }
            chunks = &chunks[batch_chunks.len()..];
        }
        embeddings
            .iter()
            .map(|chunks| aggregate_embeddings(chunks))
            .collect()
    }
}

/// 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 embedding_chunks(
    tokens: &[LlamaToken],
    separator: LlamaToken,
    context_size: usize,
) -> anyhow::Result<Vec<Vec<LlamaToken>>> {
    if context_size < 2 {
        bail!("embedding context must hold BOS and EOS tokens");
    }
    let bos = *tokens
        .first()
        .context("embedding input is missing its BOS token")?;
    let content = tokens
        .get(1..tokens.len().saturating_sub(1))
        .context("embedding input is missing its EOS token")?;
    Ok(content
        .chunks(context_size - 2)
        .map(|chunk| {
            let mut sequence = Vec::with_capacity(chunk.len() + 2);
            sequence.push(bos);
            sequence.extend_from_slice(chunk);
            sequence.push(separator);
            sequence
        })
        .collect())
}

fn embedding_batch_len(chunks: &[(usize, Vec<LlamaToken>)], context_size: usize) -> usize {
    let mut token_count = 0;
    chunks
        .iter()
        .take_while(|(_, chunk)| {
            if chunk.len() > context_size - token_count {
                return false;
            }
            token_count += chunk.len();
            true
        })
        .count()
}

fn aggregate_embeddings(embeddings: &[(Vec<f32>, usize)]) -> anyhow::Result<Vec<f32>> {
    let Some((first, _)) = embeddings.first() else {
        bail!("embedding input produced no chunks");
    };
    let mut aggregate = vec![0.0; first.len()];
    let mut weight = 0.0_f32;
    for (embedding, len) in embeddings {
        if embedding.len() != aggregate.len() {
            bail!("embedding chunks have inconsistent dimensions");
        }
        let chunk_weight = f32::from(u16::try_from(*len).context("embedding chunk length")?);
        for (sum, value) in aggregate.iter_mut().zip(embedding) {
            *sum += value * chunk_weight;
        }
        weight += chunk_weight;
    }
    if weight == 0.0 {
        bail!("embedding chunks have zero total length");
    }
    for value in &mut aggregate {
        *value /= weight;
    }
    normalize(&aggregate)
}

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)
}