Vectrum
Sign In
252 lines · 10.4 KB · Rust
//! Short generated text, such as thread titles, written by the model picked
//! under Settings → General → Text generation. Each request is a one-off,
//! tool-less agent run. Everything here blocks; call it from a background task.

use std::{path::Path, process::Stdio};

use crate::claude::{Effort, base_command};
use crate::models::{Provider, split_key};
use crate::session::Reasoning;
use crate::settings::{self, BranchNaming};

const TITLE_PROMPT: &str = "You name coding-assistant conversations. Reply with only a short title (3 to 6 words, \
sentence case, no quotes, no trailing period) for a conversation that starts with the user's message below.";

/// A title for a thread that starts with `message`, or `None` when text
/// generation is off or fails.
pub fn title(project: &Path, message: &str) -> Option<String> {
    let message: String = message.chars().take(2_000).collect();
    clean_title(&generate(project, &format!("{TITLE_PROMPT}\n\n<message>\n{message}\n</message>"))?)
}

/// A branch name for a worktree whose thread starts with `message`, as
/// T3 Code asks for one: a fragment the configured prefix goes in front of,
/// one with a semantic prefix (`feat/…`), or a whole name following the
/// user's instructions. `None` when text generation is off or fails.
pub fn branch_name(project: &Path, message: &str, naming: BranchNaming, instructions: &str) -> Option<String> {
    let mut rules = vec![
        "Branch should describe the requested work from the user message.",
        "Return a valid Git branch name without spaces.",
    ];
    match naming {
        BranchNaming::Custom => rules.push(
            "Return the complete branch name, following the user's naming instructions. No prefix or suffix will be added.",
        ),
        BranchNaming::Static | BranchNaming::Semantic => {
            rules.push("Keep it short and specific (2-6 words), in lowercase with hyphen-separated words.");
            rules.push(if naming == BranchNaming::Semantic {
                "Include a semantic prefix and a slash in the branch name, for example feat/add-search, fix/login-error, \
                 refactor/auth, docs/setup, or chore/update-deps. Choose the prefix that best describes the work."
            } else {
                "Return only the descriptive branch fragment, without a prefix or namespace. The application adds the \
                 configured prefix."
            });
        }
    }
    let message: String = message.chars().take(8_000).collect();
    let mut prompt = format!(
        "You generate concise git branch names.\nReturn a JSON object with key: branch.\nRules:\n{}\n\nUser message:\n{message}",
        rules.iter().map(|rule| format!("- {rule}")).collect::<Vec<_>>().join("\n")
    );
    if naming == BranchNaming::Custom && !instructions.trim().is_empty() {
        prompt.push_str(&format!("\n\nAdditional instructions:\n{}", instructions.trim()));
    }
    let value = json_reply(&generate(project, &prompt)?)?;
    let branch = value.get("branch")?.as_str()?.trim().to_string();
    (!branch.is_empty()).then_some(branch)
}

/// A commit message for what's staged, as T3 Code asks for one: an
/// imperative subject of at most 72 characters and an optional body.
pub fn commit_message(project: &Path, branch: Option<&str>, staged_summary: &str, staged_patch: &str, instructions: &str) -> Option<String> {
    let mut prompt = [
        "You write concise git commit messages.",
        "Return a JSON object with keys: subject, body.",
        "Rules:",
        "- subject must be imperative, <= 72 chars, and no trailing period",
        "- body can be empty string or short bullet points",
        "- capture the primary user-visible or developer-visible change",
        "",
        &format!("Branch: {}", branch.unwrap_or("(detached)")),
        "",
        "Staged files:",
        &limit(staged_summary, 6_000),
        "",
        "Staged patch:",
        &limit(staged_patch, 40_000),
    ]
    .join("\n");
    push_instructions(&mut prompt, instructions);
    let value = json_reply(&generate(project, &prompt)?)?;
    let subject = value.get("subject")?.as_str()?.trim().trim_end_matches('.').to_string();
    if subject.is_empty() {
        return None;
    }
    let body = value.get("body").and_then(|body| body.as_str()).unwrap_or_default().trim().to_string();
    Some(if body.is_empty() { subject } else { format!("{subject}\n\n{body}") })
}

/// A pull request's title and body for a branch, as T3 Code asks for them.
pub fn pull_request_content(
    project: &Path,
    base: &str,
    head: &str,
    commits: &str,
    diff_stat: &str,
    diff_patch: &str,
    template: Option<&str>,
    instructions: &str,
) -> Option<(String, String)> {
    let body_rules: &[&str] = if template.is_some() {
        &[
            "- body must be markdown and follow the repository change request template structure",
            "- fill in the template sections appropriately for this change",
            "- drop HTML comments from the template in the generated body",
            "- keep the template's markdown structure",
        ]
    } else {
        &[
            "- body must be markdown and include headings '## Summary' and '## Testing'",
            "- under Summary, provide short bullet points",
            "- under Testing, include bullet points with concrete checks or 'Not run' where appropriate",
        ]
    };
    let mut lines: Vec<String> = vec![
        "You write source control change request content.".into(),
        "Return a JSON object with keys: title, body.".into(),
        "Rules:".into(),
        "- title should be concise and specific".into(),
    ];
    lines.extend(body_rules.iter().map(|rule| rule.to_string()));
    if !instructions.trim().is_empty() {
        lines.extend(["".into(), "Additional instructions:".into(), limit(instructions.trim(), 20_000)]);
    }
    if let Some(template) = template {
        lines.extend(["".into(), "Repository change request template:".into(), limit(template, 8_000)]);
    }
    let prompt = [
        lines.join("\n").as_str(),
        "",
        &format!("Base branch: {base}"),
        &format!("Head branch: {head}"),
        "",
        "Commits:",
        &limit(commits, 12_000),
        "",
        "Diff stat:",
        &limit(diff_stat, 12_000),
        "",
        "Diff patch:",
        &limit(diff_patch, 40_000),
    ]
    .join("\n");
    let value = json_reply(&generate(project, &prompt)?)?;
    let title = value.get("title")?.as_str()?.trim().to_string();
    let body = value.get("body").and_then(|body| body.as_str()).unwrap_or_default().to_string();
    (!title.is_empty()).then_some((title, body))
}

/// T3 Code's "Additional instructions" section, when there are any.
fn push_instructions(prompt: &mut String, instructions: &str) {
    if !instructions.trim().is_empty() {
        prompt.push_str(&format!("\n\nAdditional instructions:\n{}", limit(instructions.trim(), 20_000)));
    }
}

/// The JSON object in a model's reply, which may sit in a code fence.
fn json_reply(reply: &str) -> Option<serde_json::Value> {
    let start = reply.find('{')?;
    let end = reply.rfind('}')?;
    serde_json::from_str(reply.get(start..=end)?).ok()
}

/// At most `max` characters of `text`.
fn limit(text: &str, max: usize) -> String {
    if text.chars().count() <= max {
        return text.to_string();
    }
    let mut cut: String = text.chars().take(max).collect();
    cut.push_str("\n[truncated]");
    cut
}

/// Runs `prompt` through the project's text generation model and returns
/// its reply.
fn generate(project: &Path, prompt: &str) -> Option<String> {
    let (key, reasoning) =
        settings::get_for(project, |settings| (settings.text_model.clone(), settings.text_reasoning.clone()));
    let (instance_id, model) = split_key(&key)?;
    let kind = settings::get(|settings| settings.instance(instance_id).map(|instance| instance.kind))
        .or_else(|| Provider::from_key(instance_id))?;
    let instance = settings::instance(instance_id, kind);
    if !instance.enabled {
        return None;
    }
    let effort = reasoning.as_deref().and_then(Reasoning::from_key).and_then(|reasoning| match reasoning {
        Reasoning::Level(effort) => Some(effort),
        _ => None,
    });
    if !kind.answers_one_off() {
        return None;
    }
    let (_, mut command) = base_command(&instance).ok()?;
    match kind {
        Provider::Claude => {
            command.args(["-p", "--output-format", "text", "--no-session-persistence", "--tools", ""]);
            command.args(["--strict-mcp-config", "--model", model]);
            if let Some(effort) = effort.filter(|effort| *effort >= Effort::Low) {
                command.args(["--effort", effort.as_cli()]);
            }
        }
        Provider::Pi => {
            command.args(["-p", "--no-session"]);
            if let Some((provider, model)) = model.split_once('/') {
                command.args(["--provider", provider, "--model", model]);
            }
            if let Some(effort) = effort {
                command.args(["--thinking", effort.as_cli()]);
            }
        }
        Provider::Jcode => return None,
    }
    let output = command
        .arg(prompt)
        .current_dir(std::env::temp_dir())
        .stdin(Stdio::null())
        .stderr(Stdio::null())
        .output()
        .ok()?;
    if !output.status.success() {
        return None;
    }
    Some(String::from_utf8_lossy(&output.stdout).into_owned())
}

/// The title out of a model's reply: its last non-empty line, without
/// quotes, markup or a trailing period, and not too long.
fn clean_title(reply: &str) -> Option<String> {
    let line = reply.lines().rev().map(str::trim).find(|line| !line.is_empty())?;
    let line = line.trim_start_matches(['#', '*', '-', ' ']).trim_start_matches("Title:").trim();
    let line = line.trim_matches(['"', '\'', '`', '*', '“', '”']).trim_end_matches('.').trim();
    if line.is_empty() {
        return None;
    }
    let mut title: String = line.chars().take(60).collect();
    if line.chars().count() > 60 {
        title.push('…');
    }
    Some(title)
}

#[cfg(test)]
mod tests {
    use super::clean_title;

    #[::core::prelude::v1::test]
    fn cleans_titles() {
        assert_eq!(clean_title("\"Fix login redirect.\"\n").as_deref(), Some("Fix login redirect"));
        assert_eq!(clean_title("Title: Add dark mode").as_deref(), Some("Add dark mode"));
        assert_eq!(clean_title("**Refactor settings**").as_deref(), Some("Refactor settings"));
        assert_eq!(clean_title("  \n "), None);
    }
}