GitLite/src-tauri/crates/commit_ai/src/lib.rs

mod cloud;

pub use cloud::{generate_anthropic, generate_custom, generate_openai};

use std::{
    collections::hash_map::DefaultHasher,
    hash::{Hash, Hasher},
    sync::Arc,
};

use mistralrs::{GgufModelBuilder, Model, RequestBuilder, TextMessageRole};
use tokio::sync::RwLock;

/// One selectable local (on-device) model. Larger models produce better commit messages
/// but take longer to download (first run only, then cached) and run slower on CPU.
#[derive(Debug, Clone, serde::Serialize)]
pub struct LocalModelOption {
    pub id: &'static str,
    pub label: &'static str,
    pub approx_size_mb: u32,
    repo: &'static str,
    file: &'static str,
    tokenizer_repo: &'static str,
}

pub const DEFAULT_LOCAL_MODEL_ID: &str = "qwen2.5-0.5b";

pub const LOCAL_MODELS: &[LocalModelOption] = &[
    LocalModelOption {
        id: "qwen2.5-0.5b",
        label: "Qwen2.5 0.5B Instruct — fast, lower quality",
        approx_size_mb: 490,
        repo: "Qwen/Qwen2.5-0.5B-Instruct-GGUF",
        file: "qwen2.5-0.5b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-0.5B-Instruct",
    },
    LocalModelOption {
        id: "qwen2.5-1.5b",
        label: "Qwen2.5 1.5B Instruct — recommended",
        approx_size_mb: 1050,
        repo: "Qwen/Qwen2.5-1.5B-Instruct-GGUF",
        file: "qwen2.5-1.5b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-1.5B-Instruct",
    },
    LocalModelOption {
        id: "qwen2.5-3b",
        label: "Qwen2.5 3B Instruct — best quality, slower",
        approx_size_mb: 2100,
        repo: "Qwen/Qwen2.5-3B-Instruct-GGUF",
        file: "qwen2.5-3b-instruct-q4_k_m.gguf",
        tokenizer_repo: "Qwen/Qwen2.5-3B-Instruct",
    },
];

fn find_local_model(model_id: &str) -> Option<&'static LocalModelOption> {
    LOCAL_MODELS.iter().find(|option| option.id == model_id)
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LocalGenerationProfile {
    Fast,
    Balanced,
    Detailed,
}

impl Default for LocalGenerationProfile {
    fn default() -> Self {
        Self::Fast
    }
}

impl LocalGenerationProfile {
    pub fn from_id(value: Option<&str>) -> Self {
        match value
            .unwrap_or_default()
            .trim()
            .to_ascii_lowercase()
            .as_str()
        {
            "balanced" => Self::Balanced,
            "detailed" => Self::Detailed,
            _ => Self::Fast,
        }
    }

    pub fn diff_unified_context(self) -> &'static str {
        match self {
            Self::Fast => "--unified=1",
            Self::Balanced => "--unified=2",
            Self::Detailed => "--unified=3",
        }
    }

    fn max_diff_chars(self) -> usize {
        match self {
            Self::Fast => 8_000,
            Self::Balanced => 12_000,
            Self::Detailed => 24_000,
        }
    }

    fn max_output_tokens(self) -> usize {
        match self {
            Self::Fast => 160,
            Self::Balanced => 360,
            Self::Detailed => 750,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum CommitAiPhase {
    /// Nothing has been requested yet.
    Idle,
    /// Downloading (first run only, then cached by hf-hub) and/or loading into memory.
    Loading,
    Ready,
    Error,
}

#[derive(Debug, Clone, serde::Serialize)]
pub struct CommitAiStatus {
    pub phase: CommitAiPhase,
    pub model_id: Option<String>,
    pub error: Option<String>,
}

struct Inner {
    phase: CommitAiPhase,
    model_id: Option<String>,
    error: Option<String>,
    model: Option<Arc<Model>>,
    cache: Option<GenerationCache>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct GenerationCacheKey {
    model_id: String,
    profile: LocalGenerationProfile,
    input_hash: u64,
}

#[derive(Debug, Clone)]
struct GenerationCache {
    key: GenerationCacheKey,
    message: String,
}

/// Manages the local (on-device) model only. Cloud providers are stateless HTTP calls
/// (see [`cloud`]) and don't need this — there's nothing to download or keep loaded.
#[derive(Clone)]
pub struct CommitAiEngine {
    inner: Arc<RwLock<Inner>>,
}

impl Default for CommitAiEngine {
    fn default() -> Self {
        Self {
            inner: Arc::new(RwLock::new(Inner {
                phase: CommitAiPhase::Idle,
                model_id: None,
                error: None,
                model: None,
                cache: None,
            })),
        }
    }
}

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

    pub async fn status(&self) -> CommitAiStatus {
        let guard = self.inner.read().await;
        CommitAiStatus {
            phase: guard.phase,
            model_id: guard.model_id.clone(),
            error: guard.error.clone(),
        }
    }

    /// Downloads (first run only; hf-hub caches the files afterwards) and loads the given
    /// local model. Safe to call repeatedly — a call for the model that's already
    /// ready/loading is a no-op; a call for a *different* model switches to it (the
    /// previous one is dropped once no generation is still using it).
    pub async fn ensure_loaded(&self, model_id: &str) {
        {
            let guard = self.inner.read().await;
            let same_model = guard.model_id.as_deref() == Some(model_id);
            if same_model && matches!(guard.phase, CommitAiPhase::Ready | CommitAiPhase::Loading) {
                return;
            }
        }

        let Some(option) = find_local_model(model_id) else {
            let mut guard = self.inner.write().await;
            guard.phase = CommitAiPhase::Error;
            guard.model_id = Some(model_id.to_string());
            guard.error = Some(format!("Unknown local model: {model_id}"));
            guard.cache = None;
            return;
        };

        {
            let mut guard = self.inner.write().await;
            guard.phase = CommitAiPhase::Loading;
            guard.model_id = Some(model_id.to_string());
            guard.error = None;
            guard.model = None;
            guard.cache = None;
        }

        let result = GgufModelBuilder::new(option.repo, vec![option.file])
            .with_tok_model_id(option.tokenizer_repo)
            .with_logging()
            .build()
            .await;

        let mut guard = self.inner.write().await;
        // If the user switched to yet another model while this one was loading, drop this
        // (now stale) result instead of overwriting the newer request's state.
        if guard.model_id.as_deref() != Some(model_id) {
            return;
        }
        match result {
            Ok(model) => {
                guard.model = Some(Arc::new(model));
                guard.phase = CommitAiPhase::Ready;
                guard.error = None;
                guard.cache = None;
            }
            Err(err) => {
                guard.phase = CommitAiPhase::Error;
                guard.error = Some(err.to_string());
                guard.cache = None;
            }
        }
    }

    pub async fn generate_commit_message(
        &self,
        diff: &str,
        notes: Option<&str>,
        profile: LocalGenerationProfile,
    ) -> Result<String, String> {
        let (model, cache_key) = {
            let guard = self.inner.read().await;
            match (guard.phase, &guard.model) {
                (CommitAiPhase::Ready, Some(model)) => {
                    let cache_key = GenerationCacheKey {
                        model_id: guard.model_id.clone().unwrap_or_default(),
                        profile,
                        input_hash: generation_input_hash(diff, notes),
                    };
                    if let Some(cache) = &guard.cache {
                        if cache.key == cache_key {
                            return Ok(cache.message.clone());
                        }
                    }
                    (model.clone(), cache_key)
                }
                _ => return Err("The local AI model is not ready yet.".to_string()),
            }
        };

        let (system, user) = build_local_messages(diff, notes, profile)?;
        let request = RequestBuilder::new()
            .set_sampler_max_len(profile.max_output_tokens())
            .add_message(TextMessageRole::System, system)
            .add_message(TextMessageRole::User, user);

        let response = model
            .send_chat_request(request)
            .await
            .map_err(|err| err.to_string())?;

        let content = response
            .choices
            .first()
            .and_then(|choice| choice.message.content.clone())
            .ok_or_else(|| "The model did not return a response.".to_string())?;

        let message = sanitize_message(&content);
        if message.is_empty() {
            return Err("The model did not return a response.".to_string());
        }
        if looks_like_diff_echo(&message) {
            return Err(
                "The local model returned the diff instead of a commit message. Try a larger local model (1.5B or 3B) or a cloud provider.".to_string(),
            );
        }
        {
            let mut guard = self.inner.write().await;
            guard.cache = Some(GenerationCache {
                key: cache_key,
                message: message.clone(),
            });
        }
        Ok(message)
    }
}

fn generation_input_hash(diff: &str, notes: Option<&str>) -> u64 {
    let mut hasher = DefaultHasher::new();
    diff.hash(&mut hasher);
    notes.unwrap_or("").hash(&mut hasher);
    hasher.finish()
}

/// Models occasionally ignore the "no code fences" instruction (small local models
/// especially) — strip a wrapping ``` fence and wrapping quotes so the result can go
/// straight into the commit-message box.
pub(crate) fn sanitize_message(raw: &str) -> String {
    let mut text = raw.trim().to_string();
    if text.starts_with("```") {
        text = match text.split_once('\n') {
            // Drop the opening fence line (which may carry a language tag) and the closing fence.
            Some((_fence, rest)) => rest.trim_end().trim_end_matches("```").trim().to_string(),
            None => text.trim_matches('`').trim().to_string(),
        };
    }
    let trimmed = text.trim();
    if trimmed.len() >= 2 && trimmed.starts_with('"') && trimmed.ends_with('"') {
        return trimmed[1..trimmed.len() - 1].trim().to_string();
    }
    trimmed.to_string()
}

/// Weak models (small local ones especially) sometimes just echo the prompt's diff
/// sections back instead of writing a commit message. Catch that so the UI can show a
/// clear error instead of dumping raw diff text into the commit-message box.
pub(crate) fn looks_like_diff_echo(message: &str) -> bool {
    let lower = message.to_ascii_lowercase();
    lower.contains("diff --git")
        || lower.contains("staged files:")
        || lower.contains("staged changes:")
        || lower.contains("diff stat:")
        || lower.contains("detailed diff:")
        || message.lines().any(|line| line.starts_with("@@ "))
}

fn truncate_at_char_boundary(input: &str, max_chars: usize) -> String {
    if input.len() <= max_chars {
        return input.to_string();
    }

    let mut cut = max_chars;
    while !input.is_char_boundary(cut) {
        cut -= 1;
    }
    format!("{}\n\n[... diff truncated ...]", &input[..cut])
}

pub(crate) fn build_local_messages(
    diff: &str,
    notes: Option<&str>,
    profile: LocalGenerationProfile,
) -> Result<(String, String), String> {
    if diff.trim().is_empty() {
        return Err("No staged changes available for a commit message.".to_string());
    }

    let diff = truncate_at_char_boundary(diff, profile.max_diff_chars());
    // Appended to every profile below: small local models occasionally just echo the input
    // (the "Staged files:" / "Diff stat:" / "Detailed diff:" sections built in git.rs's
    // `staged_diff_local`) instead of writing a new commit message. Naming those exact
    // section headers here makes the failure mode explicit enough for weak models to avoid.
    const ANTI_ECHO: &str = " Never repeat, quote, or paraphrase the diff or its headers — \
do not include 'diff --git', '@@', 'Staged files:', 'Diff stat:', or 'Detailed diff:' \
anywhere in your answer.";
    let system = match profile {
        LocalGenerationProfile::Fast => {
            format!(
                "You generate Git commit messages. Respond only with one Conventional Commits subject line: <type>(<scope>): <subject>. Max 72 characters. No body, bullets, preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
        LocalGenerationProfile::Balanced => {
            format!(
                "You generate Git commit messages. Respond with a Conventional Commits subject line, then a blank line, then at most two short bullet points. Keep every line under 72 characters. No preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
        LocalGenerationProfile::Detailed => {
            format!(
                "You generate Git commit messages. Respond with a Conventional Commits subject line, then a blank line, then a concise body and up to four short bullet points grouped by affected area/file. Keep every line under 72 characters. No preamble, explanation, code fences, or quotes. Answer in English.{ANTI_ECHO}"
            )
        }
    };

    let mut user = String::new();
    if let Some(n) = notes.filter(|n| !n.trim().is_empty()) {
        user.push_str(&format!("Developer notes:\n{n}\n\n"));
    }
    user.push_str(&format!("Staged changes:\n{diff}"));
    Ok((system, user))
}

pub(crate) fn build_messages(diff: &str, notes: Option<&str>) -> Result<(String, String), String> {
    if diff.trim().is_empty() {
        return Err("No staged changes available for a commit message.".to_string());
    }

    // Rough token estimate — small models often have an 8-32k context window.
    const MAX_CHARS: usize = 24_000;
    let diff = truncate_at_char_boundary(diff, MAX_CHARS);

    let system = "You are a tool that writes a Git commit message describing a staged diff. \
Output ONLY the commit message text. Never quote, restate, or paraphrase the diff itself: \
do not include lines starting with 'diff --git', '@@', '+', '-', 'index ', 'Staged files:', \
or 'Diff stat:' anywhere in your answer. \
Format: a Conventional Commits header (<type>(<scope>): <subject>) in imperative mood, \
max. 72 characters, then a blank line, then a body. \
The body is required: a short, general paragraph (2-4 sentences) summarizing what changed \
and why at a high level — do NOT enumerate every changed file individually. \
You may optionally add up to 3 bullet points (- ...) afterward, but only for the most \
significant changes overall, never one bullet or heading per file. \
Never use bold text, backticks, or markdown headings for file names. \
Lines in the body max. 72 characters. \
No preamble, no explanation, no code fences, answer in English.\n\n\
Example:\n\
Diff:\n\
diff --git a/src/auth.py b/src/auth.py\n\
+def hash_password(pw):\n\
+    return bcrypt.hash(pw)\n\
diff --git a/src/routes.py b/src/routes.py\n\
-if password == stored_password:\n\
+if bcrypt.check(password, stored_password):\n\n\
Commit message:\n\
feat(auth): hash and verify passwords with bcrypt\n\n\
Passwords were previously compared as plain text. This adds a bcrypt-based\n\
hashing helper and updates the login check to verify against the hash\n\
instead of a direct string comparison.\n\n\
- Hash passwords on write, verify with bcrypt on login"
        .to_string();

    let mut user = String::new();
    if let Some(n) = notes.filter(|n| !n.trim().is_empty()) {
        user.push_str(&format!("Developer notes:\n{n}\n\n"));
    }
    user.push_str(&format!("Staged diff:\n{diff}"));
    Ok((system, user))
}