hugo/DS4Server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: better chat support for thinking

Browse Source
This commit is contained in:
Georg Bauer
2026-07-24 18:01:53 +02:00
parent ff984bb96a
commit 5ab110ed0b
9 changed files with 463 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
PLAN.md
View File

@@ -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/
<ds4-cache-key>.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:

1
migrations/20260724230000_create_messages/down.sql Normal file
View File

@@ -0,0 +1 @@
DROP TABLE messages;

10
migrations/20260724230000_create_messages/up.sql Normal file
View File

@@ -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);

168
src/app.rs
View File

@@ -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<String>,
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<StoredMessage> 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<GenerationCommand>,
@@ -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::<Vec<_>>();
#[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<GenerationWorker, String> {
};
}
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");
}
}

50
src/app/view.rs
View File

@@ -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),

132
src/database.rs
View File

@@ -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<String>,
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<Vec<StoredMessage>, 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();
}
}

93
src/engine.rs
View File

@@ -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<String>,
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,
]
);
}

20
src/engine/tokenizer.rs
View File

@@ -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<i32>) {
if raw.is_empty() {
return;

14
src/schema.rs
View File

@@ -1,3 +1,14 @@
diesel::table! {
messages (id) {
id -> Integer,
session_id -> Integer,
user -> Bool,
reasoning -> Nullable<Text>,
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);