feat(llm): wire Phase 3 local LLM runtime via llama-cpp-2
kon-llm now owns a real LlamaBackend + LlamaModel, with three Qwen3 tiers (1.7B Q4, 4B-Instruct-2507 Q4, 14B Q5) selectable per hardware. Downloads are resumable with SHA-256 verification and stored under ~/.kon/models/llm. Engine exposes three high-level surfaces — all greedy/temp-0, GBNF-constrained where output shape matters: - cleanup_text (prompt-injection-hardened system prompt; profile terms appended as "preserve these spellings" suffix) - decompose_task (3–7 micro-steps, constrained JSON array) - extract_tasks (optional-array; empty when no explicit commitments) post_process_segments now takes an Option<&LlmEngine> and, when loaded and format_mode != Raw, joins segments → cleanup → replaces segments with the cleaned text (first segment span). Rule-based path still runs first; LLM errors log and keep rule-based output. Tauri commands: recommend_llm_tier, check_llm_model, download_llm_model, load_llm_model, unload_llm_model, delete_llm_model, get_llm_status, cleanup_transcript_text_cmd, extract_tasks_from_transcript_cmd, decompose_and_store (LLM-backed subtasks). Settings: AI tier toggle (off / cleanup / tasks), model picker with downloaded/loaded status, download progress events via kon:llm-download-progress. Dictation: ensureLlmModelLoaded on mount, cleanupTranscriptIfEnabled after stop when tier != off and format_mode != Raw, LLM task extraction when tier=tasks (regex fallback on failure). Interim: both llama-cpp-sys-2 and whisper-rs-sys statically link their own ggml, so src-tauri/build.rs emits -Wl,--allow-multiple-definition on Linux. Replace with a system-ggml shared-lib setup as a follow-up. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -4,3 +4,18 @@ version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
dirs = "6"
|
||||
encoding_rs = "0.8"
|
||||
futures-util = "0.3"
|
||||
llama-cpp-2 = { version = "0.1.144", default-features = false, features = ["openmp", "vulkan"] }
|
||||
num_cpus = "1"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "stream"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
sha2 = "0.10"
|
||||
thiserror = "2"
|
||||
tokio = { version = "1", features = ["fs", "io-util", "macros", "net", "rt-multi-thread", "sync", "time"] }
|
||||
tracing = "0.1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
|
||||
24
crates/llm/src/grammars.rs
Normal file
24
crates/llm/src/grammars.rs
Normal file
@@ -0,0 +1,24 @@
|
||||
pub const TASK_ARRAY_GRAMMAR: &str = r#"
|
||||
root ::= "[" ws string ws "," ws string ws "," ws string rest3 ws "]"
|
||||
rest3 ::= "" | "," ws string rest4
|
||||
rest4 ::= "" | "," ws string rest5
|
||||
rest5 ::= "" | "," ws string rest6
|
||||
rest6 ::= "" | "," ws string
|
||||
string ::= "\"" chars "\"" ws
|
||||
chars ::= "" | char chars
|
||||
char ::= [^"\\\n\r] | "\\" escape
|
||||
escape ::= ["\\/bfnrt] | "u" hex hex hex hex
|
||||
hex ::= [0-9a-fA-F]
|
||||
ws ::= ([ \t\n\r] ws)?
|
||||
"#;
|
||||
|
||||
pub const OPTIONAL_TASK_ARRAY_GRAMMAR: &str = r#"
|
||||
root ::= "[" ws "]" | "[" ws string tail ws "]"
|
||||
tail ::= "" | "," ws string tail
|
||||
string ::= "\"" chars "\"" ws
|
||||
chars ::= "" | char chars
|
||||
char ::= [^"\\\n\r] | "\\" escape
|
||||
escape ::= ["\\/bfnrt] | "u" hex hex hex hex
|
||||
hex ::= [0-9a-fA-F]
|
||||
ws ::= ([ \t\n\r] ws)?
|
||||
"#;
|
||||
@@ -1,58 +1,395 @@
|
||||
use std::num::NonZeroU32;
|
||||
use std::path::Path;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
struct LlmState {
|
||||
loaded: bool,
|
||||
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 model_manager::{recommend_tier, LlmModelId, LlmModelInfo};
|
||||
|
||||
const DEFAULT_CONTEXT_TOKENS: u32 = 4096;
|
||||
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("inference failed: {0}")]
|
||||
Inference(String),
|
||||
#[error("model output not valid JSON: {0}")]
|
||||
InvalidJson(String),
|
||||
}
|
||||
|
||||
/// Shared handle to the LLM engine. Cheap to clone (Arc).
|
||||
/// Phase 3 will replace the stub body with a real llama-cpp-2 model.
|
||||
#[derive(Clone)]
|
||||
#[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 {
|
||||
state: Arc<Mutex<LlmState>>,
|
||||
inner: Arc<Mutex<LlmState>>,
|
||||
}
|
||||
|
||||
impl LlmEngine {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: Arc::new(Mutex::new(LlmState { loaded: false })),
|
||||
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.state.lock().unwrap().loaded
|
||||
self.inner.lock().unwrap().model.is_some()
|
||||
}
|
||||
|
||||
/// Break a task description into 3-7 physical micro-steps.
|
||||
/// Returns Err if no model is loaded — the caller surfaces this to the UI.
|
||||
pub fn decompose_task(&self, _task_text: &str) -> Result<Vec<String>, String> {
|
||||
if !self.is_loaded() {
|
||||
return Err("Download an AI model in Settings to break down tasks.".to_string());
|
||||
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());
|
||||
}
|
||||
// Phase 3: call llama-cpp-2 with GBNF-constrained prompt here.
|
||||
Err("LLM not yet wired.".to_string())
|
||||
|
||||
let n_ctx = context_window_size(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> {
|
||||
let model = self.loaded_model_arc()?;
|
||||
let prompt = render_chat_prompt(
|
||||
&model,
|
||||
&[
|
||||
("system", prompts::DECOMPOSE_TASK_SYSTEM),
|
||||
("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> {
|
||||
if transcript.trim().is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let model = self.loaded_model_arc()?;
|
||||
let prompt = render_chat_prompt(
|
||||
&model,
|
||||
&[
|
||||
("system", prompts::EXTRACT_TASKS_SYSTEM),
|
||||
("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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for LlmEngine {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
fn context_window_size(prompt_tokens: usize, max_tokens: u32) -> u32 {
|
||||
let required = prompt_tokens
|
||||
.saturating_add(max_tokens as usize)
|
||||
.saturating_add(64);
|
||||
DEFAULT_CONTEXT_TOKENS.max(required.min(8192) as u32)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
#[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!(result.is_err());
|
||||
assert!(
|
||||
result.unwrap_err().contains("Download an AI model"),
|
||||
"error message should tell user to download a model"
|
||||
);
|
||||
assert!(matches!(result, Err(EngineError::NotLoaded)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -65,7 +402,19 @@ mod tests {
|
||||
fn engine_is_clone_and_shares_state() {
|
||||
let engine = LlmEngine::new();
|
||||
let clone = engine.clone();
|
||||
// Both point to the same Arc — neither is loaded
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
||||
447
crates/llm/src/model_manager.rs
Normal file
447
crates/llm/src/model_manager.rs
Normal file
@@ -0,0 +1,447 @@
|
||||
use std::fmt;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::str::FromStr;
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
|
||||
#[allow(non_camel_case_types)]
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum LlmModelId {
|
||||
#[serde(rename = "qwen3_1_7b")]
|
||||
Qwen3_1_7B_Q4,
|
||||
#[serde(rename = "qwen3_4b_instruct_2507")]
|
||||
Qwen3_4BInstruct2507Q4,
|
||||
#[serde(rename = "qwen3_14b")]
|
||||
Qwen3_14BQ5,
|
||||
}
|
||||
|
||||
impl LlmModelId {
|
||||
pub fn default_tier() -> Self {
|
||||
Self::Qwen3_4BInstruct2507Q4
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => "qwen3_1_7b",
|
||||
Self::Qwen3_4BInstruct2507Q4 => "qwen3_4b_instruct_2507",
|
||||
Self::Qwen3_14BQ5 => "qwen3_14b",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn display_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => "Qwen3 1.7B",
|
||||
Self::Qwen3_4BInstruct2507Q4 => "Qwen3 4B Instruct 2507",
|
||||
Self::Qwen3_14BQ5 => "Qwen3 14B",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn file_name(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => "Qwen3-1.7B-Q4_K_M.gguf",
|
||||
Self::Qwen3_4BInstruct2507Q4 => "Qwen3-4B-Instruct-2507-Q4_K_M.gguf",
|
||||
Self::Qwen3_14BQ5 => "Qwen3-14B-Q5_K_M.gguf",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn size_bytes(&self) -> u64 {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => 1_107_409_472,
|
||||
Self::Qwen3_4BInstruct2507Q4 => 2_497_281_120,
|
||||
Self::Qwen3_14BQ5 => 10_514_570_624,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn minimum_ram_bytes(&self) -> u64 {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => 8 * 1024_u64.pow(3),
|
||||
Self::Qwen3_4BInstruct2507Q4 => 16 * 1024_u64.pow(3),
|
||||
Self::Qwen3_14BQ5 => 32 * 1024_u64.pow(3),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn recommended_vram_bytes(&self) -> Option<u64> {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => None,
|
||||
Self::Qwen3_4BInstruct2507Q4 => Some(8 * 1024_u64.pow(3)),
|
||||
Self::Qwen3_14BQ5 => Some(16 * 1024_u64.pow(3)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn description(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => "Low tier for 8 GB RAM and CPU-heavy machines.",
|
||||
Self::Qwen3_4BInstruct2507Q4 => {
|
||||
"Default tier for cleanup and task extraction on 16 GB systems."
|
||||
}
|
||||
Self::Qwen3_14BQ5 => "High tier for 32 GB+ RAM and larger GPUs.",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn hf_url(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => {
|
||||
"https://huggingface.co/unsloth/Qwen3-1.7B-GGUF/resolve/d7f544eead698dbd1f15126ef60b45a1e1933222/Qwen3-1.7B-Q4_K_M.gguf"
|
||||
}
|
||||
Self::Qwen3_4BInstruct2507Q4 => {
|
||||
"https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/resolve/a06e946bb6b655725eafa393f4a9745d460374c9/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"
|
||||
}
|
||||
Self::Qwen3_14BQ5 => {
|
||||
"https://huggingface.co/unsloth/Qwen3-14B-GGUF/resolve/a04a82c4739b3ef5fa6da7d10261db2c67dd1985/Qwen3-14B-Q5_K_M.gguf"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sha256(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Qwen3_1_7B_Q4 => {
|
||||
"de942b0819216caa3bfe487180dd1bb37398fa1c98cb42bb0bbac7ab7d6e8a12"
|
||||
}
|
||||
Self::Qwen3_4BInstruct2507Q4 => {
|
||||
"bf52d44a54b81d44219833556849529ee96f09da673a38783dddc2e2eaf17881"
|
||||
}
|
||||
Self::Qwen3_14BQ5 => "6f87abc471bd509ad46aca4284b3cfa926d8114bc491bb0a7a3a7f74c16ef95b",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for LlmModelId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for LlmModelId {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(value: &str) -> Result<Self, Self::Err> {
|
||||
match value {
|
||||
"qwen3_1_7b" => Ok(Self::Qwen3_1_7B_Q4),
|
||||
"qwen3_4b_instruct_2507" => Ok(Self::Qwen3_4BInstruct2507Q4),
|
||||
"qwen3_14b" => Ok(Self::Qwen3_14BQ5),
|
||||
other => Err(format!("Unknown LLM model id: {other}")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LlmModelInfo {
|
||||
pub id: String,
|
||||
pub display_name: &'static str,
|
||||
pub file_name: &'static str,
|
||||
pub size_bytes: u64,
|
||||
pub description: &'static str,
|
||||
pub minimum_ram_bytes: u64,
|
||||
pub recommended_vram_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DownloadError {
|
||||
#[error("http error: {0}")]
|
||||
Http(String),
|
||||
#[error("io error: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("sha256 mismatch: expected {expected}, got {actual}")]
|
||||
ShaMismatch { expected: String, actual: String },
|
||||
#[error("resume failed: server does not support range requests")]
|
||||
ResumeUnsupported,
|
||||
}
|
||||
|
||||
const ALL_MODELS: &[LlmModelId] = &[
|
||||
LlmModelId::Qwen3_1_7B_Q4,
|
||||
LlmModelId::Qwen3_4BInstruct2507Q4,
|
||||
LlmModelId::Qwen3_14BQ5,
|
||||
];
|
||||
|
||||
pub fn all_models() -> &'static [LlmModelId] {
|
||||
ALL_MODELS
|
||||
}
|
||||
|
||||
pub fn model_info(id: LlmModelId) -> LlmModelInfo {
|
||||
LlmModelInfo {
|
||||
id: id.as_str().to_string(),
|
||||
display_name: id.display_name(),
|
||||
file_name: id.file_name(),
|
||||
size_bytes: id.size_bytes(),
|
||||
description: id.description(),
|
||||
minimum_ram_bytes: id.minimum_ram_bytes(),
|
||||
recommended_vram_bytes: id.recommended_vram_bytes(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn recommend_tier(total_ram_bytes: u64, total_vram_bytes: Option<u64>) -> LlmModelId {
|
||||
if total_vram_bytes.unwrap_or(0) >= 16 * 1024_u64.pow(3)
|
||||
&& total_ram_bytes >= 32 * 1024_u64.pow(3)
|
||||
{
|
||||
LlmModelId::Qwen3_14BQ5
|
||||
} else if total_vram_bytes.unwrap_or(0) >= 8 * 1024_u64.pow(3)
|
||||
|| total_ram_bytes >= 16 * 1024_u64.pow(3)
|
||||
{
|
||||
LlmModelId::Qwen3_4BInstruct2507Q4
|
||||
} else {
|
||||
LlmModelId::Qwen3_1_7B_Q4
|
||||
}
|
||||
}
|
||||
|
||||
pub fn model_dir() -> PathBuf {
|
||||
if cfg!(target_os = "windows") {
|
||||
std::env::var("LOCALAPPDATA")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| PathBuf::from("."))
|
||||
.join("kon")
|
||||
.join("models")
|
||||
.join("llm")
|
||||
} else {
|
||||
dirs::home_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join(".kon")
|
||||
.join("models")
|
||||
.join("llm")
|
||||
}
|
||||
}
|
||||
|
||||
pub fn model_path(id: LlmModelId) -> PathBuf {
|
||||
model_dir().join(id.file_name())
|
||||
}
|
||||
|
||||
pub fn partial_download_path(id: LlmModelId) -> PathBuf {
|
||||
model_path(id).with_extension("gguf.part")
|
||||
}
|
||||
|
||||
pub fn is_downloaded(id: LlmModelId) -> bool {
|
||||
model_path(id).exists()
|
||||
}
|
||||
|
||||
pub fn delete_model(id: LlmModelId) -> io::Result<()> {
|
||||
let final_path = model_path(id);
|
||||
let partial_path = partial_download_path(id);
|
||||
|
||||
if final_path.exists() {
|
||||
std::fs::remove_file(final_path)?;
|
||||
}
|
||||
if partial_path.exists() {
|
||||
std::fs::remove_file(partial_path)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn download_model<F>(id: LlmModelId, on_progress: F) -> Result<(), DownloadError>
|
||||
where
|
||||
F: FnMut(u64, u64) + Send + 'static,
|
||||
{
|
||||
let dest = model_path(id);
|
||||
tokio::fs::create_dir_all(model_dir()).await?;
|
||||
|
||||
if dest.exists() {
|
||||
let actual = sha256_file(&dest).await?;
|
||||
if actual == id.sha256() {
|
||||
return Ok(());
|
||||
}
|
||||
tokio::fs::remove_file(&dest).await?;
|
||||
}
|
||||
|
||||
download_impl(id.hf_url(), id.sha256(), &dest, on_progress).await
|
||||
}
|
||||
|
||||
async fn sha256_file(path: &Path) -> Result<String, io::Error> {
|
||||
let mut hasher = Sha256::new();
|
||||
let mut file = tokio::fs::File::open(path).await?;
|
||||
let mut buffer = [0u8; 8192];
|
||||
|
||||
loop {
|
||||
let count = file.read(&mut buffer).await?;
|
||||
if count == 0 {
|
||||
break;
|
||||
}
|
||||
hasher.update(&buffer[..count]);
|
||||
}
|
||||
|
||||
Ok(format!("{:x}", hasher.finalize()))
|
||||
}
|
||||
|
||||
async fn download_impl<F>(
|
||||
url: &str,
|
||||
expected_sha: &str,
|
||||
dest: &Path,
|
||||
mut on_progress: F,
|
||||
) -> Result<(), DownloadError>
|
||||
where
|
||||
F: FnMut(u64, u64) + Send + 'static,
|
||||
{
|
||||
let tmp = dest.with_extension("gguf.part");
|
||||
let resume_from = tokio::fs::metadata(&tmp)
|
||||
.await
|
||||
.ok()
|
||||
.map(|m| m.len())
|
||||
.unwrap_or(0);
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.user_agent("kon/0.1.0")
|
||||
.connect_timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.map_err(|e| DownloadError::Http(e.to_string()))?;
|
||||
|
||||
let mut request = client.get(url);
|
||||
if resume_from > 0 {
|
||||
request = request.header(reqwest::header::RANGE, format!("bytes={resume_from}-"));
|
||||
}
|
||||
|
||||
let response = request
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| DownloadError::Http(e.to_string()))?;
|
||||
if resume_from > 0 && response.status() != reqwest::StatusCode::PARTIAL_CONTENT {
|
||||
return Err(DownloadError::ResumeUnsupported);
|
||||
}
|
||||
if !response.status().is_success() && response.status() != reqwest::StatusCode::PARTIAL_CONTENT
|
||||
{
|
||||
return Err(DownloadError::Http(format!("status {}", response.status())));
|
||||
}
|
||||
|
||||
let total = if resume_from > 0 {
|
||||
response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_RANGE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.rsplit('/').next())
|
||||
.and_then(|value| value.parse::<u64>().ok())
|
||||
.unwrap_or_else(|| response.content_length().unwrap_or(0) + resume_from)
|
||||
} else {
|
||||
response.content_length().unwrap_or(0)
|
||||
};
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
if resume_from > 0 {
|
||||
let mut partial = tokio::fs::File::open(&tmp).await?;
|
||||
let mut buffer = [0u8; 8192];
|
||||
loop {
|
||||
let count = partial.read(&mut buffer).await?;
|
||||
if count == 0 {
|
||||
break;
|
||||
}
|
||||
hasher.update(&buffer[..count]);
|
||||
}
|
||||
}
|
||||
|
||||
let mut output = tokio::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(&tmp)
|
||||
.await?;
|
||||
|
||||
let mut downloaded = resume_from;
|
||||
let mut stream = response.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| DownloadError::Http(e.to_string()))?;
|
||||
output.write_all(&chunk).await?;
|
||||
hasher.update(&chunk);
|
||||
downloaded += chunk.len() as u64;
|
||||
on_progress(downloaded, total);
|
||||
}
|
||||
output.flush().await?;
|
||||
drop(output);
|
||||
|
||||
let actual = format!("{:x}", hasher.finalize());
|
||||
if actual != expected_sha {
|
||||
tokio::fs::remove_file(&tmp).await.ok();
|
||||
return Err(DownloadError::ShaMismatch {
|
||||
expected: expected_sha.to_string(),
|
||||
actual,
|
||||
});
|
||||
}
|
||||
|
||||
tokio::fs::rename(&tmp, dest).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tempfile::tempdir;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
#[test]
|
||||
fn model_path_contains_model_dir_and_filename() {
|
||||
let path = model_path(LlmModelId::Qwen3_1_7B_Q4);
|
||||
assert!(path.to_string_lossy().ends_with("Qwen3-1.7B-Q4_K_M.gguf"));
|
||||
assert!(path.starts_with(model_dir()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recommend_tier_prefers_mid_by_default() {
|
||||
let tier = recommend_tier(16 * 1024_u64.pow(3), None);
|
||||
assert_eq!(tier, LlmModelId::Qwen3_4BInstruct2507Q4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn download_impl_supports_resume_and_sha_verification() {
|
||||
let fixture = b"hello resumed download".to_vec();
|
||||
let expected_sha = format!("{:x}", Sha256::digest(&fixture));
|
||||
let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = server.local_addr().unwrap();
|
||||
let content = fixture.clone();
|
||||
|
||||
let server_task = tokio::spawn(async move {
|
||||
let (mut socket, _) = server.accept().await.unwrap();
|
||||
let mut request = vec![0u8; 2048];
|
||||
let size = socket.read(&mut request).await.unwrap();
|
||||
let request = String::from_utf8_lossy(&request[..size]).to_lowercase();
|
||||
let range_start = request
|
||||
.lines()
|
||||
.find_map(|line| line.strip_prefix("range: bytes="))
|
||||
.and_then(|line| line.strip_suffix('-'))
|
||||
.and_then(|line| line.trim().parse::<usize>().ok());
|
||||
|
||||
if let Some(start) = range_start {
|
||||
let body = &content[start..];
|
||||
let response = format!(
|
||||
"HTTP/1.1 206 Partial Content\r\nContent-Length: {}\r\nContent-Range: bytes {}-{}/{}\r\nAccept-Ranges: bytes\r\n\r\n",
|
||||
body.len(),
|
||||
start,
|
||||
content.len() - 1,
|
||||
content.len()
|
||||
);
|
||||
socket.write_all(response.as_bytes()).await.unwrap();
|
||||
socket.write_all(body).await.unwrap();
|
||||
} else {
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nAccept-Ranges: bytes\r\n\r\n",
|
||||
content.len()
|
||||
);
|
||||
socket.write_all(response.as_bytes()).await.unwrap();
|
||||
socket.write_all(&content).await.unwrap();
|
||||
}
|
||||
});
|
||||
|
||||
let dir = tempdir().unwrap();
|
||||
let dest = dir.path().join("fixture.gguf");
|
||||
let part = dest.with_extension("gguf.part");
|
||||
tokio::fs::write(&part, &fixture[..10]).await.unwrap();
|
||||
|
||||
let progress = Arc::new(Mutex::new(Vec::new()));
|
||||
let progress_clone = progress.clone();
|
||||
download_impl(
|
||||
&format!("http://{addr}/fixture.gguf"),
|
||||
&expected_sha,
|
||||
&dest,
|
||||
move |done, total| progress_clone.lock().unwrap().push((done, total)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let saved = tokio::fs::read(&dest).await.unwrap();
|
||||
assert_eq!(saved, fixture);
|
||||
assert!(!part.exists());
|
||||
assert!(!progress.lock().unwrap().is_empty());
|
||||
|
||||
server_task.await.unwrap();
|
||||
}
|
||||
}
|
||||
12
crates/llm/src/prompts.rs
Normal file
12
crates/llm/src/prompts.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
pub const DECOMPOSE_TASK_SYSTEM: &str = "\
|
||||
You are a task-decomposition assistant. Given a task description, produce \
|
||||
between 3 and 7 concrete, physical micro-steps. Each step must be a short \
|
||||
imperative sentence, actionable today, with no commentary. Output ONLY a \
|
||||
JSON array of strings.";
|
||||
|
||||
pub const EXTRACT_TASKS_SYSTEM: &str = "\
|
||||
You are a task-extraction assistant. Given a transcript of spoken notes, \
|
||||
output a JSON array of action items the speaker committed to. Each item must \
|
||||
be a short imperative sentence. Omit observations, wishes, and background \
|
||||
context that are not explicit commitments. Output an empty array if there are \
|
||||
no action items.";
|
||||
62
crates/llm/tests/smoke.rs
Normal file
62
crates/llm/tests/smoke.rs
Normal file
@@ -0,0 +1,62 @@
|
||||
//! Smoke test: load a GGUF model and exercise the high-level wrappers.
|
||||
//!
|
||||
//! Verified against llama-cpp-2 `0.1.144` using:
|
||||
//! - `llama_backend::LlamaBackend`
|
||||
//! - `model::LlamaModel`
|
||||
//! - `context::params::LlamaContextParams`
|
||||
//! - `sampling::LlamaSampler`
|
||||
//!
|
||||
//! The test is gated behind `KON_LLM_TEST_MODEL`.
|
||||
|
||||
use std::env;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use kon_llm::LlmEngine;
|
||||
use kon_llm::LlmModelId;
|
||||
|
||||
#[test]
|
||||
fn llama_cpp_2_smoke_generates_and_wraps() {
|
||||
let model_path = match env::var("KON_LLM_TEST_MODEL") {
|
||||
Ok(path) => PathBuf::from(path),
|
||||
Err(_) => {
|
||||
eprintln!("KON_LLM_TEST_MODEL not set — skipping");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let engine = LlmEngine::new();
|
||||
engine
|
||||
.load_model(LlmModelId::Qwen3_1_7B_Q4, &model_path, true)
|
||||
.expect("load model");
|
||||
|
||||
let completion = engine
|
||||
.generate(
|
||||
"Write exactly one short greeting.",
|
||||
&kon_llm::GenerationConfig {
|
||||
max_tokens: 32,
|
||||
temperature: 0.0,
|
||||
stop_sequences: vec!["\n".to_string()],
|
||||
grammar: None,
|
||||
},
|
||||
)
|
||||
.expect("generate");
|
||||
assert!(!completion.trim().is_empty());
|
||||
|
||||
let cleaned = engine
|
||||
.cleanup_text(
|
||||
"You are a transcript cleanup assistant. Remove fillers and output only cleaned text.",
|
||||
"um hello there like general kenobi",
|
||||
)
|
||||
.expect("cleanup_text");
|
||||
assert!(!cleaned.trim().is_empty());
|
||||
|
||||
let tasks = engine
|
||||
.extract_tasks("I need to call the plumber tomorrow and buy milk.")
|
||||
.expect("extract_tasks");
|
||||
assert!(!tasks.is_empty());
|
||||
|
||||
let steps = engine
|
||||
.decompose_task("Plan a weekend trip to the coast")
|
||||
.expect("decompose_task");
|
||||
assert!((3..=7).contains(&steps.len()));
|
||||
}
|
||||
Reference in New Issue
Block a user