goosedump 0.12.21

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

//! Local model layer. The GGUF model is downloaded to the local cache on
//! first use, then loaded for local inference.

mod bge;
mod gguf;
mod gpt_oss;
mod kernels;
mod tokenizer;
mod types;
mod wordpiece;

use std::env;
use std::fs::File;
use std::io::Read as _;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};

use anyhow::{Context as _, bail};
use rayon::{ThreadPool, ThreadPoolBuilder};
use sha2::{Digest as _, Sha256};

pub(crate) use self::bge::EMBEDDING_DIMENSIONS;
use self::bge::EmbeddingModel;
use self::gpt_oss::{Generation, GenerationStopReason, TextModel};

/// Cache subdirectory and Hugging Face source of the fixed text model.
const TEXTGEN_NAME: &str = "gpt-oss-20b-mxfp4-gguf";
const TEXTGEN_WEIGHTS_REPO: &str = "ggml-org/gpt-oss-20b-GGUF";
const TEXTGEN_WEIGHTS_FILE: &str = "gpt-oss-20b-MXFP4.gguf";

/// Cache identity and pinned Hugging Face source of the sentence embedding model.
pub(crate) const EMBEDDING_NAME: &str = "bge-small-en-v1.5-q8_0-gguf";
pub(crate) const EMBEDDING_WEIGHTS_REPO: &str = "ggml-org/bge-small-en-v1.5-Q8_0-GGUF";
pub(crate) const EMBEDDING_WEIGHTS_FILE: &str = "bge-small-en-v1.5-q8_0.gguf";
pub(crate) const EMBEDDING_WEIGHTS_SHA256: &str =
    "f046db1dc724cf4f6f0a0c5917e922823b73eb1d27b8f9a9c2797f7866974804";
pub(crate) const EMBEDDING_MODEL_ID: &str = "bge-small-en-v1.5-q8_0-pipeline-v1@f046db1dc724cf4f6f0a0c5917e922823b73eb1d27b8f9a9c2797f7866974804";

const TEXTGEN_CONTEXT: u16 = 4_096;

const INFERENCE_THREADS_ENV: &str = "GOOSEDUMP_INFERENCE_THREADS";
const EMBEDDING_SINGLE_THREADS: usize = 2;

/// Fixed GPT-OSS-20B MXFP4 model used for directive selection and memory
/// extraction.
pub struct TextGen {
    model: TextModel,
    pool: ThreadPool,
}

/// Fixed BGE-small-en-v1.5 model used for semantic retrieval.
pub struct Embedder {
    model: EmbeddingModel,
    single_pool: OnceLock<ThreadPool>,
    batch_pool: OnceLock<ThreadPool>,
    pool_init: Mutex<()>,
    max_threads: usize,
}

impl TextGen {
    /// Download and load the text model from the local model cache.
    ///
    /// # Errors
    /// Returns an error if downloading, loading, or configuring inference fails.
    pub fn load() -> anyhow::Result<Self> {
        let dir = model_cache_dir(TEXTGEN_NAME);
        ensure_files(TEXTGEN_WEIGHTS_REPO, &[TEXTGEN_WEIGHTS_FILE], &dir)?;
        let model = TextModel::load(dir.join(TEXTGEN_WEIGHTS_FILE))?;
        let threads = inference_threads()?;
        let thread_count = usize::try_from(threads).context("text inference thread count")?;
        let pool = ThreadPoolBuilder::new()
            .num_threads(thread_count)
            .build()
            .context("create text inference thread pool")?;
        Ok(Self { model, pool })
    }

    /// Return whether the formatted prompt and requested generation fit the context.
    ///
    /// # Errors
    /// Returns an error if tokenization fails.
    pub fn completion_fits(
        &self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<bool> {
        let prompt_tokens = self.model.prompt_tokens(&completion_prompt(system, user))?;
        Ok(prompt_tokens
            .checked_add(max_tokens)
            .is_some_and(|tokens| tokens <= usize::from(TEXTGEN_CONTEXT)))
    }

    /// Greedy instruction completion of `system` + `user`, capped at
    /// `max_tokens`.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails, or generation is stopped by a safety guard.
    pub fn complete(
        &mut self,
        system: &str,
        user: &str,
        max_tokens: usize,
    ) -> anyhow::Result<String> {
        let prompt = completion_prompt(system, user);
        let generation = self
            .pool
            .install(|| self.model.generate(&prompt, max_tokens, TEXTGEN_CONTEXT))?;
        if env::var_os("GOOSEDUMP_PROFILE_COMPACT").is_some() {
            let prefill_seconds = generation.prefill_duration.as_secs_f64();
            let decode_seconds = generation.decode_duration.as_secs_f64();
            eprintln!(
                "[inference-profile] prefill: tokens={} duration_ms={:.1} tokens_per_second={:.1}",
                generation.prompt_tokens,
                prefill_seconds * 1_000.0,
                token_rate(generation.prompt_tokens, prefill_seconds)?,
            );
            eprintln!(
                "[inference-profile] decode: tokens={} duration_ms={:.1} tokens_per_second={:.1}",
                generation.generated_tokens,
                decode_seconds * 1_000.0,
                token_rate(generation.generated_tokens, decode_seconds)?,
            );
        }
        finish_generation(generation)
    }
}

fn finish_generation(generation: Generation) -> anyhow::Result<String> {
    match generation.stop_reason {
        GenerationStopReason::EndOfSequence | GenerationStopReason::TokenLimit => {
            Ok(generation.text)
        }
        GenerationStopReason::TokenCycle => bail!(
            "text generation stopped after detecting a repeating token cycle ({} tokens)",
            generation.generated_tokens
        ),
        GenerationStopReason::TimeBudget => bail!(
            "text generation exceeded its decode time budget ({} tokens)",
            generation.generated_tokens
        ),
    }
}

impl Embedder {
    /// Download and load the embedding model from the local model cache.
    pub fn load() -> anyhow::Result<Self> {
        let dir = pull_embedding_model()?;
        Self::load_path(&dir.join(EMBEDDING_WEIGHTS_FILE))
    }

    /// Load the embedding model only when it is already cached.
    pub fn load_cached() -> anyhow::Result<Option<Self>> {
        let path = model_cache_dir(EMBEDDING_NAME).join(EMBEDDING_WEIGHTS_FILE);
        if !path.is_file() {
            return Ok(None);
        }
        verify_sha256(&path, EMBEDDING_WEIGHTS_SHA256)?;
        Self::load_path(&path).map(Some)
    }

    /// Produce one normalized sentence embedding.
    pub fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>> {
        self.pool(false)?.install(|| self.model.embed(text))
    }

    /// Produce normalized embeddings for an independent batch of sentences.
    pub fn embed_batch(&self, texts: &[&str]) -> anyhow::Result<Vec<Vec<f32>>> {
        if texts.len() <= 1 {
            return texts.first().map_or_else(
                || Ok(Vec::new()),
                |text| self.embed(text).map(|value| vec![value]),
            );
        }
        self.pool(true)?.install(|| self.model.embed_batch(texts))
    }

    fn pool(&self, batch: bool) -> anyhow::Result<&ThreadPool> {
        let single_threads = self.max_threads.min(EMBEDDING_SINGLE_THREADS);
        let (pool, threads) = if !batch || self.max_threads == single_threads {
            (&self.single_pool, single_threads)
        } else {
            (&self.batch_pool, self.max_threads)
        };
        if let Some(pool) = pool.get() {
            return Ok(pool);
        }
        let _guard = self
            .pool_init
            .lock()
            .map_err(|_| anyhow::anyhow!("embedding inference pool initialization is poisoned"))?;
        if let Some(pool) = pool.get() {
            return Ok(pool);
        }
        let candidate = ThreadPoolBuilder::new()
            .num_threads(threads)
            .build()
            .context("create embedding inference thread pool")?;
        pool.set(candidate)
            .map_err(|_| anyhow::anyhow!("embedding inference pool initialized twice"))?;
        pool.get()
            .context("initialize embedding inference thread pool")
    }

    fn load_path(path: &Path) -> anyhow::Result<Self> {
        let model = EmbeddingModel::load(path)?;
        let threads = inference_threads()?;
        let max_threads = usize::try_from(threads).context("embedding inference thread count")?;
        Ok(Self {
            model,
            single_pool: OnceLock::new(),
            batch_pool: OnceLock::new(),
            pool_init: Mutex::new(()),
            max_threads,
        })
    }
}

fn completion_prompt(system: &str, user: &str) -> String {
    format!(
        "<|start|>system<|message|>{system}<|end|><|start|>user<|message|>{user}<|end|><|start|>assistant<|channel|>final<|message|>"
    )
}

fn token_rate(tokens: usize, seconds: f64) -> anyhow::Result<f64> {
    if seconds == 0.0 {
        Ok(0.0)
    } else {
        let tokens = u32::try_from(tokens).context("token count exceeds u32")?;
        Ok(f64::from(tokens) / seconds)
    }
}

pub(crate) fn inference_threads() -> anyhow::Result<i32> {
    let available = std::thread::available_parallelism()
        .context("detect available parallelism")?
        .get();
    configured_threads(INFERENCE_THREADS_ENV, available)
}

fn configured_threads(name: &str, default: usize) -> anyhow::Result<i32> {
    let value = match env::var_os(name) {
        Some(value) => {
            let value = value
                .into_string()
                .map_err(|_| anyhow::anyhow!("{name} is not valid UTF-8"))?;
            let threads = value
                .parse::<usize>()
                .with_context(|| format!("parse {name}"))?;
            if threads == 0 {
                bail!("{name} must be greater than zero");
            }
            threads
        }
        None => default,
    };
    i32::try_from(value).with_context(|| format!("{name} 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 the fixed text-generation model.
pub(crate) fn pull_text_model() -> anyhow::Result<PathBuf> {
    let dir = model_cache_dir(TEXTGEN_NAME);
    ensure_files(TEXTGEN_WEIGHTS_REPO, &[TEXTGEN_WEIGHTS_FILE], &dir)?;
    Ok(dir)
}

/// Fetch and verify the pinned embedding model.
pub(crate) fn pull_embedding_model() -> anyhow::Result<PathBuf> {
    let dir = model_cache_dir(EMBEDDING_NAME);
    let path = dir.join(EMBEDDING_WEIGHTS_FILE);
    if path.is_file() {
        if verify_sha256(&path, EMBEDDING_WEIGHTS_SHA256).is_ok() {
            return Ok(dir);
        }
        std::fs::remove_file(&path)
            .with_context(|| format!("remove invalid model {}", path.display()))?;
    }
    ensure_files(EMBEDDING_WEIGHTS_REPO, &[EMBEDDING_WEIGHTS_FILE], &dir)?;
    verify_sha256(&path, EMBEDDING_WEIGHTS_SHA256)?;
    Ok(dir)
}

fn verify_sha256(path: &Path, expected: &str) -> anyhow::Result<()> {
    let mut file = File::open(path).with_context(|| format!("open {}", path.display()))?;
    let mut digest = Sha256::new();
    let mut buffer = vec![0_u8; 64 * 1_024];
    loop {
        let read = file
            .read(&mut buffer)
            .with_context(|| format!("read {}", path.display()))?;
        if read == 0 {
            break;
        }
        digest.update(&buffer[..read]);
    }
    let actual = format!("{:x}", digest.finalize());
    if actual != expected {
        bail!(
            "model {} has SHA-256 {actual}, expected {expected}; remove the file and pull it again",
            path.display()
        );
    }
    Ok(())
}

/// Fetch `files` from the Hugging Face `repo_id` into `dest` (flat layout).
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 target = dest.join(file);
        if target.is_file() {
            continue;
        }
        let cached = repo
            .get(file)
            .with_context(|| format!("download {repo_id}/{file}"))?;
        let temporary = dest.join(format!(".{file}.{}.part", std::process::id()));
        std::fs::copy(&cached, &temporary)
            .with_context(|| format!("write {}", temporary.display()))?;
        std::fs::rename(&temporary, &target)
            .with_context(|| format!("install {}", target.display()))?;
    }
    Ok(())
}

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