Vectrum
Sign In
96 lines · 3.8 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::process::Stdio;

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

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(message: &str) -> Option<String> {
    let (key, reasoning) = settings::get(|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,
    });
    let message: String = message.chars().take(2_000).collect();
    let prompt = format!("{TITLE_PROMPT}\n\n<message>\n{message}\n</message>");

    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;
    }
    clean_title(&String::from_utf8_lossy(&output.stdout))
}

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