diff --git a/README.md b/README.md index 4b7d3a6..63a59e1 100644 --- a/README.md +++ b/README.md @@ -109,7 +109,7 @@ Old sessions are automatically pruned on each `wake shell` start. ### LLM Summarization -Wake automatically summarizes command outputs using a local LLM (Qwen2.5-0.5B). Summaries appear in `wake_list_commands` output, helping Claude quickly understand what happened without reading full output. +Wake automatically summarizes command outputs using a local LLM (Qwen3-0.6B). Summaries appear in `wake_list_commands` output, helping Claude quickly understand what happened without reading full output. **Enabled by default.** On first run, the model (~468MB) downloads automatically. diff --git a/crates/wake-cli/src/commands/llm.rs b/crates/wake-cli/src/commands/llm.rs index e2e6c67..c82aa6e 100644 --- a/crates/wake-cli/src/commands/llm.rs +++ b/crates/wake-cli/src/commands/llm.rs @@ -13,7 +13,7 @@ pub async fn status() -> Result<()> { println!(); // Model info - println!("Model: Qwen2.5-0.5B-Instruct (Q4_K_M)"); + println!("Model: Qwen3-0.6B (Q4_K_M)"); println!("Model path: {}", llm.model_path().display()); // Model status @@ -47,7 +47,7 @@ pub async fn download() -> Result<()> { return Ok(()); } - println!("Downloading summarization model (Qwen2.5-0.5B)..."); + println!("Downloading summarization model (Qwen3-0.6B)..."); println!("Destination: {}", llm.model_path().display()); println!(); diff --git a/crates/wake-cli/tests/cli_tests.rs b/crates/wake-cli/tests/cli_tests.rs index f7a3b28..085db1f 100644 --- a/crates/wake-cli/tests/cli_tests.rs +++ b/crates/wake-cli/tests/cli_tests.rs @@ -231,7 +231,7 @@ fn test_llm_status_shows_model_info() { assert!(output.status.success()); let stdout = String::from_utf8_lossy(&output.stdout); assert!(stdout.contains("Model:")); - assert!(stdout.contains("Qwen2.5")); + assert!(stdout.contains("Qwen3")); assert!(stdout.contains("Status:")); } diff --git a/crates/wake-llm/src/lib.rs b/crates/wake-llm/src/lib.rs index 6d317e1..027ad7c 100644 --- a/crates/wake-llm/src/lib.rs +++ b/crates/wake-llm/src/lib.rs @@ -1,6 +1,6 @@ //! Local LLM summarization for wake terminal recorder //! -//! This crate provides command output summarization using a local Qwen2.5-0.5B model +//! This crate provides command output summarization using a local Qwen3-0.6B model //! via llama.cpp. The model is downloaded on first use from HuggingFace. //! //! # Features @@ -23,10 +23,11 @@ use std::sync::Arc; use thiserror::Error; use tokio::sync::RwLock; -/// Default model to use for summarization (Qwen2.5-0.5B - small and fast) -pub const DEFAULT_MODEL: &str = "qwen2.5-0.5b-instruct-q4_k_m.gguf"; -pub const DEFAULT_MODEL_URL: &str = "https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct-GGUF/resolve/main/qwen2.5-0.5b-instruct-q4_k_m.gguf"; -pub const DEFAULT_MODEL_SIZE: u64 = 491_000_000; // ~491MB +/// Default model to use for summarization (Qwen3-0.6B - small and fast) +pub const DEFAULT_MODEL: &str = "Qwen3-0.6B-Q4_K_M.gguf"; +pub const DEFAULT_MODEL_URL: &str = + "https://huggingface.co/unsloth/Qwen3-0.6B-GGUF/resolve/main/Qwen3-0.6B-Q4_K_M.gguf"; +pub const DEFAULT_MODEL_SIZE: u64 = 397_000_000; // ~397MB #[derive(Debug, Error)] pub enum LlmError { @@ -177,7 +178,7 @@ mod tests { } assert!(DEFAULT_MODEL.ends_with(".gguf")); assert!(DEFAULT_MODEL_URL.starts_with("https://")); - assert!(DEFAULT_MODEL_URL.contains("Qwen")); + assert!(DEFAULT_MODEL_URL.contains("Qwen3")); #[allow(clippy::assertions_on_constants)] { assert!(DEFAULT_MODEL_SIZE > 100_000_000); // > 100MB diff --git a/crates/wake-llm/src/summarize.rs b/crates/wake-llm/src/summarize.rs index 61a0de2..0df8543 100644 --- a/crates/wake-llm/src/summarize.rs +++ b/crates/wake-llm/src/summarize.rs @@ -12,12 +12,7 @@ pub enum SummarizeError { } /// System prompt for summarization -const SYSTEM_PROMPT: &str = r#"You are a terminal output summarizer. Given a command and its output, provide a brief, informative summary in 1-2 sentences. Focus on: -- What happened (success/failure) -- Key results or changes -- Any errors or warnings - -Be concise and technical. Do not include the command itself in the summary."#; +const SYSTEM_PROMPT: &str = "Summarize this terminal output in 1-2 sentences."; /// Minimum output length (bytes) worth summarizing pub const MIN_OUTPUT_BYTES: usize = 100; @@ -46,7 +41,8 @@ pub fn summarize(model: &Model, command: &str, output: &str) -> Result Result String { + // Strip any ... blocks (Qwen3 reasoning mode) + let summary = strip_think_blocks(summary); + summary .trim() // Remove any leading "Summary:" or similar .strip_prefix("Summary:") - .unwrap_or(summary) + .unwrap_or(&summary) .trim() // Limit to reasonable length .chars() @@ -75,6 +74,22 @@ fn clean_summary(summary: &str) -> String { .to_string() } +/// Strip ... blocks from Qwen3 output +fn strip_think_blocks(text: &str) -> String { + let mut result = text.to_string(); + while let Some(start) = result.find("") { + if let Some(end) = result.find("") { + // Remove the think block including tags + result = format!("{}{}", &result[..start], &result[end + 8..]); + } else { + // Unclosed think tag - remove from to end + result = result[..start].to_string(); + break; + } + } + result +} + #[cfg(test)] mod tests { use super::*; @@ -110,6 +125,23 @@ mod tests { assert_eq!(clean_summary(text), text); } + #[test] + fn test_strip_think_blocks() { + assert_eq!(strip_think_blocks("reasoning hereThe answer"), "The answer"); + assert_eq!(strip_think_blocks("Before reasoning after"), "Before after"); + assert_eq!(strip_think_blocks("No think blocks here"), "No think blocks here"); + // Unclosed think tag + assert_eq!(strip_think_blocks("Start incomplete"), "Start "); + } + + #[test] + fn test_clean_summary_strips_think_blocks() { + assert_eq!( + clean_summary("Let me analyze this...Build succeeded with 5 warnings."), + "Build succeeded with 5 warnings." + ); + } + #[test] fn test_min_output_bytes_constant() { assert_eq!(MIN_OUTPUT_BYTES, 100); @@ -126,6 +158,6 @@ mod tests { { assert!(!SYSTEM_PROMPT.is_empty()); } - assert!(SYSTEM_PROMPT.contains("summarizer")); + assert!(SYSTEM_PROMPT.contains("Summarize")); } }