Jake
3410d3c586
Recorder-style auto-titling for transcripts. Mirrors the Phase 9 content-tags pipeline so the same prompt-injection-hardened pattern, spawn_blocking discipline, and sanitisation-after-generation shape get reused; the user-facing surface (auto on save + on-demand button) lands in a follow-up commit. - crates/llm/src/prompts.rs: new TRANSCRIPT_TITLE_SYSTEM constant. Same injection guard wording as ai-formatting's CLEANUP_PROMPT — dictated speech is data, not instructions. Rules constrain output shape: 4-8 words, Title Case, no quotes, no terminal punctuation, "Untitled" fallback for empty input. - crates/llm/src/lib.rs: LlmEngine::generate_title returns Result<String, EngineError>. Mirrors extract_content_tags shape: trailing-2000-char UTF-8-boundary truncation, temperature 0, max_tokens 24, free-form output (no GBNF — titles are prose, not a closed set). Sanitisation runs server-side via the new private sanitize_title helper, which handles the real Qwen3 failure modes: surrounding curly + ASCII quotes, leading "Title:" prefix, multi-line output, trailing "." / "!" / "?", whitespace runs, 100-char cap, literal "Untitled" → None. Three unit tests cover composite real-world outputs end-to-end. kon-llm test suite goes 15 → 18 passing. The Tauri wrapper, invoke_handler registration, and frontend wiring follow in subsequent commits.
717 lines
25 KiB
Rust
717 lines
25 KiB
Rust
use std::num::NonZeroU32;
|
|
use std::path::Path;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use encoding_rs::UTF_8;
|
|
use llama_cpp_2::context::params::LlamaContextParams;
|
|
use llama_cpp_2::llama_backend::LlamaBackend;
|
|
use llama_cpp_2::llama_batch::LlamaBatch;
|
|
use llama_cpp_2::model::params::LlamaModelParams;
|
|
use llama_cpp_2::model::{AddBos, LlamaChatMessage, LlamaChatTemplate, LlamaModel};
|
|
use llama_cpp_2::sampling::LlamaSampler;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
pub mod grammars;
|
|
pub mod model_manager;
|
|
pub mod prompts;
|
|
|
|
pub use grammars::CONTENT_TAGS_GRAMMAR;
|
|
pub use model_manager::{recommend_tier, LlmModelId, LlmModelInfo};
|
|
pub use prompts::{
|
|
is_valid_intent, ContentTags, CONTENT_TAGS_SYSTEM, INTENT_CLOSED_SET, TRANSCRIPT_TITLE_SYSTEM,
|
|
};
|
|
|
|
const DEFAULT_CONTEXT_TOKENS: u32 = 4096;
|
|
const MAX_CONTEXT_TOKENS: u32 = 8192;
|
|
const CONTEXT_RESERVE_TOKENS: u32 = 64;
|
|
const GENERATION_SEED: u32 = 0;
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum EngineError {
|
|
#[error("LLM not loaded. Download an AI model in Settings.")]
|
|
NotLoaded,
|
|
#[error("LLM load failed: {0}")]
|
|
LoadFailed(String),
|
|
#[error(
|
|
"prompt too long: {prompt_tokens} prompt tokens exceed the {available_prompt_tokens}-token prompt budget for an {context_window}-token context with {max_tokens} reserved response tokens"
|
|
)]
|
|
PromptTooLong {
|
|
prompt_tokens: usize,
|
|
max_tokens: u32,
|
|
available_prompt_tokens: u32,
|
|
context_window: u32,
|
|
},
|
|
#[error("inference failed: {0}")]
|
|
Inference(String),
|
|
#[error("model output not valid JSON: {0}")]
|
|
InvalidJson(String),
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct GenerationConfig {
|
|
pub max_tokens: u32,
|
|
pub temperature: f32,
|
|
pub stop_sequences: Vec<String>,
|
|
pub grammar: Option<String>,
|
|
}
|
|
|
|
impl Default for GenerationConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_tokens: 1024,
|
|
temperature: 0.0,
|
|
stop_sequences: Vec::new(),
|
|
grammar: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct LoadedModelState {
|
|
pub model_id: String,
|
|
pub model_path: String,
|
|
pub use_gpu: bool,
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct LlmState {
|
|
backend: Option<Arc<LlamaBackend>>,
|
|
model: Option<Arc<LlamaModel>>,
|
|
loaded: Option<LoadedModelState>,
|
|
}
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct LlmEngine {
|
|
inner: Arc<Mutex<LlmState>>,
|
|
}
|
|
|
|
impl LlmEngine {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
pub fn load(&self, model_path: &Path) -> Result<(), EngineError> {
|
|
self.load_model(LlmModelId::default_tier(), model_path, true)
|
|
}
|
|
|
|
pub fn load_model(
|
|
&self,
|
|
model_id: LlmModelId,
|
|
model_path: &Path,
|
|
use_gpu: bool,
|
|
) -> Result<(), EngineError> {
|
|
let mut guard = self.inner.lock().unwrap();
|
|
|
|
if let Some(loaded) = &guard.loaded {
|
|
if loaded.model_id == model_id.as_str()
|
|
&& loaded.model_path == model_path.display().to_string()
|
|
&& loaded.use_gpu == use_gpu
|
|
{
|
|
return Ok(());
|
|
}
|
|
}
|
|
|
|
let backend = match guard.backend.clone() {
|
|
Some(existing) => existing,
|
|
None => Arc::new(
|
|
LlamaBackend::init()
|
|
.map_err(|e| EngineError::LoadFailed(format!("backend init: {e}")))?,
|
|
),
|
|
};
|
|
|
|
let gpu_layers = if use_gpu { u32::MAX } else { 0 };
|
|
let params = LlamaModelParams::default().with_n_gpu_layers(gpu_layers);
|
|
let model = LlamaModel::load_from_file(&backend, model_path, ¶ms)
|
|
.map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))?;
|
|
|
|
guard.backend = Some(backend);
|
|
guard.model = Some(Arc::new(model));
|
|
guard.loaded = Some(LoadedModelState {
|
|
model_id: model_id.as_str().to_string(),
|
|
model_path: model_path.display().to_string(),
|
|
use_gpu,
|
|
});
|
|
Ok(())
|
|
}
|
|
|
|
pub fn unload(&self) -> Result<(), EngineError> {
|
|
let mut guard = self.inner.lock().unwrap();
|
|
guard.model = None;
|
|
guard.backend = None;
|
|
guard.loaded = None;
|
|
Ok(())
|
|
}
|
|
|
|
pub fn is_loaded(&self) -> bool {
|
|
self.inner.lock().unwrap().model.is_some()
|
|
}
|
|
|
|
pub fn loaded_model(&self) -> Option<LoadedModelState> {
|
|
self.inner.lock().unwrap().loaded.clone()
|
|
}
|
|
|
|
pub fn loaded_model_id(&self) -> Option<String> {
|
|
self.loaded_model().map(|loaded| loaded.model_id)
|
|
}
|
|
|
|
pub fn generate(&self, prompt: &str, config: &GenerationConfig) -> Result<String, EngineError> {
|
|
let (backend, model) = self.loaded_handles()?;
|
|
let prompt_tokens = model
|
|
.str_to_token(prompt, AddBos::Never)
|
|
.map_err(|e| EngineError::Inference(format!("tokenize: {e}")))?;
|
|
if prompt_tokens.is_empty() {
|
|
return Ok(String::new());
|
|
}
|
|
|
|
let n_ctx = preflight_context_window(prompt_tokens.len(), config.max_tokens)?;
|
|
let thread_count = i32::try_from(num_cpus::get().max(1)).unwrap_or(4);
|
|
let ctx_params = LlamaContextParams::default()
|
|
.with_n_ctx(Some(
|
|
NonZeroU32::new(n_ctx).expect("n_ctx must be non-zero"),
|
|
))
|
|
.with_n_batch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
|
|
.with_n_ubatch(prompt_tokens.len().max(512).min(n_ctx as usize) as u32)
|
|
.with_n_threads(thread_count)
|
|
.with_n_threads_batch(thread_count);
|
|
let mut ctx = model
|
|
.new_context(&backend, ctx_params)
|
|
.map_err(|e| EngineError::Inference(format!("context: {e}")))?;
|
|
|
|
let mut batch = LlamaBatch::new(prompt_tokens.len().max(1), 1);
|
|
for (index, token) in prompt_tokens.iter().enumerate() {
|
|
batch
|
|
.add(*token, index as i32, &[0], index + 1 == prompt_tokens.len())
|
|
.map_err(|e| EngineError::Inference(format!("batch add: {e}")))?;
|
|
}
|
|
ctx.decode(&mut batch)
|
|
.map_err(|e| EngineError::Inference(format!("prefill decode: {e}")))?;
|
|
|
|
let mut sampler = self.build_sampler(&model, config)?;
|
|
let mut decoder = UTF_8.new_decoder();
|
|
let mut generated = String::new();
|
|
let mut cursor = prompt_tokens.len() as i32;
|
|
|
|
for _ in 0..config.max_tokens {
|
|
let next = sampler.sample(&ctx, batch.n_tokens() - 1);
|
|
if model.is_eog_token(next) || next == model.token_eos() {
|
|
break;
|
|
}
|
|
|
|
let piece = model
|
|
.token_to_piece(next, &mut decoder, true, None)
|
|
.map_err(|e| EngineError::Inference(format!("detokenize: {e}")))?;
|
|
generated.push_str(&piece);
|
|
sampler.accept(next);
|
|
|
|
if let Some(stop_index) = first_stop_index(&generated, &config.stop_sequences) {
|
|
generated.truncate(stop_index);
|
|
break;
|
|
}
|
|
|
|
batch.clear();
|
|
batch
|
|
.add(next, cursor, &[0], true)
|
|
.map_err(|e| EngineError::Inference(format!("sample batch: {e}")))?;
|
|
cursor += 1;
|
|
ctx.decode(&mut batch)
|
|
.map_err(|e| EngineError::Inference(format!("sample decode: {e}")))?;
|
|
}
|
|
|
|
Ok(generated.trim().to_string())
|
|
}
|
|
|
|
pub fn cleanup_text(
|
|
&self,
|
|
system_prompt: &str,
|
|
transcript: &str,
|
|
) -> Result<String, EngineError> {
|
|
if transcript.trim().is_empty() {
|
|
return Ok(String::new());
|
|
}
|
|
let model = self.loaded_model_arc()?;
|
|
let prompt =
|
|
render_chat_prompt(&model, &[("system", system_prompt), ("user", transcript)])?;
|
|
self.generate(
|
|
&prompt,
|
|
&GenerationConfig {
|
|
max_tokens: 1024,
|
|
temperature: 0.0,
|
|
stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
|
|
grammar: None,
|
|
},
|
|
)
|
|
}
|
|
|
|
pub fn decompose_task(&self, task_text: &str) -> Result<Vec<String>, EngineError> {
|
|
self.decompose_task_with_feedback(task_text, &[])
|
|
}
|
|
|
|
/// Same as `decompose_task` but allows callers to pass recent HITL
|
|
/// feedback rows so the system prompt gets conditioned on the
|
|
/// user's preferred decomposition style. The `examples` vec is
|
|
/// rendered into a few-shot block appended to the base system
|
|
/// prompt by `prompts::build_conditioned_system_prompt`.
|
|
///
|
|
/// Callers should pass most-recent-first; older examples still
|
|
/// participate but weigh less because of their position in the
|
|
/// prompt. Empty slice keeps behaviour identical to `decompose_task`.
|
|
pub fn decompose_task_with_feedback(
|
|
&self,
|
|
task_text: &str,
|
|
examples: &[prompts::FeedbackExample],
|
|
) -> Result<Vec<String>, EngineError> {
|
|
let model = self.loaded_model_arc()?;
|
|
let system =
|
|
prompts::build_conditioned_system_prompt(prompts::DECOMPOSE_TASK_SYSTEM, examples);
|
|
let prompt = render_chat_prompt(
|
|
&model,
|
|
&[
|
|
("system", system.as_str()),
|
|
("user", &format!("Task: {task_text}")),
|
|
],
|
|
)?;
|
|
let raw = self.generate(
|
|
&prompt,
|
|
&GenerationConfig {
|
|
max_tokens: 512,
|
|
temperature: 0.0,
|
|
stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
|
|
grammar: Some(grammars::TASK_ARRAY_GRAMMAR.to_string()),
|
|
},
|
|
)?;
|
|
parse_string_array(&raw)
|
|
}
|
|
|
|
pub fn extract_tasks(&self, transcript: &str) -> Result<Vec<String>, EngineError> {
|
|
self.extract_tasks_with_feedback(transcript, &[])
|
|
}
|
|
|
|
/// Phase 9 content-tag extraction. Emits a single (topic, intent)
|
|
/// pair under the `CONTENT_TAGS_GRAMMAR` GBNF. Truncates to the
|
|
/// trailing 2000 chars of the transcript so the prompt budget
|
|
/// stays well under any model's context window. Determinism is
|
|
/// enforced by temperature 0.0 and the closed-set intent grammar
|
|
/// rule; on the rare case the model emits a parse-able-but-out-of-
|
|
/// set intent, we re-validate with `is_valid_intent` and bubble
|
|
/// `InvalidJson` so the frontend toasts a clear error.
|
|
pub fn extract_content_tags(
|
|
&self,
|
|
transcript: &str,
|
|
) -> Result<prompts::ContentTags, EngineError> {
|
|
if transcript.trim().is_empty() {
|
|
return Err(EngineError::Inference("empty transcript".into()));
|
|
}
|
|
|
|
// Truncate to the last 2000 chars on a UTF-8 char boundary so
|
|
// we don't slice through a multi-byte sequence.
|
|
const MAX_CHARS: usize = 2000;
|
|
let tail = if transcript.len() > MAX_CHARS {
|
|
let mut adj = transcript.len() - MAX_CHARS;
|
|
while adj < transcript.len() && !transcript.is_char_boundary(adj) {
|
|
adj += 1;
|
|
}
|
|
&transcript[adj..]
|
|
} else {
|
|
transcript
|
|
};
|
|
|
|
let model = self.loaded_model_arc()?;
|
|
let prompt = render_chat_prompt(
|
|
&model,
|
|
&[
|
|
("system", prompts::CONTENT_TAGS_SYSTEM),
|
|
("user", &format!("Transcript:\n{tail}")),
|
|
],
|
|
)?;
|
|
let raw = self.generate(
|
|
&prompt,
|
|
&GenerationConfig {
|
|
max_tokens: 96,
|
|
temperature: 0.0,
|
|
stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
|
|
grammar: Some(grammars::CONTENT_TAGS_GRAMMAR.to_string()),
|
|
},
|
|
)?;
|
|
|
|
let tags: prompts::ContentTags = serde_json::from_str(raw.trim())
|
|
.map_err(|e| EngineError::InvalidJson(format!("{e}: raw={raw:?}")))?;
|
|
if !prompts::is_valid_intent(&tags.intent) {
|
|
return Err(EngineError::InvalidJson(format!(
|
|
"intent out of closed set: {}",
|
|
tags.intent,
|
|
)));
|
|
}
|
|
Ok(tags)
|
|
}
|
|
|
|
/// Generate a short scannable title for a transcript. Free-form
|
|
/// 4-8 word string, post-processed by [`sanitize_title`] to strip
|
|
/// the model's occasional "Title:" prefix, surrounding quotes,
|
|
/// trailing terminal punctuation, and to collapse internal
|
|
/// whitespace runs. Mirrors the `extract_content_tags` shape:
|
|
/// truncates input to the trailing 2000 chars on a UTF-8 boundary,
|
|
/// temperature 0, no GBNF (output is free-form prose).
|
|
///
|
|
/// Returns `Err(EngineError::Inference("could not derive title"))`
|
|
/// when the model emits an empty / "Untitled" response after
|
|
/// sanitisation; the caller (auto-trigger in the frontend) treats
|
|
/// that as a silent skip and leaves the row untitled.
|
|
pub fn generate_title(&self, transcript: &str) -> Result<String, EngineError> {
|
|
if transcript.trim().is_empty() {
|
|
return Err(EngineError::Inference("empty transcript".into()));
|
|
}
|
|
|
|
// Mirrors `extract_content_tags`: keep only the trailing 2000
|
|
// chars, snapped to a UTF-8 char boundary so we don't slice
|
|
// through a multi-byte sequence.
|
|
const MAX_CHARS: usize = 2000;
|
|
let tail = if transcript.len() > MAX_CHARS {
|
|
let mut adj = transcript.len() - MAX_CHARS;
|
|
while adj < transcript.len() && !transcript.is_char_boundary(adj) {
|
|
adj += 1;
|
|
}
|
|
&transcript[adj..]
|
|
} else {
|
|
transcript
|
|
};
|
|
|
|
let model = self.loaded_model_arc()?;
|
|
let prompt = render_chat_prompt(
|
|
&model,
|
|
&[
|
|
("system", prompts::TRANSCRIPT_TITLE_SYSTEM),
|
|
("user", &format!("Transcript:\n{tail}")),
|
|
],
|
|
)?;
|
|
let raw = self.generate(
|
|
&prompt,
|
|
&GenerationConfig {
|
|
max_tokens: 24,
|
|
temperature: 0.0,
|
|
stop_sequences: vec![
|
|
"\n".to_string(),
|
|
"<|im_end|>".to_string(),
|
|
"<|im_end_of_text|>".to_string(),
|
|
],
|
|
grammar: None,
|
|
},
|
|
)?;
|
|
|
|
sanitize_title(&raw)
|
|
.ok_or_else(|| EngineError::Inference("could not derive title".into()))
|
|
}
|
|
|
|
/// Feedback-conditioned variant of `extract_tasks`. See
|
|
/// `decompose_task_with_feedback` for the `examples` semantics.
|
|
pub fn extract_tasks_with_feedback(
|
|
&self,
|
|
transcript: &str,
|
|
examples: &[prompts::FeedbackExample],
|
|
) -> Result<Vec<String>, EngineError> {
|
|
if transcript.trim().is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let model = self.loaded_model_arc()?;
|
|
let system =
|
|
prompts::build_conditioned_system_prompt(prompts::EXTRACT_TASKS_SYSTEM, examples);
|
|
let prompt = render_chat_prompt(
|
|
&model,
|
|
&[
|
|
("system", system.as_str()),
|
|
("user", &format!("Transcript:\n{transcript}")),
|
|
],
|
|
)?;
|
|
let raw = self.generate(
|
|
&prompt,
|
|
&GenerationConfig {
|
|
max_tokens: 768,
|
|
temperature: 0.0,
|
|
stop_sequences: vec!["<|im_end|>".to_string(), "<|im_end_of_text|>".to_string()],
|
|
grammar: Some(grammars::OPTIONAL_TASK_ARRAY_GRAMMAR.to_string()),
|
|
},
|
|
)?;
|
|
parse_string_array(&raw)
|
|
}
|
|
|
|
fn loaded_handles(&self) -> Result<(Arc<LlamaBackend>, Arc<LlamaModel>), EngineError> {
|
|
let guard = self.inner.lock().unwrap();
|
|
let backend = guard.backend.clone().ok_or(EngineError::NotLoaded)?;
|
|
let model = guard.model.clone().ok_or(EngineError::NotLoaded)?;
|
|
Ok((backend, model))
|
|
}
|
|
|
|
fn loaded_model_arc(&self) -> Result<Arc<LlamaModel>, EngineError> {
|
|
self.loaded_handles().map(|(_, model)| model)
|
|
}
|
|
|
|
fn build_sampler(
|
|
&self,
|
|
model: &LlamaModel,
|
|
config: &GenerationConfig,
|
|
) -> Result<LlamaSampler, EngineError> {
|
|
let mut samplers = Vec::new();
|
|
|
|
if let Some(grammar) = &config.grammar {
|
|
samplers.push(
|
|
LlamaSampler::grammar(model, grammar, "root")
|
|
.map_err(|e| EngineError::Inference(format!("grammar: {e}")))?,
|
|
);
|
|
}
|
|
|
|
if config.temperature <= f32::EPSILON {
|
|
samplers.push(LlamaSampler::greedy());
|
|
} else {
|
|
samplers.push(LlamaSampler::temp(config.temperature));
|
|
samplers.push(LlamaSampler::dist(GENERATION_SEED));
|
|
}
|
|
|
|
Ok(if samplers.len() == 1 {
|
|
samplers.remove(0)
|
|
} else {
|
|
LlamaSampler::chain_simple(samplers)
|
|
})
|
|
}
|
|
}
|
|
|
|
fn context_window_size(prompt_tokens: usize, max_tokens: u32) -> u32 {
|
|
let required = prompt_tokens
|
|
.saturating_add(max_tokens as usize)
|
|
.saturating_add(CONTEXT_RESERVE_TOKENS as usize);
|
|
DEFAULT_CONTEXT_TOKENS.max(required.min(MAX_CONTEXT_TOKENS as usize) as u32)
|
|
}
|
|
|
|
fn preflight_context_window(prompt_tokens: usize, max_tokens: u32) -> Result<u32, EngineError> {
|
|
let required = prompt_tokens
|
|
.saturating_add(max_tokens as usize)
|
|
.saturating_add(CONTEXT_RESERVE_TOKENS as usize);
|
|
if required > MAX_CONTEXT_TOKENS as usize {
|
|
let available_prompt_tokens =
|
|
MAX_CONTEXT_TOKENS.saturating_sub(max_tokens.saturating_add(CONTEXT_RESERVE_TOKENS));
|
|
return Err(EngineError::PromptTooLong {
|
|
prompt_tokens,
|
|
max_tokens,
|
|
available_prompt_tokens,
|
|
context_window: MAX_CONTEXT_TOKENS,
|
|
});
|
|
}
|
|
|
|
Ok(context_window_size(prompt_tokens, max_tokens))
|
|
}
|
|
|
|
fn first_stop_index(text: &str, stop_sequences: &[String]) -> Option<usize> {
|
|
stop_sequences
|
|
.iter()
|
|
.filter(|stop| !stop.is_empty())
|
|
.filter_map(|stop| text.find(stop))
|
|
.min()
|
|
}
|
|
|
|
fn render_chat_prompt(
|
|
model: &LlamaModel,
|
|
messages: &[(&str, &str)],
|
|
) -> Result<String, EngineError> {
|
|
let chat_messages = messages
|
|
.iter()
|
|
.map(|(role, content)| {
|
|
LlamaChatMessage::new((*role).to_string(), (*content).to_string())
|
|
.map_err(|e| EngineError::Inference(format!("chat message: {e}")))
|
|
})
|
|
.collect::<Result<Vec<_>, _>>()?;
|
|
|
|
match model.chat_template(None) {
|
|
Ok(template) => model
|
|
.apply_chat_template(&template, &chat_messages, true)
|
|
.map_err(|e| EngineError::Inference(format!("chat template apply: {e}"))),
|
|
Err(err) => {
|
|
tracing::warn!("model chat template unavailable, falling back to ChatML: {err}");
|
|
let template = LlamaChatTemplate::new("chatml")
|
|
.map_err(|e| EngineError::Inference(format!("chatml template: {e}")))?;
|
|
model
|
|
.apply_chat_template(&template, &chat_messages, true)
|
|
.map_err(|e| EngineError::Inference(format!("chatml template apply: {e}")))
|
|
}
|
|
}
|
|
}
|
|
|
|
fn parse_string_array(raw: &str) -> Result<Vec<String>, EngineError> {
|
|
let parsed = serde_json::from_str::<Vec<String>>(raw.trim())
|
|
.map_err(|e| EngineError::InvalidJson(format!("{e} in: {raw:?}")))?;
|
|
|
|
let mut seen = std::collections::HashSet::new();
|
|
let normalized = parsed
|
|
.into_iter()
|
|
.map(|item| item.trim().to_string())
|
|
.filter(|item| !item.is_empty())
|
|
.filter(|item| seen.insert(item.to_lowercase()))
|
|
.collect();
|
|
|
|
Ok(normalized)
|
|
}
|
|
|
|
/// Normalise a model-generated title into something safe to persist.
|
|
///
|
|
/// Real-world failure modes from low-temp Qwen3 runs that this catches:
|
|
/// - Surrounding quotes (smart and ASCII): `"My Title"` → `My Title`.
|
|
/// - A leading `Title:` / `TITLE:` prefix where the model echoed the
|
|
/// output schema instead of just emitting the value.
|
|
/// - Trailing terminal punctuation (`.`, `!`, `?`) — titles do not
|
|
/// take it; the prompt forbids it but the model occasionally adds
|
|
/// one anyway.
|
|
/// - Multi-line output where the first stop sequence is a newline:
|
|
/// we kept the first line via `stop_sequences`, but defensively
|
|
/// collapse internal whitespace runs here too.
|
|
/// - Length over 100 chars (cap defensively; `max_tokens: 24` already
|
|
/// bounds this in practice).
|
|
/// - Empty after stripping, or the literal `Untitled` the prompt
|
|
/// instructs the model to emit for empty/filler input — caller
|
|
/// treats `None` as "no usable title".
|
|
fn sanitize_title(raw: &str) -> Option<String> {
|
|
let mut t = raw.trim();
|
|
|
|
// First-line only — defence in depth on top of `stop_sequences`.
|
|
if let Some((first, _)) = t.split_once('\n') {
|
|
t = first.trim();
|
|
}
|
|
|
|
// Strip a leading "Title:" / "TITLE:" prefix.
|
|
let lower = t.to_ascii_lowercase();
|
|
if let Some(rest) = lower.strip_prefix("title:") {
|
|
let consumed = t.len() - rest.len();
|
|
t = t[consumed..].trim_start();
|
|
}
|
|
|
|
// Strip surrounding quotes — ASCII and the curly variants Qwen
|
|
// sometimes emits. A quote-only string like `""` collapses to empty;
|
|
// the final-empty check below treats that as "no usable title".
|
|
const QUOTES: &[char] = &['"', '\'', '\u{201C}', '\u{201D}', '\u{2018}', '\u{2019}'];
|
|
while t.starts_with(QUOTES) && t.ends_with(QUOTES) && t.chars().count() >= 2 {
|
|
let start = t.chars().next().unwrap().len_utf8();
|
|
let end = t.chars().next_back().unwrap().len_utf8();
|
|
if t.len() <= start + end {
|
|
t = "";
|
|
break;
|
|
}
|
|
t = t[start..t.len() - end].trim();
|
|
}
|
|
|
|
// Drop trailing terminal punctuation. Titles don't take it.
|
|
let trimmed_tail: String = t.trim_end_matches(['.', '!', '?']).to_string();
|
|
|
|
// Collapse internal whitespace runs to single spaces.
|
|
let collapsed: String = trimmed_tail.split_whitespace().collect::<Vec<_>>().join(" ");
|
|
|
|
// Cap at 100 chars on a UTF-8 char boundary.
|
|
let capped: String = if collapsed.chars().count() > 100 {
|
|
collapsed.chars().take(100).collect()
|
|
} else {
|
|
collapsed
|
|
};
|
|
|
|
let final_title = capped.trim();
|
|
if final_title.is_empty() || final_title.eq_ignore_ascii_case("untitled") {
|
|
return None;
|
|
}
|
|
Some(final_title.to_string())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn generate_fails_when_not_loaded() {
|
|
let engine = LlmEngine::new();
|
|
let err = engine
|
|
.generate("hello", &GenerationConfig::default())
|
|
.unwrap_err();
|
|
assert!(matches!(err, EngineError::NotLoaded));
|
|
}
|
|
|
|
#[test]
|
|
fn decompose_returns_error_when_not_loaded() {
|
|
let engine = LlmEngine::new();
|
|
assert!(!engine.is_loaded());
|
|
let result = engine.decompose_task("Write a blog post");
|
|
assert!(matches!(result, Err(EngineError::NotLoaded)));
|
|
}
|
|
|
|
#[test]
|
|
fn default_creates_unloaded_engine() {
|
|
let engine = LlmEngine::default();
|
|
assert!(!engine.is_loaded());
|
|
}
|
|
|
|
#[test]
|
|
fn engine_is_clone_and_shares_state() {
|
|
let engine = LlmEngine::new();
|
|
let clone = engine.clone();
|
|
assert!(!clone.is_loaded());
|
|
}
|
|
|
|
#[test]
|
|
fn parse_string_array_trims_and_dedupes() {
|
|
let parsed = parse_string_array(r#"[" Buy milk ", "buy milk", "Call plumber"]"#).unwrap();
|
|
assert_eq!(parsed, vec!["Buy milk", "Call plumber"]);
|
|
}
|
|
|
|
#[test]
|
|
fn first_stop_index_finds_earliest_match() {
|
|
let text = "hello<|im_end|>trailing";
|
|
let index = first_stop_index(text, &["<|im_end|>".into(), "zzz".into()]);
|
|
assert_eq!(index, Some(5));
|
|
}
|
|
|
|
#[test]
|
|
fn prompt_preflight_rejects_oversized_prompt_tokens() {
|
|
let err = preflight_context_window(7_105, 1_024).unwrap_err();
|
|
assert!(matches!(
|
|
err,
|
|
EngineError::PromptTooLong {
|
|
prompt_tokens: 7_105,
|
|
max_tokens: 1_024,
|
|
available_prompt_tokens: 7_104,
|
|
context_window: MAX_CONTEXT_TOKENS,
|
|
}
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn prompt_preflight_keeps_prompts_within_budget() {
|
|
let n_ctx = preflight_context_window(7_104, 1_024).unwrap();
|
|
assert_eq!(n_ctx, MAX_CONTEXT_TOKENS);
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_title_strips_quotes_label_and_terminal_punctuation() {
|
|
// Composite of the three real-world failure modes from low-temp
|
|
// Qwen3 runs: surrounding curly quotes, "Title:" prefix, and a
|
|
// trailing period. All three must be removed in one pass.
|
|
let cleaned = sanitize_title(" Title: \u{201C}Sales Call With ACME.\u{201D} ").unwrap();
|
|
assert_eq!(cleaned, "Sales Call With ACME");
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_title_collapses_whitespace_and_keeps_first_line() {
|
|
// Multi-line output should keep only the first line (defence on
|
|
// top of `\n` stop_sequence). Internal whitespace runs must
|
|
// collapse to a single space so a model that double-spaces
|
|
// doesn't produce a weird-looking row.
|
|
let cleaned =
|
|
sanitize_title(" Roadmap Review\nignore me\nstill ignored ").unwrap();
|
|
assert_eq!(cleaned, "Roadmap Review");
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_title_returns_none_for_untitled_or_empty() {
|
|
// The prompt instructs the model to emit "Untitled" when the
|
|
// transcript is empty/filler. Treat that as no-usable-title.
|
|
// Same for empty / whitespace-only / quote-only output.
|
|
assert!(sanitize_title("Untitled").is_none());
|
|
assert!(sanitize_title("untitled.").is_none());
|
|
assert!(sanitize_title(" ").is_none());
|
|
assert!(sanitize_title("\"\"").is_none());
|
|
}
|
|
}
|