Files
DS4Server/src/app/generation.rs
2026-07-25 14:57:17 +02:00

604 lines
23 KiB
Rust

use super::*;
/// Asks for a title short enough to fit a sidebar row.
const TITLE_INSTRUCTION: &str = "Give this conversation a short title of at most six words. \
Reply with the title alone: no quotes, no trailing period, no explanation.";
const TITLE_MAX_TOKENS: i32 = 48;
const TITLE_MAX_CHARS: usize = 60;
/// A one-shot generation that produces a session title. It runs against the
/// transient KV cache, so it never creates or touches a stored session.
#[cfg(target_os = "macos")]
pub(super) struct TitleRequest {
session_id: i32,
/// Title the session had when the request was queued. If it changed since,
/// somebody renamed it by hand and that newer intent wins.
expected: String,
active: ActiveGeneration,
content: String,
}
#[derive(Clone, Debug)]
pub(crate) struct ChatMessage {
pub(super) id: i32,
pub(super) user: bool,
pub(super) tool: bool,
pub(super) reasoning: Option<String>,
pub(super) reasoning_complete: bool,
pub(super) reasoning_open: bool,
pub(super) content: String,
pub(super) markdown: Vec<markdown::Item>,
}
impl ChatMessage {
pub(super) 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);
}
}
pub(super) fn refresh_markdown(&mut self) {
if !self.user && !self.tool {
let visible = crate::agent::visible_content(&self.content);
let content = if self.reasoning.is_some() {
visible.trim_start()
} else {
visible
};
self.markdown = markdown::parse(content).collect();
}
}
}
impl From<StoredMessage> for ChatMessage {
fn from(message: StoredMessage) -> Self {
let mut message = Self {
id: message.id,
user: message.user,
tool: message.tool,
reasoning: message.reasoning,
reasoning_complete: message.reasoning_complete,
reasoning_open: false,
content: message.content,
markdown: Vec::new(),
};
message.refresh_markdown();
message
}
}
impl App {
pub(super) fn start_generation(&mut self) {
if self.generating || self.selected_project.is_none() {
return;
}
let prompt = self.composer.trim().to_owned();
if prompt.is_empty() {
return;
}
let model = match ModelChoice::from_id(&self.preferences.selected_model) {
Some(model) => model,
None => {
self.error = Some("The selected model is not supported.".into());
return;
}
};
let effective = self.preferences.generation().and_then(|generation| {
self.preferences.runtime().and_then(|runtime| {
crate::settings::effective_settings(model, &generation, &runtime, &models_path())
})
});
let mut effective = match effective {
Ok(settings) => settings,
Err(error) => {
self.error = Some(error);
return;
}
};
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
#[cfg(target_os = "macos")]
let mut messages = self
.conversation
.iter()
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
skip_previous_eos: false,
reasoning: message.reasoning.clone(),
reasoning_complete: message.reasoning_complete,
content: message.content.clone(),
})
.collect::<Vec<_>>();
#[cfg(target_os = "macos")]
messages.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: prompt.clone(),
});
#[cfg(target_os = "macos")]
{
// A draft session only reaches the database once there is a turn to store.
let opening_turn = self.selected_session.is_none();
let session_id = match self.selected_session {
Some(session_id) => session_id,
None => {
let Some(project_id) = self.selected_project else {
return;
};
match self.persist_session(project_id) {
Ok(session_id) => session_id,
Err(error) => {
self.error = Some(format!("Could not create the session: {error}"));
return;
}
}
}
};
let Some(service) = &self.generation_service else {
self.error = Some("The model runtime is unavailable.".into());
return;
};
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 idle_timeout =
Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60);
self.active_generation = match service.generate(
effective.engine,
effective.turn,
messages,
CheckpointTarget::Local(session_checkpoint_path(session_id)),
idle_timeout,
) {
Ok(active) => Some(active),
Err(error) => {
self.generation_service = None;
self.error = Some(error);
return;
}
};
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.tokens_per_second = None;
self.error = None;
if opening_turn {
// Queued behind the reply, so the answer is never delayed by it.
self.request_first_title(session_id);
}
}
#[cfg(not(target_os = "macos"))]
{
let _ = effective;
self.error = Some("Local Metal generation requires macOS.".into());
return;
}
}
pub(super) fn poll_generation(&mut self) -> bool {
#[cfg(target_os = "macos")]
if let Some(active) = &self.active_tools {
match crate::agent::try_tool_result(active) {
Ok(Some(result)) => {
let cancelled = active.cancel.load(Ordering::Relaxed);
self.active_tools = None;
if cancelled {
self.generating = false;
return false;
}
if let Err(error) = self.continue_after_tool_result(&result) {
self.generating = false;
self.error = Some(error);
}
return true;
}
Ok(None) => return false,
Err(error) => {
self.active_tools = None;
self.generating = false;
self.error = Some(error);
return false;
}
}
}
#[cfg(target_os = "macos")]
let Some(active) = &mut self.active_generation else {
self.generating = false;
return false;
};
#[cfg(target_os = "macos")]
let mut transcript_changed = false;
#[cfg(target_os = "macos")]
let mut context_changed = false;
#[cfg(target_os = "macos")]
loop {
match active.events.try_recv() {
Ok(GenerationEvent::Loading) => {}
Ok(GenerationEvent::Chunk { reasoning, content }) => {
if let Some(message) = self.conversation.last_mut()
&& !message.user
{
message.append(reasoning, &content);
transcript_changed = true;
}
}
Ok(GenerationEvent::Context {
used,
limit,
tokens_per_second,
}) => {
self.context_used = used;
self.context_limit = limit;
self.tokens_per_second = tokens_per_second;
context_changed = true;
}
Ok(GenerationEvent::Finished(result)) => {
match result {
Ok(_) => {
let model = ModelChoice::from_id(&self.preferences.selected_model)
.unwrap_or_default();
let content = self
.conversation
.last()
.map(|message| message.content.clone())
.unwrap_or_default();
match crate::agent::parse_tool_calls(model, &content) {
Ok((_, calls)) if !calls.is_empty() => {
if let Err(error) = self.start_agent_tools(calls) {
self.generating = false;
self.error = Some(error);
}
}
Ok(_) => self.generating = false,
Err(error) => {
self.active_tools = Some(crate::agent::error_async(error));
}
}
}
Err(error) => {
self.generating = false;
self.error = Some(error);
}
}
self.active_generation = None;
break;
}
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
self.generating = false;
self.active_generation = None;
self.error = Some("The model runtime stopped unexpectedly.".into());
break;
}
}
}
#[cfg(target_os = "macos")]
if transcript_changed && let Some(message) = self.conversation.last_mut() {
message.refresh_markdown();
}
#[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(active) = &self.active_generation {
active.cancel.store(true, Ordering::Relaxed);
}
self.error = Some(format!("Could not save generated chat text: {error}"));
}
#[cfg(target_os = "macos")]
if context_changed
&& let Some(session_id) = self.selected_session
&& let Some(database) = &mut self.database
{
if let Err(error) = database.update_session_context(
session_id,
self.context_used,
self.context_limit,
self.tokens_per_second,
) {
self.error = Some(format!("Could not save context usage: {error}"));
} else if let Some(session) = self
.projects
.iter_mut()
.flat_map(|project| &mut project.sessions)
.find(|session| session.id == session_id)
{
session.context_used = self.context_used as i32;
session.context_limit = self.context_limit as i32;
session.last_tokens_per_second = self.tokens_per_second;
}
}
#[cfg(target_os = "macos")]
return transcript_changed;
#[cfg(not(target_os = "macos"))]
false
}
#[cfg(target_os = "macos")]
fn start_agent_tools(&mut self, calls: Vec<crate::agent::ToolCall>) -> Result<(), String> {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
if self.agent_tools.as_ref().map(|(id, _)| *id) != Some(session_id) {
let project_id = self
.selected_project
.ok_or_else(|| "The active project is unavailable.".to_owned())?;
let root = self
.projects
.iter()
.find(|project| project.project.id == project_id)
.map(|project| PathBuf::from(&project.project.path))
.ok_or_else(|| "The active project is unavailable.".to_owned())?;
let tools = crate::agent::Tools::new(&root, self.preferences.context_tokens)?;
self.agent_tools = Some((session_id, Arc::new(Mutex::new(tools))));
}
let tools = Arc::clone(&self.agent_tools.as_ref().unwrap().1);
self.active_tools = Some(crate::agent::execute_async(tools, calls));
Ok(())
}
#[cfg(target_os = "macos")]
fn continue_after_tool_result(&mut self, result: &str) -> Result<(), String> {
let session_id = self
.selected_session
.ok_or_else(|| "The active session is unavailable.".to_owned())?;
let model = ModelChoice::from_id(&self.preferences.selected_model)
.ok_or_else(|| "The selected model is not supported.".to_owned())?;
let generation = self.preferences.generation()?;
let runtime = self.preferences.runtime()?;
let mut effective =
crate::settings::effective_settings(model, &generation, &runtime, &models_path())?;
effective.turn.system_prompt =
crate::agent::system_prompt(model, &effective.turn.system_prompt);
let assistant_reasoning = effective.turn.reasoning_mode != ReasoningMode::Direct;
let saved = self
.database
.as_mut()
.ok_or_else(|| "The project database is unavailable.".to_owned())?
.continue_tool_turn(session_id, result, assistant_reasoning)
.map_err(|error| format!("Could not save the tool turn: {error}"))?;
self.conversation.push(ChatMessage::from(saved.0));
let messages = self
.conversation
.iter()
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
skip_previous_eos: false,
reasoning: message.reasoning.clone(),
reasoning_complete: message.reasoning_complete,
content: message.content.clone(),
})
.collect();
let mut assistant = ChatMessage::from(saved.1);
assistant.reasoning_open = assistant_reasoning;
self.conversation.push(assistant);
let idle_timeout =
Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60);
self.active_generation = Some(
self.generation_service
.as_ref()
.ok_or_else(|| "The model runtime is unavailable.".to_owned())?
.generate(
effective.engine,
effective.turn,
messages,
CheckpointTarget::Local(session_checkpoint_path(session_id)),
idle_timeout,
)?,
);
self.tokens_per_second = None;
Ok(())
}
/// Asks the model for a session title without opening a session for it: the
/// turn runs against the shared transient KV cache and only its text is kept.
///
/// Reports failures to the caller instead of the error banner: the automatic
/// first-chat titling must stay silent, while the menu action shows them.
#[cfg(target_os = "macos")]
fn request_title(&mut self, session_id: i32) -> Result<(), String> {
if self.active_titling.is_some() {
return Err("A title is already being generated.".into());
}
let model = ModelChoice::from_id(&self.preferences.selected_model)
.ok_or_else(|| "The selected model is not supported.".to_owned())?;
let mut effective = self.preferences.generation().and_then(|generation| {
self.preferences.runtime().and_then(|runtime| {
crate::settings::effective_settings(model, &generation, &runtime, &models_path())
})
})?;
let database = self
.database
.as_mut()
.ok_or_else(|| "The project database is unavailable.".to_owned())?;
let stored = database
.load_messages(session_id)
.map_err(|error| format!("Could not read the session: {error}"))?;
// The assistant row is inserted empty when a turn starts; an empty turn
// is never useful context, and skipping it is what lets the automatic
// pass summarize the user's opening message on its own.
let mut messages = stored
.into_iter()
.filter(|message| !message.content.trim().is_empty())
.map(|message| ChatTurn {
user: message.user,
tool: message.tool,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: message.content,
})
.collect::<Vec<_>>();
if messages.is_empty() {
return Err("This session has no messages to summarize yet.".into());
}
messages.push(ChatTurn {
user: true,
tool: false,
skip_previous_eos: false,
reasoning: None,
reasoning_complete: true,
content: TITLE_INSTRUCTION.to_owned(),
});
effective.turn.max_generated_tokens = TITLE_MAX_TOKENS;
effective.turn.reasoning_mode = ReasoningMode::Direct;
let service = self
.generation_service
.as_ref()
.ok_or_else(|| "The model runtime is unavailable.".to_owned())?;
let idle_timeout =
Duration::from_secs(self.preferences.idle_timeout_minutes.max(1) as u64 * 60);
match service.generate(
effective.engine,
effective.turn,
messages,
CheckpointTarget::OneShot(transient_cache_path()),
idle_timeout,
) {
Ok(active) => {
self.active_titling = Some(TitleRequest {
session_id,
expected: self.session_title(session_id).unwrap_or_default(),
active,
content: String::new(),
});
Ok(())
}
Err(error) => {
self.generation_service = None;
Err(error)
}
}
}
/// The "Retitle with AI" menu action, which reports why it could not run.
#[cfg(target_os = "macos")]
pub(super) fn retitle_session(&mut self, session_id: i32) {
match self.request_title(session_id) {
Ok(()) => self.error = None,
Err(error) => self.error = Some(error),
}
}
/// Titles a session from its opening message. Queued behind the reply that
/// is already running, and skipped without complaint if it cannot start —
/// the user did not ask for this, so it must never interrupt them.
#[cfg(target_os = "macos")]
pub(super) fn request_first_title(&mut self, session_id: i32) {
let _ = self.request_title(session_id);
}
/// Drains the one-shot title generation. Returns true while it is running so
/// the caller keeps ticking.
#[cfg(target_os = "macos")]
pub(super) fn poll_titling(&mut self) -> bool {
let Some(mut request) = self.active_titling.take() else {
return false;
};
let failure = loop {
match request.active.events.try_recv() {
Ok(GenerationEvent::Chunk {
reasoning: false,
content,
}) => request.content.push_str(&content),
Ok(GenerationEvent::Finished(Err(error))) => break Some(error),
Ok(GenerationEvent::Finished(Ok(_))) => break None,
Ok(_) => {}
Err(TryRecvError::Empty) => {
self.active_titling = Some(request);
return true;
}
Err(TryRecvError::Disconnected) => {
break Some("The model runtime stopped unexpectedly.".to_owned());
}
}
};
if let Some(error) = failure {
self.error = Some(format!("Could not generate a title: {error}"));
} else if let Some(title) = session_title(&request.content) {
self.apply_title(&request, &title);
} else {
self.error = Some("The model did not return a usable title.".into());
}
false
}
#[cfg(target_os = "macos")]
fn session_title(&self, session_id: i32) -> Option<String> {
self.projects
.iter()
.flat_map(|project| &project.sessions)
.find(|session| session.id == session_id)
.map(|session| session.title.clone())
}
#[cfg(target_os = "macos")]
fn apply_title(&mut self, request: &TitleRequest, title: &str) {
let Some(current) = self.session_title(request.session_id) else {
return;
};
if current != request.expected {
return;
}
let Some(database) = &mut self.database else {
return;
};
match database.rename_session(request.session_id, title) {
Ok(()) => {
self.error = None;
self.reload_projects();
}
Err(error) => self.error = Some(format!("Could not save the session title: {error}")),
}
}
}
/// Reduces a model reply to a single sidebar-sized line, or `None` if nothing
/// usable came back.
pub(super) fn session_title(reply: &str) -> Option<String> {
let title = reply
.lines()
.map(str::trim)
.find(|line| !line.is_empty())?
.trim_matches(|character: char| {
character.is_whitespace() || matches!(character, '"' | '\'' | '`' | '*' | '#' | '.')
})
.trim();
if title.is_empty() {
return None;
}
Some(match title.char_indices().nth(TITLE_MAX_CHARS) {
Some((index, _)) => format!("{}", title[..index].trim_end()),
None => title.to_owned(),
})
}