From 5ab110ed0b8f0d260681f294e03b36d10d4465d6 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Fri, 24 Jul 2026 18:01:53 +0200 Subject: [PATCH] feat: better chat support for thinking --- PLAN.md | 43 +++-- .../20260724230000_create_messages/down.sql | 1 + .../20260724230000_create_messages/up.sql | 10 ++ src/app.rs | 168 +++++++++++++++--- src/app/view.rs | 50 +++++- src/database.rs | 132 +++++++++++++- src/engine.rs | 93 +++++++++- src/engine/tokenizer.rs | 20 ++- src/schema.rs | 14 +- 9 files changed, 463 insertions(+), 68 deletions(-) create mode 100644 migrations/20260724230000_create_messages/down.sql create mode 100644 migrations/20260724230000_create_messages/up.sql diff --git a/PLAN.md b/PLAN.md index d16b45b..22d528c 100644 --- a/PLAN.md +++ b/PLAN.md @@ -13,11 +13,11 @@ sessions, tools, and turn orchestration. | Area | Status | Available now | Still open | | --- | --- | --- | --- | | App shell | Implemented | Native multi-window Iced app, Codex-inspired layout, native Application/Edit/Window menus, Preferences panel, Model Manager, and streaming composer | Native app release packaging | -| Projects and sessions | Implemented for metadata | Native folder picker, project-name dialog, create/select/delete projects and sessions | Transcripts, KV payloads, rename/archive, and model binding | -| Persistence | Implemented for metadata and preferences | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults and CRUD tests | Messages, downloaded-artifact metadata, and session-runtime migrations | +| Projects and sessions | Implemented with transcripts | Native folder picker, project-name dialog, create/select/delete projects and sessions, and durable structured chat history | KV payloads, rename/archive, and model binding | +| Persistence | Implemented for metadata, preferences, and chat | SQLite in Application Support, Diesel ORM, embedded migrations, constrained defaults, streamed message updates, and reopen tests | Downloaded-artifact metadata and session-runtime migrations | | Model runtime | DeepSeek Flash vertical implemented | Rust-owned mmap/model/session graph, reused Metal kernels, 2K compressed-attention generation, DS4 sampling, lazy background loading, cancellation, and idle unloading | Ratio-4 sparse indexer beyond 2K, DSpark/SSD/steering, GLM/Pro, KV checkpoint persistence, and shared acquisition for HTTP | | Local endpoint | Not started | Nothing listening | OpenAI-compatible HTTP and streaming surfaces | -| Agent | Basic local chat implemented | Multi-turn role rendering, streamed DeepSeek Flash text, composer send/stop, and in-memory transcript | Durable transcripts, reasoning presentation, tools, approvals, compaction, statistics, and KV-backed resume | +| Agent | Basic local chat implemented | Multi-turn role rendering, structured reasoning disclosure, streamed DeepSeek Flash text, composer send/stop, and durable transcript rehydration | Tools, approvals, compaction, statistics, and KV-backed resume | | A2UI | Future extension | bDS2 provides the reference structured render-tool contract | Native inline surfaces for local chat after core chat and tools are stable | | Dev brain | Not started | No vault access | Obsidian vault selection, durable access, and agent knowledge/memory tools | | Release | Not started | Development binary builds and tests | `.app` packaging, signing, entitlements, and notarization | @@ -26,9 +26,10 @@ The application manages durable project/session metadata and model preferences. Its Model Manager downloads, validates, and deletes managed GGUF artifacts without blocking the UI. Local chat now lazily acquires the Rust DeepSeek Flash executor, streams generated text, supports cancellation and multi-turn role -replay, and unloads model resources after the configured idle timeout. The +replay, persists reasoning and answers while they stream, rehydrates selected +sessions, and unloads model resources after the configured idle timeout. The current Rust graph supports DS4 compressed attention through 2K context; the sparse ratio-4 indexer, -durable transcripts/KV checkpoints, HTTP service, and agent tools remain open. +KV checkpoints, HTTP service, and agent tools remain open. ## Phase 0 — project and session shell @@ -292,8 +293,8 @@ or KV frontier. ## Phase 3 — durable agent sessions -Status: **partially open**. Project/session metadata rows exist; transcripts, -model binding, and KV payload persistence are not implemented. +Status: **partially implemented**. Project/session metadata and structured +transcripts are durable; model binding and KV payload persistence remain open. Extend each metadata-only session with transcript and shared-checkpoint metadata: @@ -308,10 +309,11 @@ Application Support/DS4Server.rfc1437.de/ .kv ``` -- Add migrated message/transcript tables and an optional current-checkpoint key - to SQLite. Keep large engine-owned KV payloads only in the process-wide Phase - 1 store, written with atomic replacement; do not introduce separate - agent-only or HTTP-only cache formats. +- **Partially implemented:** migrated message tables persist user text, + reasoning state, and assistant output while generation streams. Add the + optional current-checkpoint key to SQLite. Keep large engine-owned KV + payloads only in the process-wide Phase 1 store, written with atomic + replacement; do not introduce separate agent-only or HTTP-only cache formats. - Record model identity, context size, rendered token history, title, timestamps, and working directory. - Restore compatible KV immediately. If KV is absent or incompatible, rebuild @@ -329,13 +331,20 @@ the same conversation without prefill when its KV payload is compatible. ## Phase 4 — agent chat and tools Status: **basic local chat implemented**. The conversation area streams -multi-turn DeepSeek Flash output and supports Stop; durable messages, tool -turns, and the coding-agent runtime remain open. +multi-turn DeepSeek Flash output, persists and rehydrates structured reasoning +and answers, and supports Stop; tool turns and the coding-agent runtime remain +open. -1. **Partially implemented:** build the classic chat surface: transcript, reasoning disclosure, tool-call - cards, composer, queued input, stop button, prefill progress, and generation - statistics. Local turns must call the same session registry and KV path used - by the HTTP endpoint. +Every chat feature must persist its semantic state and rehydrate it when a +session is reopened in the same implementation slice. In-memory-only chat +features are incomplete by definition; disclosure state and other purely +ephemeral presentation preferences are the only exception. + +1. **Partially implemented:** build the classic chat surface. Transcript, + structured reasoning disclosure, composer, streamed output, and Stop are + complete; tool-call cards, queued input, prefill progress, and generation + statistics remain. Local turns must call the same session registry and KV + path used by the HTTP endpoint. 2. Port the native agent turn loop from `ds4_agent.c`, including model-specific system prompts and DeepSeek DSML/GLM tool-call parsing. 3. Implement the local tool set with the project directory as its boundary: diff --git a/migrations/20260724230000_create_messages/down.sql b/migrations/20260724230000_create_messages/down.sql new file mode 100644 index 0000000..f8f44a8 --- /dev/null +++ b/migrations/20260724230000_create_messages/down.sql @@ -0,0 +1 @@ +DROP TABLE messages; diff --git a/migrations/20260724230000_create_messages/up.sql b/migrations/20260724230000_create_messages/up.sql new file mode 100644 index 0000000..36d6ea0 --- /dev/null +++ b/migrations/20260724230000_create_messages/up.sql @@ -0,0 +1,10 @@ +CREATE TABLE messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, + session_id INTEGER NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + user BOOLEAN NOT NULL CHECK (user IN (0, 1)), + reasoning TEXT, + reasoning_complete BOOLEAN NOT NULL DEFAULT 1 CHECK (reasoning_complete IN (0, 1)), + content TEXT NOT NULL +); + +CREATE INDEX messages_session_id ON messages(session_id); diff --git a/src/app.rs b/src/app.rs index 9001bb7..b5e0f13 100644 --- a/src/app.rs +++ b/src/app.rs @@ -2,7 +2,7 @@ mod view; pub(crate) use view::app_theme; -use crate::database::{AppPreferences, Database, ProjectWithSessions}; +use crate::database::{AppPreferences, Database, ProjectWithSessions, StoredMessage}; #[cfg(target_os = "macos")] use crate::engine::{ChatTurn, Generator}; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; @@ -410,11 +410,38 @@ pub(crate) struct App { #[derive(Clone, Debug)] pub(super) struct ChatMessage { + id: i32, pub(super) user: bool, - reasoning: bool, + pub(super) reasoning: Option, + reasoning_complete: bool, + pub(super) reasoning_open: bool, pub(super) content: String, } +impl ChatMessage { + fn append(&mut self, reasoning: bool, chunk: &str) { + if reasoning { + self.reasoning.get_or_insert_default().push_str(chunk); + } else { + self.reasoning_complete = true; + self.content.push_str(chunk); + } + } +} + +impl From for ChatMessage { + fn from(message: StoredMessage) -> Self { + Self { + id: message.id, + user: message.user, + reasoning: message.reasoning, + reasoning_complete: message.reasoning_complete, + reasoning_open: false, + content: message.content, + } + } +} + #[cfg(target_os = "macos")] struct GenerationWorker { commands: mpsc::Sender, @@ -436,7 +463,7 @@ enum GenerationCommand { #[cfg(target_os = "macos")] enum GenerationEvent { Loading, - Chunk(String), + Chunk { reasoning: bool, content: String }, Finished(Result<(), String>), } @@ -519,6 +546,7 @@ pub(crate) enum Message { StopModelDownload, DownloadProgressTick, ComposerChanged(String), + ToggleReasoning(usize), SubmitPrompt, StopGeneration, GenerationTick, @@ -870,6 +898,13 @@ impl App { } Message::DownloadProgressTick => self.update_download_progress(), Message::ComposerChanged(value) => self.composer = value, + Message::ToggleReasoning(index) => { + if let Some(message) = self.conversation.get_mut(index) + && message.reasoning.is_some() + { + message.reasoning_open = !message.reasoning_open; + } + } Message::SubmitPrompt => self.start_generation(), Message::StopGeneration => { #[cfg(target_os = "macos")] @@ -916,6 +951,8 @@ impl App { } self.selected_project = Some(project_id); self.selected_session = None; + self.conversation.clear(); + self.composer.clear(); self.error = None; } Message::DeleteProject(project_id) => { @@ -930,6 +967,8 @@ impl App { if self.selected_project == Some(project_id) { self.selected_project = None; self.selected_session = None; + self.conversation.clear(); + self.composer.clear(); } self.reload_projects(); } @@ -944,13 +983,24 @@ impl App { Some("Stop the active generation before changing sessions.".into()); return Task::none(); } - if self.selected_session != Some(session_id) { - self.conversation.clear(); - self.composer.clear(); + if self.selected_session == Some(session_id) { + return Task::none(); + } + let Some(database) = &mut self.database else { + return Task::none(); + }; + match database.load_messages(session_id) { + Ok(messages) => { + self.conversation = messages.into_iter().map(ChatMessage::from).collect(); + self.composer.clear(); + self.selected_project = Some(project_id); + self.selected_session = Some(session_id); + self.error = None; + } + Err(error) => { + self.error = Some(format!("Could not load the chat session: {error}")); + } } - self.selected_project = Some(project_id); - self.selected_session = Some(session_id); - self.error = None; } Message::DeleteSession(session_id) => { if self.generating && self.selected_session == Some(session_id) { @@ -963,6 +1013,8 @@ impl App { Ok(()) => { if self.selected_session == Some(session_id) { self.selected_session = None; + self.conversation.clear(); + self.composer.clear(); } self.reload_projects(); } @@ -1223,6 +1275,8 @@ impl App { match database.create_session(project_id, &title) { Ok(session) => { self.selected_session = Some(session.id); + self.conversation.clear(); + self.composer.clear(); self.error = None; self.reload_projects(); } @@ -1333,20 +1387,25 @@ impl App { } }; let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct; + let session_id = self + .selected_session + .expect("a selected session was checked"); #[cfg(target_os = "macos")] let mut messages = self .conversation .iter() .map(|message| ChatTurn { user: message.user, - reasoning: message.reasoning, + reasoning: message.reasoning.clone(), + reasoning_complete: message.reasoning_complete, content: message.content.clone(), }) .collect::>(); #[cfg(target_os = "macos")] messages.push(ChatTurn { user: true, - reasoning: false, + reasoning: None, + reasoning_complete: true, content: prompt.clone(), }); @@ -1361,6 +1420,16 @@ impl App { } } } + let Some(database) = &mut self.database else { + return; + }; + let saved = match database.start_chat_turn(session_id, &prompt, assistant_reasoning) { + Ok(turn) => turn, + Err(error) => { + self.error = Some(format!("Could not save the chat turn: {error}")); + return; + } + }; let cancel = Arc::new(AtomicBool::new(false)); let idle_timeout = Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60); @@ -1378,6 +1447,14 @@ impl App { return; } worker.cancel = Some(cancel); + let user = ChatMessage::from(saved.0); + let mut assistant = ChatMessage::from(saved.1); + assistant.reasoning_open = assistant_reasoning; + self.composer.clear(); + self.conversation.push(user); + self.conversation.push(assistant); + self.generating = true; + self.error = None; } #[cfg(not(target_os = "macos"))] { @@ -1385,19 +1462,6 @@ impl App { self.error = Some("Local Metal generation requires macOS.".into()); return; } - self.composer.clear(); - self.conversation.push(ChatMessage { - user: true, - reasoning: false, - content: prompt, - }); - self.conversation.push(ChatMessage { - user: false, - reasoning: assistant_reasoning, - content: String::new(), - }); - self.generating = true; - self.error = None; } fn poll_generation(&mut self) { @@ -1407,14 +1471,17 @@ impl App { return; }; #[cfg(target_os = "macos")] + let mut transcript_changed = false; + #[cfg(target_os = "macos")] loop { match worker.events.try_recv() { Ok(GenerationEvent::Loading) => {} - Ok(GenerationEvent::Chunk(chunk)) => { + Ok(GenerationEvent::Chunk { reasoning, content }) => { if let Some(message) = self.conversation.last_mut() && !message.user { - message.content.push_str(&chunk); + message.append(reasoning, &content); + transcript_changed = true; } } Ok(GenerationEvent::Finished(result)) => { @@ -1434,6 +1501,26 @@ impl App { } } } + #[cfg(target_os = "macos")] + if transcript_changed + && let Some(message) = self.conversation.last() + && let Some(database) = &mut self.database + && let Err(error) = database.update_message( + message.id, + message.reasoning.as_deref(), + message.reasoning_complete, + &message.content, + ) + { + if let Some(cancel) = self + .generation_worker + .as_ref() + .and_then(|worker| worker.cancel.as_ref()) + { + cancel.store(true, Ordering::Relaxed); + } + self.error = Some(format!("Could not save generated chat text: {error}")); + } } } @@ -1472,9 +1559,15 @@ fn spawn_generation_worker() -> Result { }; } if let Some((_, generator)) = &mut loaded { - let result = generator.generate(&messages, &turn, &cancel, |chunk| { - let _ = event_sender.send(GenerationEvent::Chunk(chunk)); - }); + let result = generator.generate( + &messages, + &turn, + &cancel, + |reasoning, content| { + let _ = event_sender + .send(GenerationEvent::Chunk { reasoning, content }); + }, + ); let _ = event_sender.send(GenerationEvent::Finished(result)); last_used = Instant::now(); } @@ -1588,4 +1681,21 @@ mod tests { assert!(parse_streaming_cache("1.5GB").is_err()); assert_eq!(parse_optional_gib("Memory", "8GB").unwrap(), Some(8)); } + + #[test] + fn assistant_stream_keeps_reasoning_separate_from_the_answer() { + let mut message = ChatMessage { + id: 1, + user: false, + reasoning: Some(String::new()), + reasoning_complete: false, + reasoning_open: true, + content: String::new(), + }; + message.append(true, "working it out"); + message.append(false, "final answer"); + assert_eq!(message.reasoning.as_deref(), Some("working it out")); + assert!(message.reasoning_complete); + assert_eq!(message.content, "final answer"); + } } diff --git a/src/app/view.rs b/src/app/view.rs index f829279..da40eff 100644 --- a/src/app/view.rs +++ b/src/app/view.rs @@ -234,16 +234,50 @@ impl App { .spacing(8), ); } else { - for message in &self.conversation { + for (index, message) in self.conversation.iter().enumerate() { let label = if message.user { "You" } else { "DS4" }; - let content = if !message.user && message.content.is_empty() && self.generating - { - "Loading model…" - } else { - &message.content - }; + let active = self.generating && index + 1 == self.conversation.len(); + let mut body = column![text(label).size(11)].spacing(5); + if let Some(reasoning) = &message.reasoning { + let reasoning_label = + match (message.reasoning_open, message.reasoning_complete, active) { + (true, false, true) => "▾ Thinking", + (false, false, true) => "› Thinking", + (true, false, false) => "▾ Reasoning (stopped)", + (false, false, false) => "› Reasoning (stopped)", + (true, true, _) => "▾ Reasoning", + (false, true, _) => "› Reasoning", + }; + body = body.push( + button(text(reasoning_label).size(12)) + .padding(0) + .style(button::text) + .on_press(Message::ToggleReasoning(index)), + ); + if message.reasoning_open { + body = body.push( + text(if reasoning.is_empty() && active { + "Thinking…" + } else { + reasoning + }) + .size(13) + .color(muted_text()), + ); + } + } + if !message.content.is_empty() { + let content = if message.reasoning.is_some() { + message.content.trim_start() + } else { + &message.content + }; + body = body.push(text(content).size(14)); + } else if active && message.reasoning.is_none() { + body = body.push(text("Loading model…").size(14)); + } messages = messages.push( - container(column![text(label).size(11), text(content).size(14)].spacing(5)) + container(body) .padding(14) .width(Length::Fill) .style(overview_style), diff --git a/src/database.rs b/src/database.rs index 1ed9eb4..f4b8a1e 100644 --- a/src/database.rs +++ b/src/database.rs @@ -4,7 +4,7 @@ use std::collections::HashMap; use std::fs; use std::path::Path; -use crate::schema::{preferences, projects, sessions}; +use crate::schema::{messages, preferences, projects, sessions}; use crate::settings::{ DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, @@ -269,6 +269,28 @@ struct NewSession<'a> { title: &'a str, } +#[derive(Clone, Debug, Identifiable, Queryable, Selectable)] +#[diesel(table_name = messages)] +#[diesel(check_for_backend(diesel::sqlite::Sqlite))] +pub struct StoredMessage { + pub id: i32, + pub session_id: i32, + pub user: bool, + pub reasoning: Option, + pub reasoning_complete: bool, + pub content: String, +} + +#[derive(Insertable)] +#[diesel(table_name = messages)] +struct NewMessage<'a> { + session_id: i32, + user: bool, + reasoning: Option<&'a str>, + reasoning_complete: bool, + content: &'a str, +} + #[derive(Debug)] pub struct ProjectWithSessions { pub project: Project, @@ -408,6 +430,16 @@ impl Database { pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> { self.connection .transaction(|connection| { + diesel::delete( + messages::table.filter( + messages::session_id.eq_any( + sessions::table + .filter(sessions::project_id.eq(project_id)) + .select(sessions::id), + ), + ), + ) + .execute(connection)?; diesel::delete(sessions::table.filter(sessions::project_id.eq(project_id))) .execute(connection)?; diesel::delete(projects::table.find(project_id)).execute(connection)?; @@ -425,7 +457,71 @@ impl Database { } pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> { - diesel::delete(sessions::table.find(session_id)) + self.connection + .transaction(|connection| { + diesel::delete(messages::table.filter(messages::session_id.eq(session_id))) + .execute(connection)?; + diesel::delete(sessions::table.find(session_id)).execute(connection)?; + Ok(()) + }) + .map_err(|error: diesel::result::Error| error.to_string()) + } + + pub fn load_messages(&mut self, session_id: i32) -> Result, String> { + messages::table + .filter(messages::session_id.eq(session_id)) + .order(messages::id.asc()) + .select(StoredMessage::as_select()) + .load(&mut self.connection) + .map_err(|error| error.to_string()) + } + + pub fn start_chat_turn( + &mut self, + session_id: i32, + prompt: &str, + reasoning: bool, + ) -> Result<(StoredMessage, StoredMessage), String> { + self.connection + .transaction(|connection| { + let user = diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: true, + reasoning: None, + reasoning_complete: true, + content: prompt, + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?; + let assistant = diesel::insert_into(messages::table) + .values(NewMessage { + session_id, + user: false, + reasoning: reasoning.then_some(""), + reasoning_complete: !reasoning, + content: "", + }) + .returning(StoredMessage::as_returning()) + .get_result(connection)?; + Ok((user, assistant)) + }) + .map_err(|error: diesel::result::Error| error.to_string()) + } + + pub fn update_message( + &mut self, + id: i32, + reasoning: Option<&str>, + reasoning_complete: bool, + content: &str, + ) -> Result<(), String> { + diesel::update(messages::table.find(id)) + .set(( + messages::reasoning.eq(reasoning), + messages::reasoning_complete.eq(reasoning_complete), + messages::content.eq(content), + )) .execute(&mut self.connection) .map(|_| ()) .map_err(|error| error.to_string()) @@ -535,4 +631,36 @@ mod tests { drop(reopened); fs::remove_file(path).unwrap(); } + + #[test] + fn chat_messages_survive_reopen_and_follow_session_deletion() { + let id = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let path = std::env::temp_dir().join(format!("ds4-chat-{id}.sqlite3")); + let mut database = Database::open(&path).unwrap(); + let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap(); + let session = database.create_session(project.id, "Chat").unwrap(); + let (_, assistant) = database + .start_chat_turn(session.id, "Question", true) + .unwrap(); + database + .update_message(assistant.id, Some("Reasoning"), true, "Answer") + .unwrap(); + drop(database); + + let mut reopened = Database::open(&path).unwrap(); + let messages = reopened.load_messages(session.id).unwrap(); + assert_eq!(messages.len(), 2); + assert!(messages[0].user); + assert_eq!(messages[0].content, "Question"); + assert_eq!(messages[1].reasoning.as_deref(), Some("Reasoning")); + assert!(messages[1].reasoning_complete); + assert_eq!(messages[1].content, "Answer"); + reopened.delete_session(session.id).unwrap(); + assert!(reopened.load_messages(session.id).unwrap().is_empty()); + drop(reopened); + fs::remove_file(path).unwrap(); + } } diff --git a/src/engine.rs b/src/engine.rs index 8215282..1967612 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -265,6 +265,20 @@ impl Model { self.tokenizer.is_stop(token) } + pub(crate) fn is_think_start_token(&self, token: i32) -> bool { + self.tokenizer.is_think_start(token) + } + + pub(crate) fn is_think_end_token(&self, token: i32) -> bool { + self.tokenizer.is_think_end(token) + } + + pub(crate) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool { + self.is_stop_token(token) + || (reasoning == ReasoningMode::Direct + && (self.is_think_start_token(token) || self.is_think_end_token(token))) + } + pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> { self.main.tensor_data(name) } @@ -278,7 +292,8 @@ pub(crate) struct Generator { #[derive(Clone)] pub(crate) struct ChatTurn { pub(crate) user: bool, - pub(crate) reasoning: bool, + pub(crate) reasoning: Option, + pub(crate) reasoning_complete: bool, pub(crate) content: String, } @@ -308,7 +323,7 @@ impl Generator { messages: &[ChatTurn], settings: &TurnSettings, cancelled: &AtomicBool, - mut emit: impl FnMut(String), + mut emit: impl FnMut(bool, String), ) -> Result<(), String> { self.executor.reset()?; let tokens = self.executor.model().render_conversation( @@ -327,6 +342,7 @@ impl Generator { )); } let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645)); + let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct; for token in tokens { if cancelled.load(Ordering::Relaxed) { return Ok(()); @@ -348,11 +364,20 @@ impl Generator { settings.min_p, &mut rng, ); - if self.executor.model().is_stop_token(token) { + if self + .executor + .model() + .is_stop_token_for_reasoning(token, settings.reasoning_mode) + { return Ok(()); } - if let Some(bytes) = self.executor.model().token_bytes(token) { - emit(String::from_utf8_lossy(&bytes).into_owned()); + if self.executor.model().is_think_start_token(token) { + reasoning = true; + } else if self.executor.model().is_think_end_token(token) { + reasoning = false; + emit(false, String::new()); + } else if let Some(bytes) = self.executor.model().token_bytes(token) { + emit(reasoning, String::from_utf8_lossy(&bytes).into_owned()); } self.executor.eval(token)?; } @@ -1385,30 +1410,80 @@ mod tests { model.render_prompt("", "Hello", ReasoningMode::Direct), [0, 128_803, 19_923, 128_804, 128_822] ); + assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct)); + assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High)); + let think_start = *model + .render_prompt("", "Hello", ReasoningMode::High) + .last() + .unwrap(); assert_eq!( model.render_conversation( "", &[ ChatTurn { user: true, - reasoning: false, + reasoning: None, + reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: false, - reasoning: false, + reasoning: None, + reasoning_complete: true, content: "Hello".into(), }, ChatTurn { user: true, - reasoning: false, + reasoning: None, + reasoning_complete: true, content: "Hello".into(), }, ], ReasoningMode::Direct, ), [ - 0, 128_803, 19_923, 128_804, 128_822, 19_923, 128_803, 19_923, 128_804, 128_822, + 0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822, + ] + ); + assert_eq!( + model.render_conversation( + "", + &[ + ChatTurn { + user: true, + reasoning: None, + reasoning_complete: true, + content: "Hello".into(), + }, + ChatTurn { + user: false, + reasoning: Some("Hello".into()), + reasoning_complete: true, + content: "Hello".into(), + }, + ChatTurn { + user: true, + reasoning: None, + reasoning_complete: true, + content: "Hello".into(), + }, + ], + ReasoningMode::Direct, + ), + [ + 0, + 128_803, + 19_923, + 128_804, + think_start, + 19_923, + 128_822, + 19_923, + 1, + 128_803, + 19_923, + 128_804, + 128_822, ] ); } diff --git a/src/engine/tokenizer.rs b/src/engine/tokenizer.rs index b75895c..fbebf54 100644 --- a/src/engine/tokenizer.rs +++ b/src/engine/tokenizer.rs @@ -136,7 +136,8 @@ impl Tokenizer { system_prompt, &[ChatTurn { user: true, - reasoning: false, + reasoning: None, + reasoning_complete: true, content: prompt.to_owned(), }], reasoning, @@ -180,8 +181,12 @@ impl Tokenizer { self.assistant }); if !message.user { - if message.reasoning { + if let Some(reasoning) = &message.reasoning { output.push(self.think_start); + output.extend(self.tokenize(reasoning)); + if message.reasoning_complete { + output.push(self.think_end); + } } else if self.family == ModelFamily::Glm { output.extend([self.think_start, self.think_end]); } else { @@ -189,6 +194,9 @@ impl Tokenizer { } } output.extend(self.tokenize(&message.content)); + if !message.user && self.family == ModelFamily::DeepSeek { + output.push(self.eos); + } } output.push(self.assistant); if reasoning != ReasoningMode::Direct { @@ -211,6 +219,14 @@ impl Tokenizer { && [self.system, self.user, self.assistant, self.observation].contains(&token)) } + pub(super) fn is_think_start(&self, token: i32) -> bool { + token == self.think_start + } + + pub(super) fn is_think_end(&self, token: i32) -> bool { + token == self.think_end + } + fn emit_piece(&self, raw: &[u8], output: &mut Vec) { if raw.is_empty() { return; diff --git a/src/schema.rs b/src/schema.rs index 4e58900..431789c 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -1,3 +1,14 @@ +diesel::table! { + messages (id) { + id -> Integer, + session_id -> Integer, + user -> Bool, + reasoning -> Nullable, + reasoning_complete -> Bool, + content -> Text, + } +} + diesel::table! { projects (id) { id -> Integer, @@ -54,4 +65,5 @@ diesel::table! { } diesel::joinable!(sessions -> projects (project_id)); -diesel::allow_tables_to_appear_in_same_query!(preferences, projects, sessions); +diesel::joinable!(messages -> sessions (session_id)); +diesel::allow_tables_to_appear_in_same_query!(messages, preferences, projects, sessions);