From 376f2c76646d36768fe928b6395faa83e3707e85 Mon Sep 17 00:00:00 2001 From: Uzair Aftab Date: Mon, 23 Mar 2026 22:12:57 +0100 Subject: [PATCH] feat(core): stream llm formatting output --- src/App.zig | 190 +++++++++++++++++++++++++++++++++++----- src/llm/LlamaClient.zig | 30 +++++++ 2 files changed, 199 insertions(+), 21 deletions(-) diff --git a/src/App.zig b/src/App.zig index b1bbbed..77444de 100644 --- a/src/App.zig +++ b/src/App.zig @@ -13,6 +13,16 @@ const LlamaClient = if (has_llm) @import("llm/LlamaClient.zig") else void; const App = @This(); +const llm_model_candidates = [_][]const u8{ + "llama-3.2-1b-q4_k_m.gguf", + "llama-3.2-3b-q4_k_m.gguf", + "qwen2.5-0.5b-instruct-q4_k_m.gguf", + "qwen2.5-0.5b-q4_k_m.gguf", + "qwen2.5-1.5b-instruct-q4_k_m.gguf", + "qwen2.5-1.5b-q4_k_m.gguf", + "llama-3.2-1b-q4.gguf", +}; + allocator: std.mem.Allocator, config: c_api.RuntimeConfig, status: c_api.Status, @@ -22,6 +32,7 @@ live_transcriber: ?Transcriber, audio: ?AudioCapture, llama: if (has_llm) ?LlamaClient else void, log_store: LogStore, +custom_prompt: ?[]u8 = null, // Live transcription state live_thread: ?std.Thread = null, @@ -53,9 +64,73 @@ pub fn deinit(self: *App) void { if (has_llm) { if (self.llama) |*l| l.deinit(); } + if (self.custom_prompt) |cp| self.allocator.free(cp); self.log_store.deinit(); } +fn pathExists(path: []const u8) bool { + std.fs.accessAbsolute(path, .{}) catch { + return false; + }; + return true; +} + +fn resolveLlmModelPath(self: *App, path_buf: *[std.fs.max_path_bytes]u8) ![]const u8 { + std.debug.assert(has_llm); + + const configured_path = self.config.getLlmModelPath(); + if (std.fs.path.isAbsolute(configured_path) and pathExists(configured_path)) { + return configured_path; + } + + const models_dir = self.config.getModelsDir(); + if (!std.fs.path.isAbsolute(models_dir)) { + std.log.err("Models directory is not absolute: {s}", .{models_dir}); + self.notifyError("Models directory path is invalid."); + return error.InvalidModelsDirectoryPath; + } + + for (llm_model_candidates) |filename| { + const candidate_path = std.fmt.bufPrint(path_buf, "{s}/{s}", .{ models_dir, filename }) catch { + self.notifyError("LLM model path too long."); + return error.PathTooLong; + }; + if (pathExists(candidate_path)) { + std.log.info("Using LLM model: {s}", .{candidate_path}); + return candidate_path; + } + } + + if (std.fs.path.isAbsolute(configured_path)) { + std.log.warn("Configured LLM model not found: {s}", .{configured_path}); + } else { + std.log.warn("Configured LLM model path is not absolute: {s}", .{configured_path}); + } + self.notifyError("LLM model not found. Download a llama or qwen GGUF model."); + return error.ModelNotFound; +} + +fn ensureLlamaLoaded(self: *App) !void { + if (!has_llm or self.llama != null) { + return; + } + + var llm_model_path_buf: [std.fs.max_path_bytes]u8 = undefined; + const llm_model_path = try self.resolveLlmModelPath(&llm_model_path_buf); + + self.llama = LlamaClient.init(self.allocator, .{ + .model_path = llm_model_path, + .n_ctx = 512, + .n_threads = 4, + }) catch |err| { + std.log.err("Failed to load LLM model ({s}): {}", .{ llm_model_path, err }); + self.notifyError("Failed to load LLM model."); + return err; + }; + + std.log.info("LLM model loaded: {s}", .{llm_model_path}); +} + pub fn loadModel(self: *App, size: c_api.ModelSize) !void { const models_dir = self.config.getModelsDir(); const model_name = size.toModelName(); @@ -336,12 +411,7 @@ fn transcribeTail(self: *App, options: c_api.TranscribeOptions) !void { std.log.info("Transcription complete, text length: {d}", .{final_text.len}); - if (options.use_llm_formatting) { - try self.formatText(final_text, options.tone, self.config.on_transcript, self.config.userdata); - } else { - self.notifyTranscript(final_text, true); - self.setStatus(.ready); - } + try self.finalizeTranscript(final_text, options); } pub fn isRecording(self: *App) bool { @@ -377,12 +447,7 @@ pub fn transcribe(self: *App, options: c_api.TranscribeOptions) !void { std.log.info("Transcription complete, text length: {d}", .{raw_text.len}); - if (options.use_llm_formatting) { - try self.formatText(raw_text, options.tone, self.config.on_transcript, self.config.userdata); - } else { - self.notifyTranscript(raw_text, true); - self.setStatus(.ready); - } + try self.finalizeTranscript(raw_text, options); } pub fn formatText( @@ -401,13 +466,7 @@ pub fn formatText( return; } - if (self.llama == null) { - self.llama = try LlamaClient.init(self.allocator, .{ - .model_path = self.config.getLlmModelPath(), - .n_ctx = 512, - .n_threads = 4, - }); - } + try self.ensureLlamaLoaded(); self.setStatus(.formatting); @@ -431,7 +490,71 @@ pub fn formatText( self.setStatus(.ready); } +/// Shared finalization for transcribe() and transcribeTail(). +/// Sends raw text as a preview, runs LLM formatting if enabled, +/// delivers the final result, and persists both texts to the log store. +fn finalizeTranscript(self: *App, raw_text: []const u8, options: c_api.TranscribeOptions) !void { + if (!options.use_llm_formatting) { + self.notifyTranscript(raw_text, true); + self.log_store.appendTranscript(self.allocator, raw_text, null) catch |err| { + std.log.warn("Failed to persist transcript: {}", .{err}); + }; + self.setStatus(.ready); + return; + } + + // Send raw text as a non-final preview so the UI shows something immediately + self.notifyTranscript(raw_text, false); + + if (!has_llm) { + // iOS: no LLM support yet, finalize with raw text + self.notifyTranscript(raw_text, true); + self.log_store.appendTranscript(self.allocator, raw_text, null) catch |err| { + std.log.warn("Failed to persist transcript: {}", .{err}); + }; + self.setStatus(.ready); + return; + } + + try self.ensureLlamaLoaded(); + + self.setStatus(.formatting); + + const prompt = try self.buildFormattingPrompt(raw_text, options.tone); + defer self.allocator.free(prompt); + + var stream_context = LlmStreamContext{ .app = self }; + + const formatted = self.llama.?.generateStreaming( + prompt, + 256, + onLlmFormattingPartial, + &stream_context, + ) catch |err| { + std.log.warn("LLM formatting failed: {}, returning raw text", .{err}); + self.notifyTranscript(raw_text, true); + self.log_store.appendTranscript(self.allocator, raw_text, null) catch {}; + self.setStatus(.ready); + return; + }; + defer self.allocator.free(formatted); + + self.notifyTranscript(formatted, true); + self.log_store.appendTranscript(self.allocator, raw_text, formatted) catch |err| { + std.log.warn("Failed to persist transcript: {}", .{err}); + }; + self.setStatus(.ready); +} + fn buildFormattingPrompt(self: *App, input: []const u8, tone: c_api.Tone) ![]u8 { + if (self.custom_prompt) |cp| { + return try std.fmt.allocPrint( + self.allocator, + "{s}\n\nInput: {s}\n\nOutput:", + .{ cp, input }, + ); + } + const base_prompt = \\Clean up this transcribed speech. Remove filler words (um, uh, like, you know). \\Fix grammar and punctuation. Keep the meaning intact.{s} @@ -460,14 +583,22 @@ pub fn getAudioLevel(self: *App) f32 { pub fn writeSettings(self: *App, settings: c_api.Settings) !void { std.debug.assert(self.config.getConfigDomain().len > 0); try SettingsStore.write(self.config, settings); + + if (self.custom_prompt) |cp| { + self.allocator.free(cp); + self.custom_prompt = null; + } + if (settings.getCustomPrompt()) |prompt| { + self.custom_prompt = try self.allocator.dupe(u8, prompt); + } } pub fn clearTranscriptLog(self: *App) !void { try self.log_store.clear(); } -pub fn appendTranscriptLog(self: *App, transcript: []const u8) !void { - try self.log_store.appendTranscript(self.allocator, transcript); +pub fn appendTranscriptLog(self: *App, transcript: []const u8, formatted_text: ?[]const u8) !void { + try self.log_store.appendTranscript(self.allocator, transcript, formatted_text); } pub fn getTranscriptLogRecentJson(self: *App, limit: usize) !c_api.String { @@ -487,6 +618,10 @@ pub fn getTranscriptLogRecentJson(self: *App, limit: usize) !c_api.String { try json_buffer.writer(self.allocator).print("{d}", .{entry.created_at_unix_ms}); try json_buffer.appendSlice(self.allocator, ",\"text\":"); try appendJsonEscapedString(&json_buffer, self.allocator, entry.text); + if (entry.formatted_text) |ft| { + try json_buffer.appendSlice(self.allocator, ",\"formatted_text\":"); + try appendJsonEscapedString(&json_buffer, self.allocator, ft); + } try json_buffer.append(self.allocator, '}'); } try json_buffer.append(self.allocator, ']'); @@ -520,6 +655,19 @@ fn appendJsonEscapedString( try buffer.append(allocator, '"'); } +const LlmStreamContext = struct { + app: *App, +}; + +fn onLlmFormattingPartial(userdata: ?*anyopaque, partial_text: []const u8) void { + if (partial_text.len == 0) { + return; + } + const ptr = userdata orelse return; + const context: *LlmStreamContext = @ptrCast(@alignCast(ptr)); + context.app.notifyTranscript(partial_text, false); +} + fn setStatus(self: *App, status: c_api.Status) void { self.status = status; if (self.config.on_status_change) |cb| { diff --git a/src/llm/LlamaClient.zig b/src/llm/LlamaClient.zig index 0d45895..9cfddf4 100644 --- a/src/llm/LlamaClient.zig +++ b/src/llm/LlamaClient.zig @@ -7,6 +7,8 @@ const c = @cImport({ const LlamaClient = @This(); +pub const StreamSink = *const fn (userdata: ?*anyopaque, partial_text: []const u8) void; + allocator: std.mem.Allocator, model: ?*c.llama_model, ctx: ?*c.llama_context, @@ -77,6 +79,16 @@ fn createSampler() *c.llama_sampler { } pub fn generate(self: *LlamaClient, prompt: []const u8, max_tokens: u32) ![]u8 { + return self.generateStreaming(prompt, max_tokens, null, null); +} + +pub fn generateStreaming( + self: *LlamaClient, + prompt: []const u8, + max_tokens: u32, + stream_sink: ?StreamSink, + stream_userdata: ?*anyopaque, +) ![]u8 { _ = self.model orelse return error.NoModel; const ctx = self.ctx orelse return error.NoContext; const sampler = self.sampler orelse return error.NoSampler; @@ -106,6 +118,9 @@ pub fn generate(self: *LlamaClient, prompt: []const u8, max_tokens: u32) ![]u8 { // Generate tokens var output: std.ArrayListUnmanaged(u8) = .empty; errdefer output.deinit(self.allocator); + const stream_interval_ns: i128 = 40 * std.time.ns_per_ms; + var last_stream_ns: i128 = std.time.nanoTimestamp() - stream_interval_ns; + var last_streamed_len: usize = 0; var n_cur: i32 = n_prompt; const n_ctx: i32 = @intCast(c.llama_n_ctx(ctx)); @@ -123,6 +138,15 @@ pub fn generate(self: *LlamaClient, prompt: []const u8, max_tokens: u32) ![]u8 { const n = c.llama_token_to_piece(vocab, new_token, &buf, buf.len, 0, true); if (n > 0) { try output.appendSlice(self.allocator, buf[0..@intCast(n)]); + + if (stream_sink) |sink| { + const now_ns = std.time.nanoTimestamp(); + if (output.items.len > last_streamed_len and now_ns - last_stream_ns >= stream_interval_ns) { + sink(stream_userdata, output.items); + last_stream_ns = now_ns; + last_streamed_len = output.items.len; + } + } } // Decode next token @@ -135,6 +159,12 @@ pub fn generate(self: *LlamaClient, prompt: []const u8, max_tokens: u32) ![]u8 { n_cur += 1; } + if (stream_sink) |sink| { + if (output.items.len > last_streamed_len) { + sink(stream_userdata, output.items); + } + } + // Clear memory for next generation c.llama_memory_clear(c.llama_get_memory(ctx), true);