goosedump 0.11.2

// 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.

use std::env;
use std::path::{Path, PathBuf};

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

use crate::engine::gpt_oss::TextModel;

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

const TEXTGEN_CONTEXT: usize = 4096;

const INFERENCE_THREADS_ENV: &str = "GOOSEDUMP_INFERENCE_THREADS";

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

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

    /// 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!(
            "<|start|>system<|message|>{system}<|end|><|start|>user<|message|>{user}<|end|><|start|>assistant<|channel|>final<|message|>"
        );
        let generation = self
            .pool
            .install(|| self.model.generate(&prompt, max_tokens, TEXTGEN_CONTEXT))?;
        Ok(generation.text)
    }

    /// Merge two related memory entries into one concise factual statement.
    ///
    /// # Errors
    /// Returns an error if tokenization or inference fails.
    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}");
        self.complete(system, &user, 160)
    }
}

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

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 `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 target = dest.join(file);
        if target.is_file() {
            continue;
        }
        let cached = repo
            .get(file)
            .with_context(|| format!("download {repo_id}/{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)
}