Jake
27661c816e
rust-toolchain.toml pins to stable 1.94.1 so contributors and CI runners share the exact rustc / rustfmt / clippy versions. Without the pin, every machine surfaces a different lint set depending on its local install — six pre-existing lints showed up on 1.94.1 that 1.93-era HANDOVER reported clean. Clippy fixes (all pre-existing, not introduced by feature work): - crates/storage/src/database.rs: std::iter::repeat().take() -> repeat_n() - crates/llm/src/lib.rs (docs): "+ frontends" was parsed as a markdown bullet continuation by rustdoc, breaking doc-lazy-continuation. Reworded to "and". - crates/llm/src/lib.rs (loop): while-let-on-iterator -> for-loop. - src-tauri/src/commands/security.rs: .iter().any(|a| *a == x) -> .contains(&x). - src-tauri/src/lib.rs: io::Error::new(Other, e) -> io::Error::other(e). - src-tauri/src/tauri_app_data_migration.rs: drop function-tail `return`s inside cfg blocks; each platform's block now ends with a tail expression. cargo fmt sweep across the workspace. Mechanical layout-only changes; no semantics affected. Workspace gates after this commit: - cargo fmt --check: clean - cargo clippy --workspace --all-targets -- -D warnings: clean - cargo test --workspace: 405/0 (will become 409/0 with Phase A.1+A.2)
939 lines
34 KiB
Rust
939 lines
34 KiB
Rust
use std::num::NonZeroU32;
|
|
use std::path::Path;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
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 lumotia_core::tuning::{inference_thread_count, Workload};
|
|
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};
|
|
|
|
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("Another LLM load is already in flight; refusing to start a parallel load.")]
|
|
AlreadyLoading,
|
|
#[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>>,
|
|
/// Flag held for the duration of a model load. The std::sync::Mutex
|
|
/// covers cheap state mutations (~microseconds); the multi-second
|
|
/// `LlamaModel::load_from_file` call now runs *outside* the mutex,
|
|
/// so polls like `is_loaded()` / `loaded_model_id()` (called from
|
|
/// sync Tauri handlers without `spawn_blocking`) don't park tokio
|
|
/// worker threads on a slow C++ FFI call. This Atomic also doubles
|
|
/// as a TOCTOU guard: two concurrent `load_model` invocations on
|
|
/// the same engine will not both reach the heavy load — the second
|
|
/// returns `EngineError::AlreadyLoading`.
|
|
loading: Arc<AtomicBool>,
|
|
}
|
|
|
|
/// RAII guard that clears the `loading` flag on drop, including on
|
|
/// panic / early-return. Prevents the engine getting stuck in a
|
|
/// "permanently loading" state if a load fails midway.
|
|
struct LoadingGuard {
|
|
flag: Arc<AtomicBool>,
|
|
}
|
|
|
|
impl Drop for LoadingGuard {
|
|
fn drop(&mut self) {
|
|
self.flag.store(false, Ordering::Release);
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
// instrument: the load is multi-second (`LlamaBackend::init` +
|
|
// mmap + GPU layer init). Tagging events with `model_id` and
|
|
// `use_gpu` lets the operator separate the GPU sequential-guard
|
|
// logs and llama-backend init lines from the LLM transcription
|
|
// pipeline by structured field rather than by adjacency.
|
|
#[tracing::instrument(skip_all, fields(model_id = %model_id.as_str(), use_gpu = use_gpu))]
|
|
pub fn load_model(
|
|
&self,
|
|
model_id: LlmModelId,
|
|
model_path: &Path,
|
|
use_gpu: bool,
|
|
) -> Result<(), EngineError> {
|
|
self.load_model_with(model_id, model_path, use_gpu, |backend, path, params| {
|
|
LlamaModel::load_from_file(backend, path, params)
|
|
.map_err(|e| EngineError::LoadFailed(format!("model load: {e}")))
|
|
})
|
|
}
|
|
|
|
/// Core load implementation with a swappable file-loader closure.
|
|
/// Production callers use `load_model`, which delegates here with
|
|
/// the real `LlamaModel::load_from_file`. Tests inject a sleepy /
|
|
/// counting closure to exercise the locking discipline without
|
|
/// pulling a real GGUF off disk.
|
|
///
|
|
/// Locking discipline (the whole point of this function):
|
|
/// 1. Take the mutex briefly to compare against the currently
|
|
/// loaded triple — if it matches, return early. No-op fast path.
|
|
/// 2. CAS the `loading` flag from false → true. If another load is
|
|
/// already in flight, refuse with `AlreadyLoading` rather than
|
|
/// starting a parallel one. A `LoadingGuard` ensures the flag
|
|
/// is cleared on every exit path including panic.
|
|
/// 3. Take the mutex briefly to drop the OLD model Arc (frees its
|
|
/// VRAM via `llama_free_model`) before the new load begins.
|
|
/// The backend Arc is preserved — `LlamaBackend::init()` is a
|
|
/// one-shot per process (an `AtomicBool` in llama-cpp-2 enforces
|
|
/// `BackendAlreadyInitialized` on a second call), so we must
|
|
/// never drop the backend while the process keeps running.
|
|
/// Note: `is_loaded()` reports false during the swap window —
|
|
/// that is the correct semantics. Callers wanting "model X is
|
|
/// loaded" must check `loaded_model_id()` against their target.
|
|
/// 4. Initialise the backend if absent (first-ever load only) and
|
|
/// run the slow `load_from_file` call — both OUTSIDE the mutex.
|
|
/// 5. Take the mutex briefly to install the new backend (if just
|
|
/// initialised) and the new model Arc.
|
|
fn load_model_with<F>(
|
|
&self,
|
|
model_id: LlmModelId,
|
|
model_path: &Path,
|
|
use_gpu: bool,
|
|
loader: F,
|
|
) -> Result<(), EngineError>
|
|
where
|
|
F: FnOnce(&LlamaBackend, &Path, &LlamaModelParams) -> Result<LlamaModel, EngineError>,
|
|
{
|
|
// Step 1: short crit section — already-loaded fast path.
|
|
{
|
|
let 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(());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Step 2: claim the loading slot. Refuse if a parallel load is
|
|
// already mid-flight rather than starting a second slow load
|
|
// and silently overwriting the first.
|
|
if self
|
|
.loading
|
|
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
|
.is_err()
|
|
{
|
|
return Err(EngineError::AlreadyLoading);
|
|
}
|
|
let _loading_guard = LoadingGuard {
|
|
flag: Arc::clone(&self.loading),
|
|
};
|
|
|
|
// Step 3: short crit section — drop the OLD model so its VRAM is
|
|
// released BEFORE we allocate the new one. Without this, a swap
|
|
// briefly holds two models resident (Lifecycle-1: an
|
|
// ~17 GB Q4 27B swap on a 24 GB card OOMs even though either
|
|
// model fits alone). Keep the backend Arc — see locking notes.
|
|
let existing_backend = {
|
|
let mut guard = self.inner.lock().unwrap();
|
|
guard.model = None;
|
|
guard.loaded = None;
|
|
guard.backend.clone()
|
|
};
|
|
|
|
// Step 4: heavy work OUTSIDE the mutex. `is_loaded()` and
|
|
// `loaded_model_id()` can be polled freely here without parking
|
|
// tokio worker threads.
|
|
let backend = match existing_backend {
|
|
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 = loader(&backend, model_path, ¶ms)?;
|
|
|
|
// Step 5: short crit section — install the new state.
|
|
{
|
|
let mut guard = self.inner.lock().unwrap();
|
|
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,
|
|
});
|
|
}
|
|
// `_loading_guard` drops here and clears the flag.
|
|
Ok(())
|
|
}
|
|
|
|
pub fn unload(&self) -> Result<(), EngineError> {
|
|
let mut guard = self.inner.lock().unwrap();
|
|
guard.model = None;
|
|
// Backend is process-singleton (llama-cpp-2 enforces this via
|
|
// `LLAMA_BACKEND_INITIALIZED`). Dropping the Arc here would call
|
|
// `llama_backend_free` and a subsequent `init` would succeed, but
|
|
// we keep it resident to avoid the init/free churn on every
|
|
// load/unload cycle.
|
|
guard.loaded = None;
|
|
Ok(())
|
|
}
|
|
|
|
/// True iff a model load is currently in flight. Exposed for tests
|
|
/// and frontends that want to render a "loading…" state without
|
|
/// polling `is_loaded()` (which returns false during a swap).
|
|
pub fn is_loading(&self) -> bool {
|
|
self.loading.load(Ordering::Acquire)
|
|
}
|
|
|
|
/// Test-only harness: runs `op` while holding the same locking
|
|
/// discipline as `load_model_with` (loading flag claimed, model
|
|
/// state cleared, slow op runs OUTSIDE the inner mutex, new state
|
|
/// installed at the end). Used by the regression test to verify
|
|
/// that `is_loaded()` / `loaded_model_id()` don't block on the
|
|
/// slow section. Not part of the public API.
|
|
#[cfg(test)]
|
|
pub(crate) fn __test_run_with_lock_discipline<F>(&self, op: F) -> Result<(), EngineError>
|
|
where
|
|
F: FnOnce(),
|
|
{
|
|
if self
|
|
.loading
|
|
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
|
|
.is_err()
|
|
{
|
|
return Err(EngineError::AlreadyLoading);
|
|
}
|
|
let _loading_guard = LoadingGuard {
|
|
flag: Arc::clone(&self.loading),
|
|
};
|
|
{
|
|
let mut guard = self.inner.lock().unwrap();
|
|
guard.model = None;
|
|
guard.loaded = None;
|
|
}
|
|
op();
|
|
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 use_gpu = self.loaded_model().map(|s| s.use_gpu).unwrap_or(false);
|
|
let gpu_layers = if use_gpu { u32::MAX } else { 0u32 };
|
|
// Trivially true today (gpu_layers = u32::MAX when use_gpu), but the
|
|
// explicit comparison documents intent. True residency observability
|
|
// (parsing llama.cpp's "offloaded N/M layers" log) is tracked as a
|
|
// follow-up in docs/superpowers/specs/2026-05-09-battery-gpu-aware-
|
|
// thread-tuning-design.md (§ Out of scope).
|
|
let gpu_offloaded = use_gpu && gpu_layers >= model.n_layer();
|
|
let thread_count =
|
|
i32::try_from(inference_thread_count(Workload::Llm, gpu_offloaded)).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 config.grammar.is_none() && json_envelope_complete(&generated) {
|
|
break;
|
|
}
|
|
|
|
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 as JSON. 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;
|
|
/// the parsed intent is re-validated with `is_valid_intent` and
|
|
/// invalid JSON bubbles as `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: None,
|
|
},
|
|
)?;
|
|
|
|
let tags: prompts::ContentTags = parse_json_payload(&raw)?;
|
|
if !prompts::is_valid_intent(&tags.intent) {
|
|
return Err(EngineError::InvalidJson(format!(
|
|
"intent out of closed set: {}",
|
|
tags.intent,
|
|
)));
|
|
}
|
|
Ok(tags)
|
|
}
|
|
|
|
/// 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 json_envelope_complete(text: &str) -> bool {
|
|
extract_json_envelope(text) == Some(text.trim())
|
|
}
|
|
|
|
fn extract_json_envelope(text: &str) -> Option<&str> {
|
|
let start = text
|
|
.char_indices()
|
|
.find_map(|(idx, ch)| (ch == '{' || ch == '[').then_some(idx))?;
|
|
let mut chars = text[start..].char_indices();
|
|
let (_, first) = chars.next()?;
|
|
|
|
let mut stack = vec![match first {
|
|
'{' => '}',
|
|
'[' => ']',
|
|
_ => unreachable!(),
|
|
}];
|
|
let mut in_string = false;
|
|
let mut escaped = false;
|
|
|
|
for (offset, ch) in chars {
|
|
if in_string {
|
|
if escaped {
|
|
escaped = false;
|
|
} else if ch == '\\' {
|
|
escaped = true;
|
|
} else if ch == '"' {
|
|
in_string = false;
|
|
}
|
|
continue;
|
|
}
|
|
|
|
match ch {
|
|
'"' => in_string = true,
|
|
'{' => stack.push('}'),
|
|
'[' => stack.push(']'),
|
|
'}' | ']' => {
|
|
if stack.pop() != Some(ch) {
|
|
return None;
|
|
}
|
|
if stack.is_empty() {
|
|
let end = start + offset + ch.len_utf8();
|
|
return Some(&text[start..end]);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn parse_json_payload<T: for<'de> Deserialize<'de>>(raw: &str) -> Result<T, EngineError> {
|
|
let payload = extract_json_envelope(raw).unwrap_or(raw.trim());
|
|
serde_json::from_str(payload).map_err(|e| EngineError::InvalidJson(format!("{e}: raw={raw:?}")))
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
#[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 json_envelope_complete_detects_finished_object() {
|
|
assert!(json_envelope_complete(
|
|
r#"{"topic":"meeting","intent":"planning"}"#
|
|
));
|
|
}
|
|
|
|
#[test]
|
|
fn json_envelope_complete_detects_finished_array() {
|
|
assert!(json_envelope_complete(r#"["Call plumber","Buy milk"]"#));
|
|
}
|
|
|
|
#[test]
|
|
fn json_envelope_complete_ignores_braces_inside_strings() {
|
|
assert!(!json_envelope_complete(r#"{"topic":"literal } brace""#));
|
|
}
|
|
|
|
#[test]
|
|
fn json_envelope_complete_rejects_prefixes_and_trailing_text() {
|
|
assert!(!json_envelope_complete(r#"{"topic":"meeting""#));
|
|
assert!(!json_envelope_complete(r#"{"topic":"meeting"} extra"#));
|
|
}
|
|
|
|
#[test]
|
|
fn extract_json_envelope_skips_qwen_thinking_prefix() {
|
|
let raw =
|
|
"<think>\n\n</think>\n\n{\"topic\":\"grant-application\",\"intent\":\"planning\"}";
|
|
assert_eq!(
|
|
extract_json_envelope(raw),
|
|
Some("{\"topic\":\"grant-application\",\"intent\":\"planning\"}"),
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn extract_json_envelope_handles_arrays_and_trailing_stop_text() {
|
|
assert_eq!(
|
|
extract_json_envelope("prefix [\"Call plumber\",\"Buy milk\"]<|im_end|>"),
|
|
Some("[\"Call plumber\",\"Buy milk\"]"),
|
|
);
|
|
}
|
|
|
|
#[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);
|
|
}
|
|
|
|
/// Race-3 regression. The inner `std::sync::Mutex` MUST NOT be held
|
|
/// across the slow `LlamaModel::load_from_file` call: sync Tauri
|
|
/// command handlers like `get_llm_status`, `check_llm_model`,
|
|
/// `delete_llm_model`, and `test_llm_model` call `is_loaded()` /
|
|
/// `loaded_model_id()` from tokio worker threads without
|
|
/// `spawn_blocking`. If the lock is held for the duration of a
|
|
/// 5-15 s load, parallel status polls from the frontend park the
|
|
/// tokio executor and the whole UI deadlocks.
|
|
///
|
|
/// Pre-fix this test FAILS — both probes time out because the load
|
|
/// holds the mutex. Post-fix it PASSES — probes return in ≤50 ms.
|
|
#[test]
|
|
fn is_loaded_does_not_block_on_slow_load() {
|
|
use std::sync::mpsc;
|
|
use std::thread;
|
|
use std::time::{Duration, Instant};
|
|
|
|
let engine = LlmEngine::new();
|
|
let load_started = Arc::new(std::sync::Barrier::new(2));
|
|
let release_load = Arc::new(std::sync::Barrier::new(2));
|
|
|
|
let engine_for_loader = engine.clone();
|
|
let load_started_for_loader = Arc::clone(&load_started);
|
|
let release_load_for_loader = Arc::clone(&release_load);
|
|
|
|
let loader_handle = thread::spawn(move || {
|
|
engine_for_loader
|
|
.__test_run_with_lock_discipline(|| {
|
|
// Signal the probe thread that the load is mid-flight
|
|
// (loading flag claimed, inner mutex released).
|
|
load_started_for_loader.wait();
|
|
// Wait until the probe thread says it's done so the
|
|
// load's "duration" is bounded by the probes.
|
|
release_load_for_loader.wait();
|
|
})
|
|
.unwrap();
|
|
});
|
|
|
|
// Wait until the loader is inside its slow section.
|
|
load_started.wait();
|
|
|
|
// Now probe `is_loaded()` and `loaded_model_id()` from this
|
|
// thread. They MUST return without contending on the lock.
|
|
let probe_deadline = Duration::from_millis(50);
|
|
let (tx, rx) = mpsc::channel();
|
|
let engine_for_probe = engine.clone();
|
|
let probe_handle = thread::spawn(move || {
|
|
let start = Instant::now();
|
|
let loaded = engine_for_probe.is_loaded();
|
|
let id = engine_for_probe.loaded_model_id();
|
|
let loading = engine_for_probe.is_loading();
|
|
let elapsed = start.elapsed();
|
|
tx.send((loaded, id, loading, elapsed)).unwrap();
|
|
});
|
|
|
|
let result = rx
|
|
.recv_timeout(probe_deadline)
|
|
.expect("is_loaded / loaded_model_id probe must return within 50 ms");
|
|
let (loaded, id, loading, elapsed) = result;
|
|
|
|
assert!(
|
|
!loaded,
|
|
"is_loaded() should report false while a load is in flight"
|
|
);
|
|
assert_eq!(id, None, "loaded_model_id() should be None mid-load");
|
|
assert!(loading, "is_loading() should report true mid-load");
|
|
assert!(
|
|
elapsed < probe_deadline,
|
|
"probe took {elapsed:?}, expected < {probe_deadline:?}"
|
|
);
|
|
|
|
probe_handle.join().unwrap();
|
|
release_load.wait();
|
|
loader_handle.join().unwrap();
|
|
|
|
// After the load completes the flag clears.
|
|
assert!(!engine.is_loading());
|
|
}
|
|
|
|
/// Race-3 / Race-4 — concurrent load attempts must not both reach
|
|
/// the heavy work. The second caller should be told `AlreadyLoading`
|
|
/// rather than starting a parallel load that silently overwrites
|
|
/// the first.
|
|
#[test]
|
|
fn second_concurrent_load_is_refused() {
|
|
use std::thread;
|
|
|
|
let engine = LlmEngine::new();
|
|
let load_started = Arc::new(std::sync::Barrier::new(2));
|
|
let release_load = Arc::new(std::sync::Barrier::new(2));
|
|
|
|
let engine_for_loader = engine.clone();
|
|
let load_started_for_loader = Arc::clone(&load_started);
|
|
let release_load_for_loader = Arc::clone(&release_load);
|
|
|
|
let loader_handle = thread::spawn(move || {
|
|
engine_for_loader
|
|
.__test_run_with_lock_discipline(|| {
|
|
load_started_for_loader.wait();
|
|
release_load_for_loader.wait();
|
|
})
|
|
.unwrap();
|
|
});
|
|
|
|
load_started.wait();
|
|
|
|
// Second concurrent attempt MUST be refused, not parallel-load.
|
|
let second = engine.__test_run_with_lock_discipline(|| {
|
|
panic!("second concurrent load should never reach its op");
|
|
});
|
|
assert!(matches!(second, Err(EngineError::AlreadyLoading)));
|
|
|
|
release_load.wait();
|
|
loader_handle.join().unwrap();
|
|
|
|
// After the first load completes, a fresh attempt is allowed.
|
|
assert!(engine.__test_run_with_lock_discipline(|| {}).is_ok());
|
|
}
|
|
}
|