diff --git a/.claude/settings.json b/.claude/settings.json index 6f779d7..d7b6e96 100644 --- a/.claude/settings.json +++ b/.claude/settings.json @@ -1,7 +1,4 @@ { - "enabledPlugins": { - "rust-lsp@claude-plugins-official": true - }, "permissions": { "allow": [ "Bash(bun run:*)", @@ -12,5 +9,9 @@ "Bash(git log:*)", "Bash(git branch:*)" ] + }, + "enabledPlugins": { + "rust-lsp@claude-plugins-official": true, + "rust-analyzer-lsp@claude-plugins-official": true } } diff --git a/.gitignore b/.gitignore index a433209..40645ca 100644 --- a/.gitignore +++ b/.gitignore @@ -27,6 +27,7 @@ dist-ssr .ai_agents/session_context/ .ai_agents/ .ai_agents/coderabbit_output.txt +/.pi-subagents/ target/ # Environment diff --git a/Cargo.lock b/Cargo.lock index 6afec06..4e5651c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -220,6 +220,7 @@ dependencies = [ "nix 0.29.0", "parakeet-rs", "realfft 0.4.0", + "regex", "reqwest", "ringbuf", "rusqlite", diff --git a/Cargo.toml b/Cargo.toml index 9ac483f..a4ae9e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,9 @@ futures-util = "0.3" sha2 = "0.10" arc-swap = "1" +# Text post-processing +regex = "1" + # GPUI and UI dependencies gpui = { git = "https://github.com/BumpyClock/gpui", rev = "4332ea7deae4838c12bad6ea64292ca22a33cf98", features = [ "font-kit", diff --git a/README.md b/README.md index d2669f7..7082f0f 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ A native cross-platform voice assistant built with Rust and GPUI. ## Architecture This is a Rust workspace. Crates (see `[workspace].members` in `Cargo.toml` for the authoritative list): + - `crates/shared` - Shared utilities and types - `crates/core` - Core backend services (audio, transcription, and platform permission/runtime bridges) - `crates/app` - GPUI-based native UI application @@ -50,6 +51,14 @@ KERNEL=="event*", SUBSYSTEM=="input", GROUP="input", MODE="0660" Then add your user to the `input` group. +## Transcript enhancement + +Settings include deterministic custom dictionary replacements, soft vocabulary boosting, +and optional AI enhancement via OpenAI-compatible providers (OpenAI, Ollama, +LM Studio). OpenAI reads API keys from an env var (default `OPENAI_API_KEY`); +local model hosting inside Ansible is deferred. See +`docs/design/ai-enhancement-dictionary-vocabulary.md`. + ## Supported Models | Model | Size | Notes | diff --git a/crates/app/src/ui/settings/sections/models/sidecar.rs b/crates/app/src/ui/settings/sections/models/sidecar.rs index 579a1d2..0059965 100644 --- a/crates/app/src/ui/settings/sections/models/sidecar.rs +++ b/crates/app/src/ui/settings/sections/models/sidecar.rs @@ -21,6 +21,8 @@ pub fn render_whisper_sidecar_card( let muted_fg = cx.theme().muted_foreground; let app_state = cx.global::().clone(); let progress = status.progress.clamp(0.0, 100.0); + let needs_update = status.needs_update; + let override_configured = status.override_configured; div() .p_3() @@ -53,20 +55,30 @@ pub fn render_whisper_sidecar_card( .text_color(muted_fg), ) .child( - Label::new("Install to show Whisper model cards.") - .text_xs() - .text_color(muted_fg), + Label::new(if override_configured { + "Managed install is unavailable while the sidecar path override is set." + } else if needs_update { + "Update the sidecar to restore Whisper model support." + } else { + "Install to show Whisper model cards." + }) + .text_xs() + .text_color(muted_fg), ), ) .child( Button::new("install-whisper-sidecar") .label(if status.installing { "Installing…" + } else if override_configured { + "Managed install unavailable" + } else if needs_update { + "Update Whisper" } else { "Install Whisper" }) .outline() - .disabled(status.installing) + .disabled(status.installing || override_configured) .on_click(move |_event, _window, _cx| { let runtime = app_state.runtime().clone(); let controller = app_state.controller().clone(); @@ -96,10 +108,12 @@ pub fn render_whisper_sidecar_card( }) .when(status.last_error.is_some(), |this| { let message = status.last_error.unwrap_or_default(); - this.child(Alert::error( - "whisper-sidecar-error", - format!("Install failed: {message}"), - )) + let message = if override_configured { + message + } else { + format!("Install failed: {message}") + }; + this.child(Alert::error("whisper-sidecar-error", message)) }) .into_any_element() } diff --git a/crates/app/src/ui/settings/sections/transcription.rs b/crates/app/src/ui/settings/sections/transcription.rs index e1518b0..047b9b7 100644 --- a/crates/app/src/ui/settings/sections/transcription.rs +++ b/crates/app/src/ui/settings/sections/transcription.rs @@ -9,6 +9,10 @@ use ansible_shared::settings::SttAccelerator; use gpui::*; use gpui_component::setting::NumberFieldOptions; use gpui_component::setting::{SettingControl, SettingGroup, SettingItem, SettingPage}; +use transcription_enhancement::{create_ai_enhancement_group, create_dictionary_group}; + +#[path = "transcription_enhancement.rs"] +mod transcription_enhancement; static STT_ACCELERATOR_OPTIONS: LazyLock> = LazyLock::new(|| { SttAccelerator::available() @@ -115,6 +119,8 @@ pub fn create_transcription_section(state: &Entity) -> SettingPag ), ), ) + .group(create_ai_enhancement_group(state)) + .group(create_dictionary_group(state)) .group( SettingGroup::new() .title("Minimum Utterance") diff --git a/crates/app/src/ui/settings/sections/transcription_enhancement.rs b/crates/app/src/ui/settings/sections/transcription_enhancement.rs new file mode 100644 index 0000000..5fab323 --- /dev/null +++ b/crates/app/src/ui/settings/sections/transcription_enhancement.rs @@ -0,0 +1,391 @@ +//! AI enhancement, custom dictionary, and vocabulary settings. + +use std::str::FromStr; +use std::sync::LazyLock; + +use crate::ui::settings::SettingsState; +use ansible_shared::settings::{AiProvider, CustomDictionaryEntry}; +use gpui::*; +use gpui_component::setting::NumberFieldOptions; +use gpui_component::setting::{SettingControl, SettingGroup, SettingItem}; + +static AI_PROVIDER_OPTIONS: LazyLock> = LazyLock::new(|| { + AiProvider::ALL + .iter() + .map(|v| (v.as_str().into(), v.label().into())) + .collect() +}); + +fn format_terms(words: &[String]) -> SharedString { + words.join(", ").into() +} + +fn parse_terms(value: &str) -> Vec { + let mut out = Vec::new(); + for term in value.split([',', ';', '\n']) { + let term = term.trim(); + if term.is_empty() || out.iter().any(|v: &String| v.eq_ignore_ascii_case(term)) { + continue; + } + out.push(term.to_string()); + } + out +} + +fn format_dictionary(entries: &[CustomDictionaryEntry]) -> SharedString { + entries + .iter() + .map(ToString::to_string) + .collect::>() + .join("; ") + .into() +} + +fn parse_dictionary(value: &str) -> Vec { + value + .split([';', '\n']) + .filter_map(|entry| { + let entry = entry.trim(); + if entry.is_empty() { + return None; + } + match CustomDictionaryEntry::from_str(entry) { + Ok(entry) => Some(entry), + Err(_) => { + log::warn!("Ignoring invalid custom dictionary entry"); + None + } + } + }) + .collect() +} + +fn blank_to_none(value: SharedString) -> Option { + let value = value.trim(); + if value.is_empty() { + None + } else { + Some(value.to_string()) + } +} + +pub fn create_ai_enhancement_group(state: &Entity) -> SettingGroup { + SettingGroup::new() + .title("AI Enhancement") + .item( + SettingItem::new( + "Enhance Transcript", + SettingControl::switch( + { + let state = state.clone(); + move |cx: &App| { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .enabled + } + }, + { + let state = state.clone(); + move |val: bool, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.enabled = val; + }); + }); + } + }, + ), + ) + .description( + "Runs a Rust-backend OpenAI-compatible post-pass after ASR. OpenAI uses the configured env var; Ollama and LM Studio use local endpoints.", + ), + ) + .item(SettingItem::new( + "Provider", + SettingControl::dropdown( + AI_PROVIDER_OPTIONS.clone(), + { + let state = state.clone(); + move |cx: &App| -> SharedString { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .provider + .as_str() + .into() + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + if let Ok(provider) = AiProvider::from_str(&val) { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + let current_model = s.transcription.ai_enhancement.model.trim(); + let should_update_model = current_model.is_empty() + || current_model + == s.transcription + .ai_enhancement + .provider + .default_model(); + s.transcription.ai_enhancement.provider = provider; + if should_update_model { + s.transcription.ai_enhancement.model = + provider.default_model().to_string(); + } + }); + }); + } + } + }, + ), + )) + .item(SettingItem::new( + "Model", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .model + .clone() + .into() + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.model = val.trim().to_string(); + }); + }); + } + }, + ), + )) + .item( + SettingItem::new( + "Base URL", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .base_url + .clone() + .unwrap_or_default() + .into() + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.base_url = blank_to_none(val); + }); + }); + } + }, + ), + ) + .description("Blank uses the provider default. Non-loopback endpoints require trusted HTTPS. Custom OpenAI URLs receive the bearer key."), + ) + .item( + SettingItem::new( + "OpenAI Key Env", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .api_key_env_var + .clone() + .into() + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.api_key_env_var = + val.trim().to_string(); + }); + }); + } + }, + ), + ) + .description("Env var name, not the secret value. Default: OPENAI_API_KEY."), + ) + .item( + SettingItem::new( + "Timeout (s)", + SettingControl::number_input( + NumberFieldOptions { + min: 1.0, + max: 120.0, + step: 1.0, + }, + { + let state = state.clone(); + move |cx: &App| { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .timeout_secs as f64 + } + }, + { + let state = state.clone(); + move |val: f64, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.timeout_secs = val as u64; + }); + }); + } + }, + ), + ) + .description("On timeout or error, Ansible uses the deterministic transcript."), + ) + .item( + SettingItem::new( + "Enhancement Prompt", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + state + .read(cx) + .settings + .transcription + .ai_enhancement + .system_prompt + .clone() + .into() + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.ai_enhancement.system_prompt = + val.trim().to_string(); + }); + }); + } + }, + ), + ) + .layout(Axis::Vertical) + .description("FluidVoice-style prompt sent as the user message prefix. Use `${transcript}` to place the transcript explicitly."), + ) +} + +pub fn create_dictionary_group(state: &Entity) -> SettingGroup { + SettingGroup::new() + .title("Custom Dictionary & Vocabulary") + .item( + SettingItem::new( + "Dictionary", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + format_dictionary(&state.read(cx).settings.transcription.custom_dictionary) + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.custom_dictionary = parse_dictionary(&val); + }); + }); + } + }, + ), + ) + .layout(Axis::Vertical) + .description("Deterministic replacements after ASR: `port key -> Portkey; teh -> the`. Applied before AI and history."), + ) + .item( + SettingItem::new( + "Vocabulary Boost", + SettingControl::switch( + { + let state = state.clone(); + move |cx: &App| { + state + .read(cx) + .settings + .transcription + .vocabulary_boost + .enabled + } + }, + { + let state = state.clone(); + move |val: bool, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.vocabulary_boost.enabled = val; + }); + }); + } + }, + ), + ) + .description("Soft boost for Whisper via initial prompt. Parakeet/Nemotron currently expose no true boost API."), + ) + .item( + SettingItem::new( + "Vocabulary Terms", + SettingControl::input( + { + let state = state.clone(); + move |cx: &App| -> SharedString { + format_terms(&state.read(cx).settings.transcription.vocabulary_boost.words) + } + }, + { + let state = state.clone(); + move |val: SharedString, cx: &mut App| { + state.update(cx, |state, cx| { + state.apply_immediate(cx, |s| { + s.transcription.vocabulary_boost.words = parse_terms(&val); + }); + }); + } + }, + ), + ) + .layout(Axis::Vertical) + .description("Comma/semicolon-separated proper nouns, product names, acronyms, and domain terms."), + ) +} diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index 2b914dc..b0d0ef5 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -38,6 +38,9 @@ sha2.workspace = true # Storage rusqlite.workspace = true +# Text post-processing +regex.workspace = true + # Utilities toml.workspace = true diff --git a/crates/core/src/ai.rs b/crates/core/src/ai.rs new file mode 100644 index 0000000..8167cd6 --- /dev/null +++ b/crates/core/src/ai.rs @@ -0,0 +1,725 @@ +//! AI transcript enhancement via OpenAI-compatible providers. +//! +//! Supported providers are external/OpenAI-compatible endpoints only: +//! OpenAI, Ollama, and LM Studio. In-app local model hosting is deferred. + +use std::{net::IpAddr, time::Duration}; + +use ansible_shared::{ + AiEnhancementSettings, AiProvider, CustomDictionaryEntry, VocabularyBoostSettings, +}; +use anyhow::{Context, Result, anyhow, bail}; +use reqwest::header::{AUTHORIZATION, CONTENT_TYPE, HeaderMap, HeaderValue}; +use serde::Deserialize; +use serde_json::{Value, json}; + +use crate::transcription::postprocess::build_vocabulary_prompt; + +const FLUID_TRANSCRIPT_PLACEHOLDER: &str = "${transcript}"; +const DICTATION_TEMPERATURE: f32 = 0.2; +const LOW_REASONING_EFFORT: &str = "low"; +const MAX_RESPONSE_BODY_BYTES: usize = 1024 * 1024; + +#[derive(Clone)] +pub(crate) struct AiEnhancer { + client: reqwest::Client, +} + +impl AiEnhancer { + pub(crate) fn new() -> Self { + Self { + client: reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .no_proxy() + .build() + .expect("AI HTTP client should build"), + } + } + + pub(crate) async fn enhance_transcript( + &self, + transcript: &str, + settings: &AiEnhancementSettings, + dictionary: &[CustomDictionaryEntry], + boost: &VocabularyBoostSettings, + ) -> Result { + if !settings.enabled { + return Ok(transcript.to_string()); + } + let transcript = transcript.trim(); + if transcript.is_empty() { + return Ok(String::new()); + } + + let config = ProviderConfig::from_settings(settings)?; + let request_kind = RequestKind::for_endpoint(&config.base_url, &config.model); + let prompt = append_hint_section(&settings.system_prompt, dictionary, boost); + let user_message = render_dictation_user_message(&prompt, transcript); + let payload = request_kind.payload(&config.model, user_message); + + let mut headers = HeaderMap::new(); + headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + if let Some(key) = config.api_key { + let value = HeaderValue::from_str(&format!("Bearer {key}")) + .context("OpenAI API key contains invalid header characters")?; + headers.insert(AUTHORIZATION, value); + } + + let response = self + .client + .post(&request_kind.url) + .headers(headers) + .timeout(Duration::from_secs(settings.timeout_secs.max(1))) + .json(&payload) + .send() + .await + .context("AI enhancement request failed")?; + + let status = response.status(); + if !status.is_success() { + bail!("AI enhancement request failed with HTTP {status}"); + } + let body = read_response_body(response).await?; + + match request_kind.api { + ApiKind::ChatCompletions => parse_chat_completion_text(&body), + ApiKind::Responses => parse_responses_text(&body), + } + } +} + +async fn read_response_body(mut response: reqwest::Response) -> Result { + if response + .content_length() + .is_some_and(|length| length > MAX_RESPONSE_BODY_BYTES as u64) + { + bail!("AI enhancement response body exceeded 1 MiB"); + } + + let mut body = Vec::new(); + while let Some(chunk) = response + .chunk() + .await + .context("AI enhancement response body read failed")? + { + if body.len().saturating_add(chunk.len()) > MAX_RESPONSE_BODY_BYTES { + bail!("AI enhancement response body exceeded 1 MiB"); + } + body.extend_from_slice(&chunk); + } + String::from_utf8(body).context("AI enhancement response body was not valid UTF-8") +} + +struct ProviderConfig { + base_url: String, + model: String, + api_key: Option, +} + +impl ProviderConfig { + fn from_settings(settings: &AiEnhancementSettings) -> Result { + let base_url = settings + .base_url + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| settings.provider.default_base_url()); + let base_url = validate_base_url(settings.provider, base_url)?; + let model = if settings.model.trim().is_empty() { + settings.provider.default_model().to_string() + } else { + settings.model.trim().to_string() + }; + let api_key = match settings.provider { + AiProvider::OpenAi => { + let env_var = settings.api_key_env_var.trim(); + if env_var.is_empty() { + bail!("OpenAI API key env var is empty"); + } + let key = std::env::var(env_var) + .with_context(|| format!("{env_var} is not set for OpenAI enhancement"))?; + let key = key.trim(); + if key.is_empty() { + bail!("{env_var} is empty for OpenAI enhancement"); + } + Some(key.to_string()) + } + AiProvider::Ollama | AiProvider::LmStudio => None, + }; + Ok(Self { + base_url, + model, + api_key, + }) + } +} + +fn validate_base_url(provider: AiProvider, base_url: &str) -> Result { + let url = reqwest::Url::parse(base_url).context("AI provider base URL is invalid")?; + if !matches!(url.scheme(), "http" | "https") { + bail!("AI provider base URL must use HTTP or HTTPS"); + } + if url.host_str().is_none() { + bail!("AI provider base URL must include a host"); + } + if !url.username().is_empty() || url.password().is_some() { + bail!("AI provider base URL must not include credentials"); + } + if url.query().is_some() { + bail!("AI provider base URL must not include a query"); + } + if url.fragment().is_some() { + bail!("AI provider base URL must not include a fragment"); + } + if url.scheme() != "https" && !is_loopback_url(&url) { + bail!( + "{} base URL must use HTTPS unless it targets loopback", + provider.label() + ); + } + Ok(url.to_string()) +} + +fn is_loopback_url(url: &reqwest::Url) -> bool { + let Some(host) = url.host_str() else { + return false; + }; + let host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .is_ok_and(|address| address.is_loopback()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ApiKind { + ChatCompletions, + Responses, +} + +#[derive(Debug)] +struct RequestKind { + api: ApiKind, + url: String, + reasoning_effort: Option<&'static str>, + temperature: Option, +} + +impl RequestKind { + fn for_endpoint(base_url: &str, model: &str) -> Self { + let trimmed = base_url.trim().trim_end_matches('/'); + let explicit_responses = trimmed.ends_with("/responses"); + let explicit_chat = trimmed.ends_with("/chat/completions"); + let reasoning_effort = low_reasoning_effort(model); + let api = if explicit_responses + || (!explicit_chat && is_openai_api(trimmed) && should_use_responses_api(model)) + { + ApiKind::Responses + } else { + ApiKind::ChatCompletions + }; + let url = match api { + ApiKind::ChatCompletions if explicit_chat => trimmed.to_string(), + ApiKind::ChatCompletions => format!("{trimmed}/chat/completions"), + ApiKind::Responses if explicit_responses => trimmed.to_string(), + ApiKind::Responses => format!("{trimmed}/responses"), + }; + Self { + api, + url, + reasoning_effort, + temperature: if is_temperature_unsupported(model) { + None + } else { + Some(DICTATION_TEMPERATURE) + }, + } + } + + fn payload(&self, model: &str, user_message: String) -> Value { + match self.api { + ApiKind::ChatCompletions => self.chat_payload(model, user_message), + ApiKind::Responses => self.responses_payload(model, user_message), + } + } + + fn chat_payload(&self, model: &str, user_message: String) -> Value { + let mut payload = json!({ + "model": model, + "messages": [ + { + "role": "user", + "content": user_message, + } + ], + "stream": false, + }); + if let Some(temperature) = self.temperature { + payload["temperature"] = json!(temperature); + } + if let Some(effort) = self.reasoning_effort { + payload["reasoning_effort"] = json!(effort); + } + payload + } + + fn responses_payload(&self, model: &str, user_message: String) -> Value { + let mut payload = json!({ + "model": model, + "input": [ + { + "role": "user", + "content": [ + { + "type": "input_text", + "text": user_message, + } + ], + } + ], + "store": false, + "stream": false, + }); + if let Some(temperature) = self.temperature { + payload["temperature"] = json!(temperature); + } + if let Some(effort) = self.reasoning_effort { + payload["reasoning"] = json!({ "effort": effort }); + } + payload + } +} + +fn is_openai_api(base_url: &str) -> bool { + reqwest::Url::parse(base_url) + .ok() + .and_then(|url| { + url.host_str() + .map(|host| host.eq_ignore_ascii_case("api.openai.com")) + }) + .unwrap_or(false) +} + +fn should_use_responses_api(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + model.starts_with("gpt-5") + || model.contains("gpt-5.") + || model.starts_with("o1") + || model.starts_with("o3") + || model.starts_with("o4") +} + +fn is_reasoning_model(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + should_use_responses_api(&model) + || model.contains("gpt-oss") + || model.starts_with("openai/") + || (model.contains("deepseek") && model.contains("reasoner")) +} + +fn supports_reasoning_effort(model: &str) -> bool { + let model = model.to_ascii_lowercase(); + should_use_responses_api(&model) || model.contains("gpt-oss") || model.starts_with("openai/") +} + +fn low_reasoning_effort(model: &str) -> Option<&'static str> { + let model = model.to_ascii_lowercase(); + if model == "gpt-5-pro" || model.starts_with("gpt-5-pro-") { + return None; + } + if supports_reasoning_effort(&model) { + Some(LOW_REASONING_EFFORT) + } else { + None + } +} + +fn is_temperature_unsupported(model: &str) -> bool { + is_reasoning_model(model) +} + +fn append_hint_section( + prompt_text: &str, + dictionary: &[CustomDictionaryEntry], + boost: &VocabularyBoostSettings, +) -> String { + let mut prompt = prompt_text.trim().to_string(); + if !dictionary.is_empty() { + prompt.push_str("\n\n## Custom Dictionary:\nApply these replacements exactly:\n"); + for entry in dictionary { + prompt.push_str("- "); + prompt.push_str(entry.from()); + prompt.push_str(" -> "); + prompt.push_str(entry.to()); + prompt.push('\n'); + } + } + if let Some(vocabulary) = build_vocabulary_prompt(boost, dictionary) { + prompt.push_str("\n\n## Vocabulary Hints:\n"); + prompt.push_str(&vocabulary); + } + prompt +} + +fn render_dictation_user_message(prompt_text: &str, transcript: &str) -> String { + if prompt_text.contains(FLUID_TRANSCRIPT_PLACEHOLDER) { + return prompt_text.replace(FLUID_TRANSCRIPT_PLACEHOLDER, transcript); + } + let trimmed_prompt = prompt_text.trim(); + if trimmed_prompt.is_empty() { + transcript.to_string() + } else { + format!("{prompt_text}\n\n{transcript}") + } +} + +fn parse_chat_completion_text(body: &str) -> Result { + let parsed: ChatCompletionResponse = serde_json::from_str(body) + .context("AI enhancement response was not valid chat-completions JSON")?; + parsed + .choices + .into_iter() + .next() + .filter(|choice| choice.finish_reason.as_deref() == Some("stop")) + .map(|choice| choice.message.content.trim().to_string()) + .filter(|text| !text.is_empty()) + .ok_or_else(|| anyhow!("AI enhancement returned no transcript")) +} + +fn parse_responses_text(body: &str) -> Result { + let parsed: Value = serde_json::from_str(body) + .context("AI enhancement response was not valid Responses API JSON")?; + if parsed.get("status").and_then(Value::as_str) != Some("completed") { + return Err(anyhow!("AI enhancement response was not completed")); + } + if let Some(text) = parsed.get("output_text").and_then(Value::as_str) { + let text = text.trim(); + if !text.is_empty() { + return Ok(text.to_string()); + } + } + let text = parsed + .get("output") + .and_then(Value::as_array) + .into_iter() + .flatten() + .flat_map(|item| { + item.get("content") + .and_then(Value::as_array) + .into_iter() + .flatten() + }) + .filter_map(|content| content.get("text").and_then(Value::as_str)) + .collect::(); + let text = text.trim(); + if text.is_empty() { + Err(anyhow!("AI enhancement returned no transcript")) + } else { + Ok(text.to_string()) + } +} + +#[derive(Debug, Deserialize)] +struct ChatCompletionResponse { + #[serde(default)] + choices: Vec, +} + +#[derive(Debug, Deserialize)] +struct ChatChoice { + message: ChatChoiceMessage, + finish_reason: Option, +} + +#[derive(Debug, Deserialize)] +struct ChatChoiceMessage { + content: String, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn chat_url_appends_chat_completions() { + let request = RequestKind::for_endpoint("http://localhost:11434/v1", "llama3.1"); + assert_eq!(request.api, ApiKind::ChatCompletions); + assert_eq!(request.url, "http://localhost:11434/v1/chat/completions"); + let request = + RequestKind::for_endpoint("http://localhost:11434/v1/chat/completions", "llama3.1"); + assert_eq!(request.url, "http://localhost:11434/v1/chat/completions"); + } + + #[test] + fn openai_reasoning_model_uses_responses_api_and_low_effort() { + let request = RequestKind::for_endpoint("https://api.openai.com/v1", "gpt-5-mini"); + assert_eq!(request.api, ApiKind::Responses); + assert_eq!(request.url, "https://api.openai.com/v1/responses"); + assert_eq!(request.reasoning_effort, Some("low")); + assert_eq!(request.temperature, None); + let payload = request.payload("gpt-5-mini", "clean this".to_string()); + assert_eq!(payload["reasoning"]["effort"], "low"); + assert!(payload.get("temperature").is_none()); + } + + #[test] + fn gpt_5_pro_omits_unsupported_low_reasoning_effort() { + for model in ["gpt-5-pro", "GPT-5-PRO-2025-08-07"] { + let request = RequestKind::for_endpoint("https://api.openai.com/v1", model); + assert_eq!(request.api, ApiKind::Responses); + assert_eq!(request.reasoning_effort, None); + let payload = request.payload(model, "clean this".to_string()); + assert!(payload.get("reasoning").is_none()); + } + } + + #[test] + fn openai_host_detection_rejects_lookalike_hosts() { + let request = + RequestKind::for_endpoint("https://api.openai.com.attacker.invalid/v1", "gpt-5-mini"); + assert_eq!(request.api, ApiKind::ChatCompletions); + assert_eq!( + request.url, + "https://api.openai.com.attacker.invalid/v1/chat/completions" + ); + } + + #[test] + fn provider_base_url_rejects_malformed_urls() { + let error = validate_base_url(AiProvider::OpenAi, "not a URL") + .unwrap_err() + .to_string(); + assert!(error.contains("base URL is invalid")); + } + + #[test] + fn provider_base_url_rejects_non_loopback_http() { + for provider in AiProvider::ALL { + let error = validate_base_url(*provider, "http://attacker.example/v1") + .unwrap_err() + .to_string(); + assert!(error.contains("must use HTTPS"), "{provider}: {error}"); + } + } + + #[test] + fn provider_base_url_rejects_credentials_query_and_fragment() { + let credentials = + validate_base_url(AiProvider::OpenAi, "https://user:password@example.com/v1") + .unwrap_err() + .to_string(); + assert!(credentials.contains("must not include credentials")); + + let query = validate_base_url(AiProvider::OpenAi, "https://example.com/v1?tenant=x") + .unwrap_err() + .to_string(); + assert!(query.contains("must not include a query")); + + let fragment = validate_base_url(AiProvider::OpenAi, "https://example.com/v1#settings") + .unwrap_err() + .to_string(); + assert!(fragment.contains("must not include a fragment")); + } + + #[test] + fn provider_base_url_accepts_trusted_https_and_loopback_http() { + for provider in AiProvider::ALL { + assert!(validate_base_url(*provider, "https://trusted.example/v1").is_ok()); + assert!(validate_base_url(*provider, "http://localhost:1234/v1").is_ok()); + assert!(validate_base_url(*provider, "http://127.0.0.1:1234/v1").is_ok()); + assert!(validate_base_url(*provider, "http://[::1]:1234/v1").is_ok()); + } + } + + #[test] + fn openai_compatible_claude_model_keeps_temperature() { + let request = RequestKind::for_endpoint("http://localhost:1234/v1", "claude-opus-4-7"); + assert_eq!(request.temperature, Some(DICTATION_TEMPERATURE)); + } + + #[tokio::test] + #[allow(clippy::disallowed_methods)] + async fn client_ignores_system_proxy_configuration() { + const CHILD_ENV: &str = "ANSIBLE_AI_PROXY_TEST_CHILD"; + + if std::env::var_os(CHILD_ENV).is_some() { + let _ = AiEnhancer::new() + .client + .get("http://192.0.2.1:9/probe") + .timeout(Duration::from_millis(200)) + .send() + .await; + return; + } + + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let proxy_url = format!("http://{}", listener.local_addr().unwrap()); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "ai::tests::client_ignores_system_proxy_configuration", + "--nocapture", + ]) + .env(CHILD_ENV, "1") + .env("HTTP_PROXY", &proxy_url) + .env("http_proxy", &proxy_url) + .env("ALL_PROXY", &proxy_url) + .env("all_proxy", &proxy_url) + .env("NO_PROXY", "") + .env("no_proxy", "") + .output() + .unwrap(); + + assert!(output.status.success()); + assert!( + matches!(listener.accept(), Err(error) if error.kind() == std::io::ErrorKind::WouldBlock), + "AI client forwarded a request through the configured system proxy" + ); + } + + #[tokio::test] + async fn client_does_not_follow_redirects() { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + socket + .write_all( + b"HTTP/1.1 302 Found\r\nLocation: http://attacker.example/final\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .unwrap(); + }); + + let response = AiEnhancer::new() + .client + .get(format!("http://{address}/start")) + .send() + .await + .unwrap(); + assert_eq!(response.status(), reqwest::StatusCode::FOUND); + server.await.unwrap(); + } + + #[tokio::test] + async fn response_body_limit_rejects_oversized_content_length() { + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + MAX_RESPONSE_BODY_BYTES + 1 + ); + socket.write_all(response.as_bytes()).await.unwrap(); + }); + + let response = AiEnhancer::new() + .client + .get(format!("http://{address}/oversized")) + .send() + .await + .unwrap(); + let error = read_response_body(response).await.unwrap_err().to_string(); + assert!(error.contains("exceeded 1 MiB")); + server.await.unwrap(); + } + + #[test] + fn malformed_response_errors_do_not_include_response_body() { + let body = "private transcript content"; + let error = parse_chat_completion_text(body).unwrap_err().to_string(); + assert!(!error.contains(body)); + + let error = parse_responses_text(body).unwrap_err().to_string(); + assert!(!error.contains(body)); + } + + #[test] + fn chat_payload_uses_fluid_voice_user_only_prompt_shape() { + let request = RequestKind::for_endpoint("http://localhost:1234/v1", "local-model"); + let payload = request.payload("local-model", "prompt\n\ntranscript".to_string()); + assert_eq!(payload["messages"].as_array().unwrap().len(), 1); + assert_eq!(payload["messages"][0]["role"], "user"); + assert_eq!(payload["temperature"], DICTATION_TEMPERATURE); + assert_eq!(payload["stream"], false); + } + + #[test] + fn render_prompt_appends_transcript_like_fluid_voice() { + let rendered = render_dictation_user_message("Prompt body", "hello world"); + assert_eq!(rendered, "Prompt body\n\nhello world"); + } + + #[test] + fn render_prompt_supports_fluid_voice_placeholder() { + let rendered = render_dictation_user_message("Fix: ${transcript}", "hello world"); + assert_eq!(rendered, "Fix: hello world"); + } + + #[test] + fn user_prompt_includes_dictionary_and_transcript() { + let prompt = append_hint_section( + "Base prompt", + &[CustomDictionaryEntry::new("port key", "Portkey").unwrap()], + &VocabularyBoostSettings::default(), + ); + let rendered = render_dictation_user_message(&prompt, "hello port key"); + assert!(rendered.contains("port key -> Portkey")); + assert!(rendered.ends_with("\n\nhello port key")); + } + + #[test] + fn chat_text_requires_completed_stop_choice() { + let text = parse_chat_completion_text( + r#"{"choices":[{"message":{"content":" Clean text. "},"finish_reason":"stop"}]}"#, + ) + .unwrap(); + assert_eq!(text, "Clean text."); + + for body in [ + r#"{"choices":[{"message":{"content":"text"}}]}"#, + r#"{"choices":[{"message":{"content":"text"},"finish_reason":"length"}]}"#, + ] { + assert!(parse_chat_completion_text(body).is_err()); + } + } + + #[test] + fn responses_text_requires_completed_status() { + for body in [ + r#"{"output_text":"text"}"#, + r#"{"status":"in_progress","output_text":"text"}"#, + ] { + assert!(parse_responses_text(body).is_err()); + } + } + + #[test] + fn responses_text_parses_output_text() { + let text = parse_responses_text(r#"{"status":"completed","output_text":" Clean text. "}"#) + .unwrap(); + assert_eq!(text, "Clean text."); + } + + #[test] + fn responses_text_parses_all_output_content_text() { + let text = parse_responses_text( + r#"{"status":"completed","output":[{"content":[{"type":"output_text","text":" Clean "},{"type":"output_text","text":"text. "}]}]}"#, + ) + .unwrap(); + assert_eq!(text, "Clean text."); + } +} diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index 0c586a8..ad8c653 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -263,6 +263,57 @@ show_settings_on_launch = true fs::remove_dir_all(temp_dir).ok(); } + #[test] + fn load_settings_filters_invalid_dictionary_entries_without_resetting() { + let temp_dir = std::env::temp_dir().join(format!( + "ansible-settings-dictionary-{}", + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + fs::create_dir_all(&temp_dir).unwrap(); + let path = temp_dir.join("settings.toml"); + fs::write( + &path, + r#" +schema_version = 1 +hotkey_press_delay_ms = 123 + +[transcription] +selected_model = "whisper-turbo" + +[[transcription.custom_dictionary]] +from = "teh" +to = "the" + +[[transcription.custom_dictionary]] +from = "" +to = "X" + +[[transcription.custom_dictionary]] +from = 42 +to = "forty-two" +"#, + ) + .unwrap(); + + let settings = load_settings_from_path(&path).unwrap(); + + assert_eq!(settings.hotkey_press_delay_ms, 123); + assert_eq!( + settings.transcription.selected_model, + Some(ansible_shared::ModelId::WhisperTurbo) + ); + assert_eq!( + settings.transcription.custom_dictionary, + vec![ansible_shared::CustomDictionaryEntry::new("teh", "the").unwrap()] + ); + assert!(!temp_dir.join("settings.toml.bak.v1").exists()); + + fs::remove_dir_all(temp_dir).ok(); + } + #[test] fn save_settings_is_atomic_round_trips_and_leaves_no_temp_file() { let temp_dir = std::env::temp_dir().join(format!( diff --git a/crates/core/src/controller.rs b/crates/core/src/controller.rs index 8caf178..81ad76d 100644 --- a/crates/core/src/controller.rs +++ b/crates/core/src/controller.rs @@ -4,10 +4,14 @@ //! Model management, download, and sidecar methods live in [`controller::models`]. use std::sync::{Arc, Mutex}; +use std::time::Instant; -use ansible_shared::{AppSettings, CoreEvent, ModelId, SessionPhase}; +use ansible_shared::{ + AppSettings, CoreEvent, CustomDictionaryEntry, ModelId, SessionPhase, TranscriptionSettings, +}; use anyhow::{Context, Result, anyhow}; +use crate::ai::AiEnhancer; use crate::audio::capture::{LivePreviewConfig, StopResult}; use crate::audio::feedback::{FeedbackSound, play_feedback}; use crate::audio::processing::{ensure_vad_model_path, resolve_vad_model_path}; @@ -17,8 +21,10 @@ use crate::config; use crate::events::EventBus; use crate::sidecar::WhisperSidecarFactory; use crate::storage::{HistoryDb, HistoryEntryInput}; +use crate::transcription::postprocess::{apply_dictionary, build_vocabulary_prompt}; use crate::transcription::{ EngineRegistry, ModelManager, NemotronFactory, ParakeetFactory, TranscriptionEngineState, + TranscriptionOptions, }; use live_preview::{LivePreviewResult, LivePreviewSession, PreviewMode}; use tokio::sync::{Mutex as AsyncMutex, RwLock, mpsc as tokio_mpsc}; @@ -60,6 +66,12 @@ fn hotkey_fields_changed(previous: &AppSettings, new: &AppSettings) -> bool { || previous.hotkey_press_delay_ms != new.hotkey_press_delay_ms } +fn finalize_enhanced_transcript(enhanced: &str, dictionary: &[CustomDictionaryEntry]) -> String { + apply_dictionary(enhanced.trim(), dictionary) + .trim() + .to_string() +} + struct SessionPhaseGuard { event_bus: EventBus, restore_to: SessionPhase, @@ -80,6 +92,13 @@ impl Drop for SessionPhaseGuard { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum WarmupState { + Idle, + Running, + Pending, +} + pub struct AppController { event_bus: EventBus, settings: Arc>, @@ -89,7 +108,9 @@ pub struct AppController { model_manager: Arc, sidecar_manager: Arc, history_db: Arc, + ai_enhancer: AiEnhancer, runtime: tokio::runtime::Handle, + warmup_state: Arc>, settings_save_lock: Arc>, preview_session: Mutex>, pending_model_activation: Mutex>, @@ -146,7 +167,9 @@ impl AppController { model_manager, sidecar_manager, history_db, + ai_enhancer: AiEnhancer::new(), runtime, + warmup_state: Arc::new(Mutex::new(WarmupState::Idle)), settings_save_lock: Arc::new(AsyncMutex::new(())), preview_session: Mutex::new(None), pending_model_activation: Mutex::new(None), @@ -252,8 +275,12 @@ impl AppController { } pub async fn start_listening(&self) -> Result<()> { - if self.event_bus.session_phase() == SessionPhase::Listening { - return Ok(()); + match self.event_bus.session_phase() { + SessionPhase::Listening => return Ok(()), + SessionPhase::Transcribing => { + return Err(anyhow!("Transcription is still in progress")); + } + SessionPhase::Idle => {} } let ( @@ -379,17 +406,20 @@ impl AppController { } = stop_result?; let preview_text = preview.text.clone(); - let (has_selected_model, min_utterance_enabled, min_utterance_ms) = { + let transcription_settings = { let settings = self.settings.read().await; - ( - settings.transcription.selected_model.is_some(), - settings.transcription.min_utterance_ms.enabled, - settings.transcription.min_utterance_ms.value, - ) + settings.transcription.clone() + }; + let min_utterance_enabled = transcription_settings.min_utterance_ms.enabled; + let min_utterance_ms = transcription_settings.min_utterance_ms.value; + // Whisper is the only engine with `initial_prompt` support and uses the + // batch path; current streaming engines expose no vocabulary-hint option. + let transcription_options = TranscriptionOptions { + initial_prompt: build_vocabulary_prompt( + &transcription_settings.vocabulary_boost, + &transcription_settings.custom_dictionary, + ), }; - if !has_selected_model { - return Err(anyhow!("No transcription model selected")); - } // Streaming preview disables VAD to preserve fixed-size chunks, so // min-utterance uses total recording duration instead of VAD-filtered speech duration. @@ -423,22 +453,53 @@ impl AppController { return Ok(()); } - let history_db = self.history_db.clone(); - let history_text = text.clone(); - let entry_id = self - .runtime - .spawn_blocking(move || { - history_db.insert_entry(HistoryEntryInput { text: history_text }) - }) + let transcription_text = self + .postprocess_transcript(text, &transcription_settings) + .await; + if transcription_text.is_empty() { + self.erase_preview_text_if_trusted( + &preview_text, + preview.trusted, + "streaming finalized post-processed to empty text", + ) + .await; + return Ok(()); + } + + self.save_history(transcription_text.clone()).await?; + + if !preview.trusted { + log::error!( + "Live preview on-screen state is untrusted after a failed delta; saving transcript to history only, skipping focused-document output (preview chars: {})", + preview_text.len() + ); + } else if preview_text == transcription_text { + // Preview text already applied to focused app. + } else if let Err(err) = self + .apply_preview_delta_blocking(preview_text.clone(), transcription_text.clone()) .await - .context("History insert task failed to join")??; - self.event_bus.emit(CoreEvent::HistoryAdded { entry_id }); - // Preview text already applied to clipboard — nothing more to output. + { + log::error!( + "Failed to apply streaming post-process preview delta; transcript saved to history: {}", + err + ); + } return Ok(()); } + if transcription_settings.selected_model.is_none() { + self.erase_preview_text_if_trusted( + &preview_text, + preview.trusted, + "selected model disappeared before batch finalization", + ) + .await; + return Err(anyhow!("No transcription model selected")); + } + let audio_processing = self.audio_processing.clone(); let transcription_engine = self.transcription.clone(); + let transcription_options_for_batch = transcription_options.clone(); let outcome = self .runtime .spawn_blocking(move || { @@ -463,7 +524,10 @@ impl AppController { } transcription_engine - .transcribe(processed_audio.samples) + .transcribe_with_options( + processed_audio.samples, + transcription_options_for_batch, + ) .map(Some) }) .await @@ -479,8 +543,8 @@ impl AppController { return Ok(()); }; - let transcription_text = text.trim().to_string(); - if transcription_text.is_empty() { + let raw_transcription_text = text.trim().to_string(); + if raw_transcription_text.is_empty() { self.erase_preview_text_if_trusted( &preview_text, preview.trusted, @@ -490,16 +554,20 @@ impl AppController { return Ok(()); } - let history_db = self.history_db.clone(); - let history_text = transcription_text.clone(); - let entry_id = self - .runtime - .spawn_blocking(move || { - history_db.insert_entry(HistoryEntryInput { text: history_text }) - }) - .await - .context("History insert task failed to join")??; - self.event_bus.emit(CoreEvent::HistoryAdded { entry_id }); + let transcription_text = self + .postprocess_transcript(raw_transcription_text, &transcription_settings) + .await; + if transcription_text.is_empty() { + self.erase_preview_text_if_trusted( + &preview_text, + preview.trusted, + "batch fallback post-processed to empty transcript", + ) + .await; + return Ok(()); + } + + self.save_history(transcription_text.clone()).await?; if !preview.trusted { log::error!( @@ -548,6 +616,64 @@ impl AppController { &self.history_db } + async fn postprocess_transcript( + &self, + text: String, + settings: &TranscriptionSettings, + ) -> String { + let dictionary_text = apply_dictionary(text.trim(), &settings.custom_dictionary) + .trim() + .to_string(); + if dictionary_text.is_empty() || !settings.ai_enhancement.enabled { + return dictionary_text; + } + + let started = Instant::now(); + match self + .ai_enhancer + .enhance_transcript( + &dictionary_text, + &settings.ai_enhancement, + &settings.custom_dictionary, + &settings.vocabulary_boost, + ) + .await + { + Ok(enhanced) => { + let enhanced = finalize_enhanced_transcript(&enhanced, &settings.custom_dictionary); + if enhanced.is_empty() { + log::warn!( + "AI enhancement returned empty transcript; using deterministic transcript" + ); + dictionary_text + } else { + log::info!( + "ASR_BENCH ai_enhance_done elapsed_ms={:.1} chars_in={} chars_out={}", + started.elapsed().as_secs_f64() * 1000.0, + dictionary_text.len(), + enhanced.len() + ); + enhanced + } + } + Err(err) => { + log::warn!("AI enhancement failed; using deterministic transcript: {err:#}"); + dictionary_text + } + } + } + + async fn save_history(&self, text: String) -> Result { + let history_db = self.history_db.clone(); + let entry_id = self + .runtime + .spawn_blocking(move || history_db.insert_entry(HistoryEntryInput { text })) + .await + .context("History insert task failed to join")??; + self.event_bus.emit(CoreEvent::HistoryAdded { entry_id }); + Ok(entry_id) + } + /// Type `text` into the focused app. The settle-focus sleep + keystroke /// injection runs on a blocking pool thread, not the async worker (P1.4). async fn output_text(&self, text: &str) { @@ -629,18 +755,95 @@ impl AppController { /// switch, explicit download-activate). `warm_up` is idempotent when the /// engine is already resident (M3). fn spawn_model_preload(&self) { - spawn_warm_up(&self.runtime, self.transcription.clone()); + spawn_warm_up( + &self.runtime, + self.transcription.clone(), + self.warmup_state.clone(), + ); + } +} + +struct WarmupStateGuard { + state: Arc>, + armed: bool, +} + +impl WarmupStateGuard { + fn new(state: Arc>) -> Self { + Self { state, armed: true } + } +} + +impl Drop for WarmupStateGuard { + fn drop(&mut self) { + if self.armed { + *lock_warmup_state(&self.state) = WarmupState::Idle; + } + } +} + +fn lock_warmup_state(state: &Mutex) -> std::sync::MutexGuard<'_, WarmupState> { + state.lock().unwrap_or_else(|err| { + log::warn!("warm-up state mutex poisoned, recovering: {err}"); + err.into_inner() + }) +} + +fn request_warm_up(state: &Mutex) -> bool { + let mut state = lock_warmup_state(state); + match *state { + WarmupState::Idle => { + *state = WarmupState::Running; + true + } + WarmupState::Running => { + *state = WarmupState::Pending; + false + } + WarmupState::Pending => false, + } +} + +fn finish_warm_up_pass(guard: &mut WarmupStateGuard) -> bool { + let mut state = lock_warmup_state(&guard.state); + match *state { + WarmupState::Pending => { + *state = WarmupState::Running; + true + } + WarmupState::Running | WarmupState::Idle => { + *state = WarmupState::Idle; + guard.armed = false; + false + } } } /// Spawn a blocking task that warms (loads) the STT model so the first utterance -/// does not pay the load cost. Idempotent when the engine is already resident. +/// does not pay the load cost. Concurrent requests coalesce into one trailing +/// pass so a settings change during an active warm-up cannot leave a stale model. /// Shared by launch/settings-switch preload (`spawn_model_preload`), /// `start_listening`, and the auto-download activation path. -fn spawn_warm_up(runtime: &tokio::runtime::Handle, transcription: Arc) { +fn spawn_warm_up( + runtime: &tokio::runtime::Handle, + transcription: Arc, + state: Arc>, +) { + if !request_warm_up(&state) { + log::debug!("STT warm-up already in flight; coalescing preload"); + return; + } + let _preload_task = runtime.spawn_blocking(move || { - if let Err(err) = transcription.warm_up() { - log::warn!("Failed to warm up STT model: {err}"); + let mut guard = WarmupStateGuard::new(state); + loop { + if let Err(err) = transcription.warm_up() { + log::warn!("Failed to warm up STT model: {err}"); + } + if !finish_warm_up_pass(&mut guard) { + break; + } + log::debug!("Running coalesced STT warm-up"); } }); } @@ -649,7 +852,21 @@ fn spawn_warm_up(runtime: &tokio::runtime::Handle, transcription: Arc { // Warm up the freshly activated model so the first // utterance is instant (M3 auto-download preload). - super::spawn_warm_up(&runtime, transcription.clone()); + super::spawn_warm_up(&runtime, transcription.clone(), warmup_state.clone()); } Err(err) => { log::warn!("Failed to activate auto-downloaded Parakeet V3: {err}"); diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index ad01fb6..1748b16 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -13,6 +13,7 @@ //! - **Storage**: Settings persistence and history database //! - **Controller**: Lifecycle orchestration and event bus +mod ai; pub mod audio; pub mod automation; pub mod config; diff --git a/crates/core/src/sidecar/client.rs b/crates/core/src/sidecar/client.rs index 5b756d2..84233fc 100644 --- a/crates/core/src/sidecar/client.rs +++ b/crates/core/src/sidecar/client.rs @@ -8,13 +8,13 @@ use std::time::Duration; use ansible_shared::ModelId; use ansible_sidecar_proto::{ - PROTO_VERSION, SidecarRequest, SidecarResponse, read_json_frame, write_f32_frame, - write_json_frame, + PROTO_VERSION, SidecarRequest, SidecarResponse, SidecarTranscribeParams, read_json_frame, + write_f32_frame, write_json_frame, }; use anyhow::{Context, Result, bail}; use crate::sidecar::whisper_sidecar_path; -use crate::transcription::engine_trait::{SttEngine, SttEngineFactory}; +use crate::transcription::engine_trait::{SttEngine, SttEngineFactory, TranscriptionOptions}; pub struct WhisperSidecarFactory; @@ -110,8 +110,19 @@ impl WhisperSidecarEngine { impl SttEngine for WhisperSidecarEngine { fn transcribe(&mut self, samples: &[f32]) -> Result { + self.transcribe_with_options(samples, &TranscriptionOptions::default()) + } + + fn transcribe_with_options( + &mut self, + samples: &[f32], + options: &TranscriptionOptions, + ) -> Result { self.request(SidecarRequest::Transcribe { sample_count: samples.len(), + params: SidecarTranscribeParams { + initial_prompt: options.initial_prompt.clone(), + }, })?; write_f32_frame(&mut self.stdin, samples)?; match self.read_response()? { diff --git a/crates/core/src/sidecar/installer.rs b/crates/core/src/sidecar/installer.rs index a2d8798..3da6d29 100644 --- a/crates/core/src/sidecar/installer.rs +++ b/crates/core/src/sidecar/installer.rs @@ -12,8 +12,9 @@ use sha2::{Digest, Sha256}; use crate::events::EventBus; use crate::progress_throttle::should_emit_progress; use crate::sidecar::{ - SidecarManifest, SidecarStatus, WHISPER_PLUGIN, current_target, installed_manifest, - sidecar_binary_name, whisper_plugin_dir, whisper_sidecar_path, + SIDECAR_ENV_PATH, SidecarCompatibility, SidecarManifest, SidecarStatus, WHISPER_PLUGIN, + current_target, installed_manifest, sidecar_binary_name, sidecar_compatibility, + whisper_plugin_dir, whisper_sidecar_path_from_env, }; use ansible_shared::CoreEvent; @@ -49,19 +50,44 @@ impl SidecarManager { pub fn status(&self) -> SidecarStatus { let state = self.lock_state(); - let path = whisper_sidecar_path().ok(); - let manifest = installed_manifest().ok().flatten(); + let env_path = std::env::var_os(SIDECAR_ENV_PATH); + let override_configured = env_path.is_some(); + let path = whisper_sidecar_path_from_env(env_path).ok(); + let manifest = if override_configured { + None + } else { + installed_manifest().ok().flatten() + }; + let binary_exists = path.as_ref().is_some_and(|path| path.is_file()); + let compatibility = sidecar_compatibility( + override_configured, + binary_exists, + manifest.as_ref().map(|manifest| manifest.proto_version), + ); + let last_error = if override_configured { + (!binary_exists).then(|| { + format!("{SIDECAR_ENV_PATH} does not point to an existing file; fix or unset it") + }) + } else { + state.last_error.clone() + }; SidecarStatus { - installed: path.as_ref().is_some_and(|p| p.is_file()), + installed: compatibility == SidecarCompatibility::Compatible, + needs_update: compatibility == SidecarCompatibility::NeedsUpdate, installing: state.installing, progress: state.progress, - version: manifest.map(|m| m.version), + version: manifest.map(|manifest| manifest.version), path, - last_error: state.last_error.clone(), + override_configured, + last_error, } } pub async fn install_whisper(&self, event_bus: &EventBus) -> Result<()> { + if std::env::var_os(SIDECAR_ENV_PATH).is_some() { + bail!("Unset {SIDECAR_ENV_PATH} before installing the managed Whisper sidecar"); + } + { let mut state = self.lock_state(); if state.installing { diff --git a/crates/core/src/sidecar/mod.rs b/crates/core/src/sidecar/mod.rs index ae48e79..3becee6 100644 --- a/crates/core/src/sidecar/mod.rs +++ b/crates/core/src/sidecar/mod.rs @@ -10,6 +10,7 @@ use std::ffi::OsString; use std::fs; use std::path::PathBuf; +use ansible_sidecar_proto::PROTO_VERSION; use anyhow::Result; use serde::{Deserialize, Serialize}; @@ -28,19 +29,55 @@ pub struct SidecarManifest { #[derive(Debug, Clone, Default)] pub struct SidecarStatus { pub installed: bool, + pub needs_update: bool, pub installing: bool, pub progress: f64, pub version: Option, pub path: Option, + pub override_configured: bool, pub last_error: Option, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SidecarCompatibility { + Missing, + Compatible, + NeedsUpdate, +} + pub(crate) fn is_plugin_installed(plugin_name: &str) -> bool { matches!(plugin_name, WHISPER_PLUGIN) && is_whisper_sidecar_installed() } fn is_whisper_sidecar_installed() -> bool { - whisper_sidecar_path().is_ok_and(|path| path.is_file()) + let env_path = std::env::var_os(SIDECAR_ENV_PATH); + let binary_exists = + whisper_sidecar_path_from_env(env_path.clone()).is_ok_and(|path| path.is_file()); + let installed_proto_version = if env_path.is_some() { + None + } else { + installed_manifest() + .ok() + .flatten() + .map(|manifest| manifest.proto_version) + }; + + sidecar_compatibility(env_path.is_some(), binary_exists, installed_proto_version) + == SidecarCompatibility::Compatible +} + +fn sidecar_compatibility( + is_env_override: bool, + binary_exists: bool, + installed_proto_version: Option, +) -> SidecarCompatibility { + if !binary_exists { + SidecarCompatibility::Missing + } else if is_env_override || installed_proto_version == Some(PROTO_VERSION) { + SidecarCompatibility::Compatible + } else { + SidecarCompatibility::NeedsUpdate + } } pub(super) fn whisper_sidecar_path() -> Result { @@ -126,4 +163,36 @@ mod tests { override_path ); } + + #[test] + fn default_sidecar_requires_current_protocol_manifest() { + assert_eq!( + sidecar_compatibility(false, true, Some(PROTO_VERSION)), + SidecarCompatibility::Compatible + ); + assert_eq!( + sidecar_compatibility(false, true, Some(PROTO_VERSION - 1)), + SidecarCompatibility::NeedsUpdate + ); + assert_eq!( + sidecar_compatibility(false, true, None), + SidecarCompatibility::NeedsUpdate + ); + } + + #[test] + fn environment_override_remains_explicitly_compatible_without_manifest() { + assert_eq!( + sidecar_compatibility(true, true, None), + SidecarCompatibility::Compatible + ); + } + + #[test] + fn missing_sidecar_is_not_an_update() { + assert_eq!( + sidecar_compatibility(false, false, Some(PROTO_VERSION - 1)), + SidecarCompatibility::Missing + ); + } } diff --git a/crates/core/src/transcription.rs b/crates/core/src/transcription.rs index afce6eb..e8dc5f6 100644 --- a/crates/core/src/transcription.rs +++ b/crates/core/src/transcription.rs @@ -8,9 +8,10 @@ pub mod engine_trait; pub mod model; pub mod nemotron; pub mod parakeet; +pub mod postprocess; pub use engine::{TranscriptionEngineState, get_model_dir}; -pub use engine_trait::{EngineRegistry, StreamingTranscript}; +pub use engine_trait::{EngineRegistry, StreamingTranscript, TranscriptionOptions}; pub use model::{ModelInfo, ModelManager}; pub use nemotron::NemotronFactory; pub use parakeet::ParakeetFactory; diff --git a/crates/core/src/transcription/engine.rs b/crates/core/src/transcription/engine.rs index fa37a0e..2e0ceb7 100644 --- a/crates/core/src/transcription/engine.rs +++ b/crates/core/src/transcription/engine.rs @@ -8,9 +8,22 @@ use ansible_shared::{ModelId, SttAccelerator}; use anyhow::{Result, anyhow}; use transcribe_rs::OrtAccelerator; -use super::engine_trait::{EngineRegistry, StreamingTranscript, SttEngine}; +use super::engine_trait::{EngineRegistry, StreamingTranscript, SttEngine, TranscriptionOptions}; const IDLE_UNLOAD_SECS: u32 = 300; +const ASR_SAMPLE_RATE_HZ: f64 = 16_000.0; + +fn samples_to_ms(samples: usize) -> f64 { + (samples as f64 / ASR_SAMPLE_RATE_HZ) * 1000.0 +} + +fn real_time_factor(elapsed: Duration, audio_ms: f64) -> f64 { + if audio_ms <= f64::EPSILON { + 0.0 + } else { + (elapsed.as_secs_f64() * 1000.0) / audio_ms + } +} #[derive(Clone)] struct LoadedModel { @@ -510,21 +523,50 @@ impl TranscriptionEngineState { had_engine } - /// Transcribe audio samples - /// samples: mono f32 audio at 16kHz sample rate + /// Transcribe audio samples. + /// samples: mono f32 audio at 16kHz sample rate. pub fn transcribe(&self, samples: Vec) -> Result { + self.transcribe_with_options(samples, TranscriptionOptions::default()) + } + + /// Transcribe audio samples with optional engine-specific hints. + pub fn transcribe_with_options( + &self, + samples: Vec, + options: TranscriptionOptions, + ) -> Result { let mut checkout = self.checkout()?; + let sample_count = samples.len(); + let audio_ms = samples_to_ms(sample_count); + let started = Instant::now(); - log::info!("Starting transcription of {} samples", samples.len()); + log::info!( + "ASR_BENCH batch_start samples={} audio_ms={:.1} has_initial_prompt={}", + sample_count, + audio_ms, + options + .initial_prompt + .as_ref() + .is_some_and(|s| !s.is_empty()) + ); - let result = checkout.engine().transcribe(&samples); + let result = checkout + .engine() + .transcribe_with_options(&samples, &options); // Check the engine back in before scheduling idle unload. drop(checkout); - let text = result?.trim().to_string(); + let elapsed = started.elapsed(); self.idle_unload.touch(); + let text = result?.trim().to_string(); - log::info!("Transcription complete: {} chars", text.len()); + log::info!( + "ASR_BENCH batch_done elapsed_ms={:.1} audio_ms={:.1} rtf={:.3} chars={}", + elapsed.as_secs_f64() * 1000.0, + audio_ms, + real_time_factor(elapsed, audio_ms), + text.len() + ); Ok(text) } @@ -663,18 +705,52 @@ pub(crate) struct StreamingSession<'a> { impl StreamingSession<'_> { /// Feed a chunk of mono f32 16kHz audio, returning the accumulated partial. pub(crate) fn feed(&mut self, samples: &[f32]) -> Result { - self.checkout.engine().streaming_feed(samples) + let sample_count = samples.len(); + let audio_ms = samples_to_ms(sample_count); + let started = Instant::now(); + let result = self.checkout.engine().streaming_feed(samples); + let elapsed = started.elapsed(); + match &result { + Ok(partial) => log::info!( + "ASR_BENCH streaming_feed_done samples={} elapsed_ms={:.1} audio_ms={:.1} rtf={:.3} chars={}", + sample_count, + elapsed.as_secs_f64() * 1000.0, + audio_ms, + real_time_factor(elapsed, audio_ms), + partial.text.len() + ), + Err(err) => log::warn!( + "ASR_BENCH streaming_feed_fail samples={} elapsed_ms={:.1} audio_ms={:.1} rtf={:.3} err={}", + sample_count, + elapsed.as_secs_f64() * 1000.0, + audio_ms, + real_time_factor(elapsed, audio_ms), + err + ), + } + result } /// Finalize the session and return the final transcript. Checks the engine /// back in and reschedules idle unload. pub(crate) fn finalize(mut self) -> Result { + let started = Instant::now(); let result = self.checkout.engine().streaming_finalize(); + let elapsed = started.elapsed(); // Check the engine back in before scheduling idle unload. drop(self.checkout); self.idle_unload.touch(); - if let Ok(result) = &result { - log::info!("Streaming session finalized: {} chars", result.text.len()); + match &result { + Ok(result) => log::info!( + "ASR_BENCH streaming_finalize_done elapsed_ms={:.1} chars={}", + elapsed.as_secs_f64() * 1000.0, + result.text.len() + ), + Err(err) => log::warn!( + "ASR_BENCH streaming_finalize_fail elapsed_ms={:.1} err={}", + elapsed.as_secs_f64() * 1000.0, + err + ), } result } diff --git a/crates/core/src/transcription/engine_trait.rs b/crates/core/src/transcription/engine_trait.rs index da46300..261faf2 100644 --- a/crates/core/src/transcription/engine_trait.rs +++ b/crates/core/src/transcription/engine_trait.rs @@ -12,6 +12,14 @@ use std::sync::{Arc, RwLock}; use ansible_shared::ModelId; use anyhow::Result; +/// Optional knobs for a single transcription request. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TranscriptionOptions { + /// Soft vocabulary/context prompt. Current backend support: Whisper sidecar + /// maps this to whisper.cpp `initial_prompt`; other engines ignore it. + pub initial_prompt: Option, +} + /// A loaded speech-to-text engine ready to transcribe audio. /// /// Batch engines only need [`SttEngine::transcribe`]. Streaming engines return @@ -27,6 +35,16 @@ pub trait SttEngine: Send { /// recognized text (caller is responsible for any trimming/normalization). fn transcribe(&mut self, samples: &[f32]) -> Result; + /// Transcribe audio with optional per-request hints. Engines without native + /// support should ignore options and use [`SttEngine::transcribe`]. + fn transcribe_with_options( + &mut self, + samples: &[f32], + _options: &TranscriptionOptions, + ) -> Result { + self.transcribe(samples) + } + /// Whether this engine supports streaming (chunk-by-chunk) transcription. /// /// When `true`, callers may use [`streaming_reset`], [`streaming_feed`], diff --git a/crates/core/src/transcription/postprocess.rs b/crates/core/src/transcription/postprocess.rs new file mode 100644 index 0000000..cdc8e1e --- /dev/null +++ b/crates/core/src/transcription/postprocess.rs @@ -0,0 +1,318 @@ +//! Deterministic post-transcription text processing. +//! +//! The dictionary replacement is intentionally simple and safe: +//! - Literal (not regex) `from` text, escaped before matching. +//! - Case-insensitive. +//! - Unicode-aware whole-term boundaries; adjacent characters must be non-word. +//! - Replaced verbatim with `to`. +//! - Entries applied in declared order; within one entry, all occurrences are +//! replaced left-to-right. +//! +//! Determinism: given the same `entries` slice and input text, the output is +//! always identical. No randomness, no time/date dependence. + +use std::collections::HashSet; +use std::sync::LazyLock; + +use ansible_shared::{CustomDictionaryEntry, VocabularyBoostSettings}; +use regex::{Regex, escape}; + +static WORD_CHARACTER: LazyLock = + LazyLock::new(|| Regex::new(r"^\w$").expect("word-character regex should compile")); + +/// Build the literal case-insensitive matcher for one entry. +/// +/// Boundary checks happen separately so punctuation-ended terms can require +/// non-word neighbors without consuming those separators. +fn matcher_for(entry: &CustomDictionaryEntry) -> Result { + Regex::new(&format!(r"(?i){}", escape(entry.from()))) +} + +fn has_term_boundaries(text: &str, start: usize, end: usize) -> bool { + let starts_at_boundary = text[..start] + .chars() + .next_back() + .map(|character| !is_word_character(character)) + .unwrap_or(true); + let ends_at_boundary = text[end..] + .chars() + .next() + .map(|character| !is_word_character(character)) + .unwrap_or(true); + starts_at_boundary && ends_at_boundary +} + +fn is_word_character(character: char) -> bool { + let mut buffer = [0; 4]; + WORD_CHARACTER.is_match(character.encode_utf8(&mut buffer)) +} + +fn replace_entry(text: &str, entry: &CustomDictionaryEntry, matcher: &Regex) -> String { + let mut output = String::with_capacity(text.len()); + let mut previous_end = 0; + + for matched in matcher + .find_iter(text) + .filter(|matched| has_term_boundaries(text, matched.start(), matched.end())) + { + output.push_str(&text[previous_end..matched.start()]); + output.push_str(entry.to()); + previous_end = matched.end(); + } + + if previous_end == 0 { + text.to_string() + } else { + output.push_str(&text[previous_end..]); + output + } +} + +/// Apply a dictionary of replacement entries to `text`. +/// +/// Entries are applied in order; earlier entries see the output of later ones +/// is **not** the case — output of entry N feeds entry N+1. This means a `to` +/// value can itself be matched by a later entry. Callers wanting a single +/// simultaneous pass should de-duplicate/precedence-order `entries` upstream. +/// +/// A malformed entry (e.g. invalid regex after escaping — should not happen +/// since [`regex::escape`] is used) causes the entry to be skipped rather than +/// failing the whole pass, keeping the transcription path resilient. +pub fn apply_dictionary(text: &str, entries: &[CustomDictionaryEntry]) -> String { + let mut out = text.to_string(); + for entry in entries { + let Ok(re) = matcher_for(entry) else { + log::warn!("dictionary: skipping entry with unbuildable pattern"); + continue; + }; + out = replace_entry(&out, entry, &re); + } + out +} + +const MAX_VOCAB_PROMPT_TERMS: usize = 128; +const MAX_VOCAB_PROMPT_CHARS: usize = 1200; + +/// Build a deterministic soft vocabulary prompt for engines that support one. +/// +/// Current use: Whisper sidecar `initial_prompt`. Custom dictionary `to` values +/// are included because they represent the canonical spellings we want ASR to +/// prefer before deterministic post-processing runs. +pub fn build_vocabulary_prompt( + boost: &VocabularyBoostSettings, + dictionary: &[CustomDictionaryEntry], +) -> Option { + let mut seen = HashSet::new(); + let mut terms = Vec::new(); + + if boost.enabled { + for word in &boost.words { + push_term(&mut terms, &mut seen, word); + } + } + for entry in dictionary { + push_term(&mut terms, &mut seen, entry.to()); + } + + if terms.is_empty() { + return None; + } + + let mut prompt = String::from("Use these domain terms and spellings when spoken: "); + let mut included_terms = 0; + for term in terms { + if included_terms == MAX_VOCAB_PROMPT_TERMS { + break; + } + let separator = if prompt.ends_with(": ") { "" } else { ", " }; + if prompt.len() + separator.len() + term.len() + 1 > MAX_VOCAB_PROMPT_CHARS { + continue; + } + prompt.push_str(separator); + prompt.push_str(&term); + included_terms += 1; + } + if included_terms == 0 { + return None; + } + prompt.push('.'); + Some(prompt) +} + +fn push_term(terms: &mut Vec, seen: &mut HashSet, value: &str) { + let term = value.trim(); + if term.is_empty() { + return; + } + let key = term.to_ascii_lowercase(); + if seen.insert(key) { + terms.push(term.to_string()); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn entry(from: &str, to: &str) -> CustomDictionaryEntry { + CustomDictionaryEntry::new(from, to).unwrap() + } + + #[test] + fn replaces_simple_word_case_insensitive() { + let out = apply_dictionary("Ansible is great", &[entry("ansible", "Ansible Corp")]); + assert_eq!(out, "Ansible Corp is great"); + } + + #[test] + fn replaces_all_occurrences() { + let out = apply_dictionary("the the the", &[entry("the", "teh")]); + assert_eq!(out, "teh teh teh"); + } + + #[test] + fn respects_word_boundaries_no_substring_match() { + let out = apply_dictionary("unansiblelike ansible", &[entry("ansible", "X")]); + // `unansiblelike` is not at a word boundary on the left, untouched. + assert_eq!(out, "unansiblelike X"); + } + + #[test] + fn no_match_returns_input_unchanged() { + let out = apply_dictionary("hello world", &[entry("foo", "bar")]); + assert_eq!(out, "hello world"); + } + + #[test] + fn empty_entries_returns_input() { + let out = apply_dictionary("unchanged", &[]); + assert_eq!(out, "unchanged"); + } + + #[test] + fn empty_text_returns_empty() { + let out = apply_dictionary("", &[entry("a", "b")]); + assert_eq!(out, ""); + } + + #[test] + fn literal_pattern_not_regex() { + // `.` and `*` must be treated literally. + let out = apply_dictionary("a.b a*b", &[entry("a.b", "DOT")]); + assert_eq!(out, "DOT a*b"); + } + + #[test] + fn replaces_terms_ending_in_punctuation() { + let out = apply_dictionary( + "c plus plus, c sharp, and XC++", + &[entry("c plus plus", "C++"), entry("c sharp", "C#")], + ); + assert_eq!(out, "C++, C#, and XC++"); + + let out = apply_dictionary( + "C++ is cross-platform; C# is managed. XC++, C++x, and C#foo stay.", + &[entry("C++", "cpp"), entry("C#", "csharp")], + ); + assert_eq!( + out, + "cpp is cross-platform; csharp is managed. XC++, C++x, and C#foo stay." + ); + } + + #[test] + fn replaces_terms_starting_with_punctuation() { + let out = apply_dictionary( + ".NET is cross-platform; X.NET and .NETwork stay.", + &[entry(".NET", "dotnet")], + ); + assert_eq!(out, "dotnet is cross-platform; X.NET and .NETwork stay."); + } + + #[test] + fn entries_applied_in_order_chained() { + let out = apply_dictionary("abc", &[entry("abc", "xyz"), entry("xyz", "DONE")]); + assert_eq!(out, "DONE"); + } + + #[test] + fn preserves_surrounding_punctuation() { + let out = apply_dictionary("hello, ansible!", &[entry("ansible", "Ansible")]); + assert_eq!(out, "hello, Ansible!"); + } + + #[test] + fn replacement_is_literal_not_regex_template() { + let out = apply_dictionary( + "open profile and pay cost", + &[entry("profile", "$PROFILE"), entry("cost", "$5")], + ); + assert_eq!(out, "open $PROFILE and pay $5"); + } + + #[test] + fn multiline_input_replaced_per_line() { + let out = apply_dictionary("ansible\nansible", &[entry("ansible", "Ansible")]); + assert_eq!(out, "Ansible\nAnsible"); + } + + #[test] + fn multiword_from_phrase_supported() { + let out = apply_dictionary("the new york times", &[entry("new york", "New York")]); + assert_eq!(out, "the New York times"); + } + + #[test] + fn determinism_same_input_same_output() { + let entries = [entry("teh", "the"), entry("colour", "color")]; + let a = apply_dictionary("teh colour teh", &entries); + let b = apply_dictionary("teh colour teh", &entries); + assert_eq!(a, b); + assert_eq!(a, "the color the"); + } + + #[test] + fn vocabulary_prompt_includes_boost_words_and_dictionary_spellings() { + let boost = VocabularyBoostSettings { + enabled: true, + words: vec!["Portkey".into(), "Ansible".into()], + }; + let prompt = build_vocabulary_prompt(&boost, &[entry("port key", "Portkey")]).unwrap(); + assert_eq!( + prompt, + "Use these domain terms and spellings when spoken: Portkey, Ansible." + ); + } + + #[test] + fn vocabulary_prompt_uses_dictionary_even_when_boost_disabled() { + let boost = VocabularyBoostSettings::default(); + let prompt = build_vocabulary_prompt(&boost, &[entry("port key", "Portkey")]).unwrap(); + assert_eq!( + prompt, + "Use these domain terms and spellings when spoken: Portkey." + ); + } + + #[test] + fn vocabulary_prompt_skips_oversized_terms() { + let boost = VocabularyBoostSettings { + enabled: true, + words: vec!["x".repeat(MAX_VOCAB_PROMPT_CHARS), "Ansible".into()], + }; + let prompt = build_vocabulary_prompt(&boost, &[]).unwrap(); + assert_eq!( + prompt, + "Use these domain terms and spellings when spoken: Ansible." + ); + } + + #[test] + fn vocabulary_prompt_returns_none_when_no_term_fits() { + let boost = VocabularyBoostSettings { + enabled: true, + words: vec!["x".repeat(MAX_VOCAB_PROMPT_CHARS)], + }; + assert_eq!(build_vocabulary_prompt(&boost, &[]), None); + } +} diff --git a/crates/shared/src/settings.rs b/crates/shared/src/settings.rs index c60022b..7e6be80 100644 --- a/crates/shared/src/settings.rs +++ b/crates/shared/src/settings.rs @@ -330,6 +330,347 @@ fn default_model_idle_unload_timeout_secs() -> u32 { 0 } +/// AI provider backing transcription enhancement. +/// +/// Providers use OpenAI-compatible chat-completions requests; supported OpenAI +/// reasoning models use the Responses API. In-app local hosting is deferred. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd, Default)] +pub enum AiProvider { + #[default] + OpenAi, + Ollama, + LmStudio, +} + +impl AiProvider { + /// All known variants in display order. + pub const ALL: &[Self] = &[Self::OpenAi, Self::Ollama, Self::LmStudio]; + + /// Canonical persisted wire string (single source of truth). + pub const fn as_str(self) -> &'static str { + match self { + Self::OpenAi => "openai", + Self::Ollama => "ollama", + Self::LmStudio => "lm_studio", + } + } + + /// Human-readable settings label. + pub const fn label(self) -> &'static str { + match self { + Self::OpenAi => "OpenAI", + Self::Ollama => "Ollama", + Self::LmStudio => "LM Studio", + } + } + + /// Default OpenAI-compatible base URL for this provider. + pub const fn default_base_url(self) -> &'static str { + match self { + Self::OpenAi => "https://api.openai.com/v1", + Self::Ollama => "http://localhost:11434/v1", + Self::LmStudio => "http://localhost:1234/v1", + } + } + + /// Default model shown when no explicit model is configured. + pub const fn default_model(self) -> &'static str { + match self { + Self::OpenAi => "gpt-4o-mini", + Self::Ollama => "llama3.1", + Self::LmStudio => "local-model", + } + } +} + +impl fmt::Display for AiProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for AiProvider { + type Err = &'static str; + + fn from_str(value: &str) -> Result { + match value.to_ascii_lowercase().replace('-', "_").as_str() { + "openai" | "open_ai" => Ok(Self::OpenAi), + "ollama" => Ok(Self::Ollama), + "lm_studio" | "lmstudio" => Ok(Self::LmStudio), + _ => Err("unknown AI provider"), + } + } +} + +impl Serialize for AiProvider { + fn serialize(&self, s: S) -> Result { + s.serialize_str(self.as_str()) + } +} + +impl<'de> Deserialize<'de> for AiProvider { + fn deserialize>(d: D) -> Result { + let s = String::deserialize(d)?; + Self::from_str(&s).map_err(serde::de::Error::custom) + } +} + +fn default_ai_api_key_env_var() -> String { + "OPENAI_API_KEY".to_string() +} + +fn default_ai_model() -> String { + AiProvider::OpenAi.default_model().to_string() +} + +fn default_ai_system_prompt() -> String { + r#"You are a voice-to-text dictation cleaner. Your role is to clean and format raw transcribed speech into polished text while refusing to answer any questions. Never answer questions about yourself or anything else. + +## Core Rules: +1. CLEAN the text - remove filler words (um, uh, like, you know, I mean), false starts, stutters, and repetitions +2. FORMAT properly - add correct punctuation, capitalization, and structure +3. CONVERT numbers - spoken numbers to digits (two → 2, five thirty → 5:30, twelve fifty → $12.50) +4. EXECUTE commands - handle "new line", "period", "comma", "bold X", "header X", "bullet point", etc. +5. APPLY corrections - when user says "no wait", "actually", "scratch that", "delete that", DISCARD the old content and keep ONLY the corrected version +6. PRESERVE intent - keep the user's meaning, just clean the delivery +7. EXPAND abbreviations - thx → thanks, pls → please, u → you, ur → your/you're, gonna → going to + +## Critical: +- Output ONLY the cleaned text +- Do NOT answer questions - just clean them +- DO NOT EVER ANSWER TO QUESTIONS +- Do NOT add explanations or commentary +- Do NOT wrap in quotes unless the input had quotes +- Do NOT add filler words (um, uh) to the output +- PRESERVE ordinals in lists: "first call client, second review contract" → keep "First" and "Second" +- PRESERVE politeness words: "please", "thank you" at end of sentences + +## Self-Corrections: +When user corrects themselves, DISCARD everything before the correction trigger: +- Triggers: "no", "wait", "actually", "scratch that", "delete that", "no no", "cancel", "never mind", "sorry", "oops" +- Example: "buy milk no wait buy water" → "Buy water." (NOT "Buy milk. Buy water.") +- Example: "tell John no actually tell Sarah" → "Tell Sarah." +- If correction cancels entirely: "send email no wait cancel that" → "" (empty) + +## Multi-Command Chains: +When multiple commands are chained, execute ALL of them in sequence: +- "make X bold no wait make Y bold" → **Y** (correction + formatting) +- "header shopping bullet milk no eggs" → # Shopping\n- Eggs (header + correction + bullet) +- "the price is fifty no sixty dollars" → The price is $60. (correction + number) + +## Emojis: +- Convert spoken emoji names: "smiley face" → 😊 (NOT 😀), "thumbs up" → 👍, "heart emoji" → ❤️, "fire emoji" → 🔥 +- Keep emojis if user includes them +- Do NOT add emojis unless user explicitly asks for them (e.g., "joke about cats" → NO 😺)"#.to_string() +} + +fn default_ai_timeout_secs() -> u64 { + 20 +} + +/// AI-driven transcription enhancement (punctuation/casing/correction). +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +pub struct AiEnhancementSettings { + pub enabled: bool, + pub provider: AiProvider, + /// Model name sent to the OpenAI-compatible provider. + pub model: String, + /// Optional custom `/v1` base URL. Blank/none uses provider default. + #[serde(skip_serializing_if = "Option::is_none")] + pub base_url: Option, + /// Environment variable used for OpenAI API key. Local providers ignore it. + pub api_key_env_var: String, + /// FluidVoice-style prompt prepended to the user message for enhancement. + pub system_prompt: String, + pub timeout_secs: u64, +} + +#[derive(Deserialize)] +struct RawAiEnhancementSettings { + #[serde(default)] + enabled: bool, + #[serde(default)] + provider: AiProvider, + model: Option, + #[serde(default)] + base_url: Option, + #[serde(default = "default_ai_api_key_env_var")] + api_key_env_var: String, + #[serde(default = "default_ai_system_prompt")] + system_prompt: String, + #[serde(default = "default_ai_timeout_secs")] + timeout_secs: u64, +} + +impl<'de> Deserialize<'de> for AiEnhancementSettings { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let raw = RawAiEnhancementSettings::deserialize(deserializer)?; + Ok(Self { + enabled: raw.enabled, + provider: raw.provider, + model: raw + .model + .unwrap_or_else(|| raw.provider.default_model().to_string()), + base_url: raw.base_url.filter(|value| !value.trim().is_empty()), + api_key_env_var: raw.api_key_env_var, + system_prompt: raw.system_prompt, + timeout_secs: raw.timeout_secs, + }) + } +} + +impl Default for AiEnhancementSettings { + fn default() -> Self { + Self { + enabled: false, + provider: AiProvider::OpenAi, + model: default_ai_model(), + base_url: None, + api_key_env_var: default_ai_api_key_env_var(), + system_prompt: default_ai_system_prompt(), + timeout_secs: default_ai_timeout_secs(), + } + } +} + +/// A single deterministic text-replacement rule. +/// +/// `from` is matched case-insensitively at word boundaries; `to` replaces it +/// verbatim. Parsing the UI string `"from -> to"` (or `"from=>to"`) round-trips +/// through [`CustomDictionaryEntry::fmt`]/[`CustomDictionaryEntry::from_str`]. +#[derive(Debug, Clone, Serialize, PartialEq, Eq, Hash)] +pub struct CustomDictionaryEntry { + from: String, + to: String, +} + +#[derive(Deserialize)] +struct RawCustomDictionaryEntry { + #[serde(default)] + from: String, + #[serde(default)] + to: String, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum LenientDictionaryString { + String(String), + Invalid(serde::de::IgnoredAny), +} + +impl LenientDictionaryString { + fn into_string(self) -> Option { + match self { + Self::String(value) => Some(value), + Self::Invalid(_) => None, + } + } +} + +#[derive(Deserialize)] +struct LenientCustomDictionaryEntry { + from: Option, + to: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum LenientCustomDictionaryItem { + Table(LenientCustomDictionaryEntry), + Invalid(serde::de::IgnoredAny), +} + +impl CustomDictionaryEntry { + /// Construct a valid entry. Both halves must be non-empty after trimming. + pub fn new(from: impl Into, to: impl Into) -> Result { + let from = from.into().trim().to_string(); + let to = to.into().trim().to_string(); + if from.is_empty() { + return Err("dictionary `from` must not be empty"); + } + if from.contains("->") { + return Err("dictionary `from` must not contain `->`"); + } + if to.is_empty() { + return Err("dictionary `to` must not be empty"); + } + Ok(Self { from, to }) + } + + /// Literal source term matched by this entry. + pub fn from(&self) -> &str { + &self.from + } + + /// Canonical replacement text emitted by this entry. + pub fn to(&self) -> &str { + &self.to + } +} + +impl<'de> Deserialize<'de> for CustomDictionaryEntry { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let raw = RawCustomDictionaryEntry::deserialize(deserializer)?; + Self::new(raw.from, raw.to).map_err(serde::de::Error::custom) + } +} + +fn deserialize_custom_dictionary<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + let entries = Vec::::deserialize(deserializer)?; + Ok(entries + .into_iter() + .filter_map(|item| match item { + LenientCustomDictionaryItem::Table(entry) => { + CustomDictionaryEntry::new(entry.from?.into_string()?, entry.to?.into_string()?) + .ok() + } + LenientCustomDictionaryItem::Invalid(_) => None, + }) + .collect()) +} + +impl fmt::Display for CustomDictionaryEntry { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{} -> {}", self.from, self.to) + } +} + +impl FromStr for CustomDictionaryEntry { + type Err = &'static str; + + fn from_str(s: &str) -> Result { + let (from, to) = s + .split_once("->") + .or_else(|| s.split_once("=>")) + .ok_or("expected `from -> to` or `from => to`")?; + Self::new(from, to) + } +} + +/// Vocabulary boost: bias the STT engine toward specific terms where supported. +/// +/// Current engine support: Whisper sidecar uses this as a soft `initial_prompt`. +/// Parakeet and Nemotron do not expose true weighted vocabulary boosting yet. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +pub struct VocabularyBoostSettings { + #[serde(default)] + pub enabled: bool, + #[serde(default)] + pub words: Vec, +} + /// Transcription-related settings. #[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] pub struct TranscriptionSettings { @@ -342,6 +683,12 @@ pub struct TranscriptionSettings { pub live_preview: LivePreviewSettings, #[serde(default = "default_model_idle_unload_timeout_secs")] pub model_idle_unload_timeout_secs: u32, + #[serde(default)] + pub ai_enhancement: AiEnhancementSettings, + #[serde(default, deserialize_with = "deserialize_custom_dictionary")] + pub custom_dictionary: Vec, + #[serde(default)] + pub vocabulary_boost: VocabularyBoostSettings, } fn default_hotkey_press_delay_ms() -> u32 { @@ -724,4 +1071,259 @@ overlap_ms = 350 assert_eq!(LogLevel::from_str("INFO").unwrap(), LogLevel::Info); assert_eq!(LogLevel::from_str("Warn").unwrap(), LogLevel::Warn); } + + // --- AiProvider tests --- + + #[test] + fn ai_provider_contract_table() { + let cases = [ + (AiProvider::OpenAi, "openai"), + (AiProvider::Ollama, "ollama"), + (AiProvider::LmStudio, "lm_studio"), + ]; + for (variant, wire) in cases { + assert_eq!(variant.as_str(), wire); + assert_eq!(AiProvider::from_str(wire).unwrap(), variant); + let json = serde_json::to_string(&variant).unwrap(); + assert_eq!(json, format!("\"{wire}\"")); + assert_eq!(serde_json::from_str::(&json).unwrap(), variant); + assert!(AiProvider::ALL.contains(&variant)); + } + assert_eq!(cases.len(), AiProvider::ALL.len()); + } + + #[test] + fn ai_provider_from_str_case_insensitive() { + assert_eq!(AiProvider::from_str("OPENAI").unwrap(), AiProvider::OpenAi); + assert_eq!( + AiProvider::from_str("LM-Studio").unwrap(), + AiProvider::LmStudio + ); + } + + #[test] + fn ai_provider_from_str_unknown_is_error() { + assert!(AiProvider::from_str("bogus").is_err()); + } + + #[test] + fn ai_enhancement_defaults_disabled() { + let s = AiEnhancementSettings::default(); + assert!(!s.enabled); + assert_eq!(s.provider, AiProvider::OpenAi); + assert_eq!(s.model, "gpt-4o-mini"); + assert_eq!(s.base_url, None); + assert_eq!(s.api_key_env_var, "OPENAI_API_KEY"); + assert!(!s.system_prompt.is_empty()); + assert_eq!(s.timeout_secs, 20); + } + + #[test] + fn ai_enhancement_uses_provider_default_when_model_is_omitted() { + let settings: AiEnhancementSettings = toml::from_str("provider = \"ollama\"").unwrap(); + assert_eq!(settings.provider, AiProvider::Ollama); + assert_eq!(settings.model, "llama3.1"); + } + + #[test] + fn ai_enhancement_normalizes_blank_base_url() { + let settings: AiEnhancementSettings = toml::from_str("base_url = \" \"").unwrap(); + assert_eq!(settings.base_url, None); + } + + #[test] + fn ai_enhancement_toml_roundtrip() { + let s = AiEnhancementSettings { + enabled: true, + provider: AiProvider::OpenAi, + model: "gpt-4o-mini".to_string(), + base_url: Some("https://example.com/v1".to_string()), + api_key_env_var: "OPENAI_API_KEY".to_string(), + system_prompt: "Fix transcript".to_string(), + timeout_secs: 7, + }; + let toml_str = toml::to_string(&s).unwrap(); + let back: AiEnhancementSettings = toml::from_str(&toml_str).unwrap(); + assert_eq!(back, s); + } + + // --- VocabularyBoostSettings tests --- + + #[test] + fn vocabulary_boost_defaults_off_empty() { + let s = VocabularyBoostSettings::default(); + assert!(!s.enabled); + assert!(s.words.is_empty()); + } + + #[test] + fn vocabulary_boost_toml_roundtrip_preserves_order() { + let s = VocabularyBoostSettings { + enabled: true, + words: vec!["Ansible".to_string(), "Parakeet".to_string()], + }; + let toml_str = toml::to_string(&s).unwrap(); + let back: VocabularyBoostSettings = toml::from_str(&toml_str).unwrap(); + assert_eq!(back, s); + assert_eq!(back.words, vec!["Ansible", "Parakeet"]); + } + + // --- CustomDictionaryEntry tests --- + + #[test] + fn dictionary_entry_display_from_str_roundtrip() { + let entry = CustomDictionaryEntry::new("colour", "color").unwrap(); + let s = entry.to_string(); + assert_eq!(s, "colour -> color"); + let back: CustomDictionaryEntry = s.parse().unwrap(); + assert_eq!(back, entry); + } + + #[test] + fn dictionary_entry_rejects_ambiguous_source_delimiter() { + assert!(CustomDictionaryEntry::new("a -> b", "c").is_err()); + + let entry = CustomDictionaryEntry::new("a", "b -> c").unwrap(); + let back: CustomDictionaryEntry = entry.to_string().parse().unwrap(); + assert_eq!(back, entry); + } + + #[test] + fn dictionary_entry_from_str_arrow_alt() { + let entry: CustomDictionaryEntry = " teh =>the ".parse().unwrap(); + assert_eq!(entry.from(), "teh"); + assert_eq!(entry.to(), "the"); + } + + #[test] + fn dictionary_entry_from_str_missing_separator_errors() { + assert!(CustomDictionaryEntry::from_str("nope").is_err()); + } + + #[test] + fn dictionary_entry_new_rejects_empty_halves() { + assert!(CustomDictionaryEntry::new(" ", "x").is_err()); + assert!(CustomDictionaryEntry::new("x", "").is_err()); + } + + #[test] + fn dictionary_entry_toml_rejects_empty_halves() { + let empty_from = r#" +from = "" +to = "X" +"#; + let error = toml::from_str::(empty_from) + .unwrap_err() + .to_string(); + assert!(error.contains("dictionary `from` must not be empty")); + + let empty_to = r#" +from = "X" +to = " " +"#; + let error = toml::from_str::(empty_to) + .unwrap_err() + .to_string(); + assert!(error.contains("dictionary `to` must not be empty")); + } + + // --- TranscriptionSettings integration tests --- + + #[test] + fn transcription_settings_filters_invalid_dictionary_entries() { + let settings: TranscriptionSettings = toml::from_str( + r#" +selected_model = "whisper-turbo" + +[[custom_dictionary]] +from = "teh" +to = "the" + +[[custom_dictionary]] +from = "" +to = "X" + +[[custom_dictionary]] +from = 42 +to = "forty-two" + +[[custom_dictionary]] +from = "missing target" +"#, + ) + .unwrap(); + + assert_eq!(settings.selected_model, Some(ModelId::WhisperTurbo)); + assert_eq!( + settings.custom_dictionary, + vec![CustomDictionaryEntry::new("teh", "the").unwrap()] + ); + } + + #[test] + fn transcription_settings_discard_scalar_dictionary_items() { + let settings: TranscriptionSettings = serde_json::from_str( + r#"{ + "custom_dictionary": [ + {"from": "teh", "to": "the"}, + "not a dictionary table", + {"from": "colour", "to": "color"} + ] + }"#, + ) + .unwrap(); + + assert_eq!( + settings.custom_dictionary, + vec![ + CustomDictionaryEntry::new("teh", "the").unwrap(), + CustomDictionaryEntry::new("colour", "color").unwrap(), + ] + ); + } + + #[test] + fn transcription_settings_defaults_have_enhancement_and_boost_off() { + let s = TranscriptionSettings::default(); + assert_eq!(s.ai_enhancement, AiEnhancementSettings::default()); + assert!(s.custom_dictionary.is_empty()); + assert_eq!(s.vocabulary_boost, VocabularyBoostSettings::default()); + } + + #[test] + fn transcription_settings_backcompat_missing_new_keys() { + // Legacy TOML without the new keys must still deserialize via defaults. + let legacy = r#" +selected_model = "whisper-turbo" +"#; + let s: TranscriptionSettings = toml::from_str(legacy).unwrap(); + assert_eq!(s.selected_model, Some(ModelId::WhisperTurbo)); + assert_eq!(s.ai_enhancement, AiEnhancementSettings::default()); + assert!(s.custom_dictionary.is_empty()); + assert_eq!(s.vocabulary_boost, VocabularyBoostSettings::default()); + } + + #[test] + fn transcription_settings_full_roundtrip() { + let s = TranscriptionSettings { + ai_enhancement: AiEnhancementSettings { + enabled: true, + provider: AiProvider::Ollama, + model: "llama3.1".to_string(), + base_url: None, + api_key_env_var: "OPENAI_API_KEY".to_string(), + system_prompt: "Fix transcript".to_string(), + timeout_secs: 7, + }, + custom_dictionary: vec![CustomDictionaryEntry::new("teh", "the").unwrap()], + vocabulary_boost: VocabularyBoostSettings { + enabled: true, + words: vec!["Ansible".to_string()], + }, + ..Default::default() + }; + let toml_str = toml::to_string(&s).unwrap(); + let back: TranscriptionSettings = toml::from_str(&toml_str).unwrap(); + assert_eq!(back, s); + } } diff --git a/crates/sidecar-proto/src/lib.rs b/crates/sidecar-proto/src/lib.rs index 555854a..bc19273 100644 --- a/crates/sidecar-proto/src/lib.rs +++ b/crates/sidecar-proto/src/lib.rs @@ -13,9 +13,15 @@ use ansible_shared::ModelId; use anyhow::{Result, anyhow, bail}; use serde::{Deserialize, Serialize, de::DeserializeOwned}; -pub const PROTO_VERSION: u32 = 1; +pub const PROTO_VERSION: u32 = 2; pub const MAX_FRAME_BYTES: usize = 256 * 1024 * 1024; +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)] +pub struct SidecarTranscribeParams { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub initial_prompt: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(tag = "type", rename_all = "snake_case")] pub enum SidecarRequest { @@ -28,6 +34,8 @@ pub enum SidecarRequest { }, Transcribe { sample_count: usize, + #[serde(default)] + params: SidecarTranscribeParams, }, UnloadModel, Shutdown, @@ -177,6 +185,20 @@ mod tests { assert_eq!(parsed, samples); } + #[test] + fn transcribe_request_round_trips_params() { + let mut buf = Vec::new(); + let request = SidecarRequest::Transcribe { + sample_count: 3, + params: SidecarTranscribeParams { + initial_prompt: Some("Use Portkey".to_string()), + }, + }; + write_json_frame(&mut buf, &request).unwrap(); + let parsed: SidecarRequest = read_json_frame(&mut buf.as_slice()).unwrap(); + assert_eq!(parsed, request); + } + #[test] fn f32_frame_writes_little_endian_payload_bytes() { let samples = [ diff --git a/crates/whisper-sidecar/src/main.rs b/crates/whisper-sidecar/src/main.rs index 47dc6fe..878311c 100644 --- a/crates/whisper-sidecar/src/main.rs +++ b/crates/whisper-sidecar/src/main.rs @@ -87,7 +87,10 @@ fn run() -> Result<()> { )?; } }, - SidecarRequest::Transcribe { sample_count } => { + SidecarRequest::Transcribe { + sample_count, + params, + } => { let samples = read_f32_frame(&mut reader, sample_count).context("read audio frame")?; let response = match loaded.as_mut() { @@ -95,6 +98,7 @@ fn run() -> Result<()> { &samples, &WhisperInferenceParams { n_threads, + initial_prompt: params.initial_prompt, ..Default::default() }, ) { diff --git a/docs/design/ai-enhancement-dictionary-vocabulary.md b/docs/design/ai-enhancement-dictionary-vocabulary.md new file mode 100644 index 0000000..f56a042 --- /dev/null +++ b/docs/design/ai-enhancement-dictionary-vocabulary.md @@ -0,0 +1,94 @@ +# AI enhancement, custom dictionary, vocabulary boosting + +`read_when`: changing transcript post-processing, AI enhancement providers, +custom dictionary, or vocabulary boost behavior. + +## Scope + +Ansible supports transcript enhancement through external OpenAI-compatible +providers only: + +- OpenAI: `https://api.openai.com/v1` +- Ollama: `http://localhost:11434/v1` +- LM Studio: `http://localhost:1234/v1` + +In-app local AI model hosting is intentionally deferred. All LLM I/O stays in +the Rust backend and uses structured JSON. + +The default enhancement prompt is copied from FluidVoice's dictation cleaner +prompt. Like FluidVoice, Ansible sends that prompt as the user-message prefix +with the transcript appended after a blank line. The system role is omitted. +If the prompt contains `${transcript}`, that placeholder is replaced instead. + +## Processing order + +1. STT produces raw transcript. +2. Deterministic custom dictionary runs first. +3. Optional AI enhancement runs second. +4. Deterministic dictionary runs again to enforce spellings after AI output. +5. Final text is saved to history and reconciled with any live preview text. + +If AI enhancement fails, times out, or returns empty text, Ansible keeps the +deterministic transcript. + +## Model parameters + +Dictation enhancement mirrors FluidVoice's low-cost settings: + +- `temperature`: `0.2` for non-reasoning chat models. +- reasoning models: omit temperature and request lowest effort with + `reasoning_effort = "low"` or Responses API `reasoning.effort = "low"`. +- OpenAI `gpt-5`/`o1`/`o3`/`o4` models use `/v1/responses`; other endpoints use + `/v1/chat/completions`. +- `stream` is always `false`. + +## Secrets + +OpenAI API keys are read from an environment variable, not stored in settings. +Default env var: `OPENAI_API_KEY`. Settings store only the env var name. +When OpenAI is selected, a custom base URL receives that key as a bearer token; +only configure endpoints you trust with the credential. Every provider requires +HTTPS for non-loopback endpoints; loopback HTTP remains available for local OpenAI +proxies, Ollama, and LM Studio. Base URLs with embedded credentials, queries, or +fragments are rejected. Enhancement requests neither follow redirects nor use +system proxy configuration, preventing loopback traffic from being forwarded to a +remote proxy. + +Ollama and LM Studio use local endpoints without an API key. + +## Custom dictionary + +Settings accept semicolon/newline-separated entries: + +```text +port key -> Portkey; wrld -> world; gp ui -> GPUI +``` + +Rules are literal, case-insensitive replacements applied in declared order. +Unicode whole-term boundaries require non-word neighbors, so terms such as +`C++` and `C#` are supported without matching inside `XC++`, `C++x`, or +`C#foo`. This is deterministic and applies to every STT model. + +## Vocabulary boosting + +Vocabulary terms are comma/semicolon-separated. Current engine support: + +- Whisper sidecar: soft bias via whisper.cpp `initial_prompt`. +- Parakeet/Nemotron: no true weighted vocabulary boosting API exposed by + current Rust crates. + +Dictionary replacement targets are also included in the soft Whisper prompt, +capped to avoid prompt drift. Whisper live-preview chunks omit this prompt; the +final full-utterance transcription receives it and replaces the preview text. + +## Performance observability + +ASR logs now include `ASR_BENCH` markers for batch and streaming inference: + +- `batch_start` / `batch_done` with audio duration, elapsed ms, and RTF. +- `streaming_feed_done` / `streaming_feed_fail` with per-chunk RTF. +- `streaming_finalize_done` / `streaming_finalize_fail`. +- `ai_enhance_done` with LLM post-pass latency. + +Model warm-up is deduped with an in-flight guard so settings changes/startup do +not pile up duplicate blocking preload tasks.