Improve chat rendering and generation summaries
This commit is contained in:
@@ -55,6 +55,7 @@ pub(crate) struct ChatMessage {
|
||||
pub(super) system: bool,
|
||||
pub(super) compaction: bool,
|
||||
pub(super) compaction_tail_start: Option<i32>,
|
||||
pub(super) generation_stats: Option<GenerationStats>,
|
||||
pub(super) reasoning: Option<String>,
|
||||
pub(super) reasoning_complete: bool,
|
||||
pub(super) reasoning_open: bool,
|
||||
@@ -66,6 +67,14 @@ pub(crate) struct ChatMessage {
|
||||
pub(super) a2ui_open_urls: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub(super) struct GenerationStats {
|
||||
pub(super) duration_ms: u64,
|
||||
pub(super) input_tokens: u32,
|
||||
pub(super) cached_tokens: u32,
|
||||
pub(super) output_tokens: u32,
|
||||
}
|
||||
|
||||
const TOOL_PROTOCOL_CORRECTION: &str = "Tool protocol correction (no tool was executed): you emitted a complete tool call inside private reasoning. End reasoning, then emit the tool call again as assistant content using the required syntax; do not merely discuss it.";
|
||||
const EMPTY_RESPONSE_CORRECTION: &str = "Protocol correction: your previous response was empty. Continue the task now with either a valid tool call or a final answer.";
|
||||
|
||||
@@ -122,6 +131,27 @@ fn correction_already_sent(conversation: &[ChatMessage], correction: &str) -> bo
|
||||
|
||||
impl From<StoredMessage> for ChatMessage {
|
||||
fn from(message: StoredMessage) -> Self {
|
||||
let generation_stats = match (
|
||||
message.generation_duration_ms,
|
||||
message.input_tokens,
|
||||
message.cached_tokens,
|
||||
message.output_tokens,
|
||||
) {
|
||||
(Some(duration_ms), Some(input_tokens), Some(cached_tokens), Some(output_tokens))
|
||||
if duration_ms >= 0
|
||||
&& input_tokens >= 0
|
||||
&& cached_tokens >= 0
|
||||
&& output_tokens >= 0 =>
|
||||
{
|
||||
Some(GenerationStats {
|
||||
duration_ms: duration_ms as u64,
|
||||
input_tokens: input_tokens as u32,
|
||||
cached_tokens: cached_tokens as u32,
|
||||
output_tokens: output_tokens as u32,
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let mut message = Self {
|
||||
id: message.id,
|
||||
user: message.user,
|
||||
@@ -129,6 +159,7 @@ impl From<StoredMessage> for ChatMessage {
|
||||
system: message.system,
|
||||
compaction: message.compaction,
|
||||
compaction_tail_start: message.compaction_tail_start,
|
||||
generation_stats,
|
||||
reasoning: message.reasoning,
|
||||
reasoning_complete: message.reasoning_complete,
|
||||
reasoning_open: false,
|
||||
@@ -661,6 +692,30 @@ impl App {
|
||||
context_changed = true;
|
||||
}
|
||||
Ok(GenerationEvent::Finished(result)) => {
|
||||
if let Ok(output) = &result {
|
||||
let stats = GenerationStats {
|
||||
duration_ms: u64::try_from(active.started_at.elapsed().as_millis())
|
||||
.unwrap_or(u64::MAX),
|
||||
input_tokens: output.prompt_tokens,
|
||||
cached_tokens: output.cached_tokens,
|
||||
output_tokens: output.completion_tokens,
|
||||
};
|
||||
if let Some(message) = self.conversation.last_mut() {
|
||||
message.generation_stats = Some(stats);
|
||||
if let Some(database) = &mut self.database
|
||||
&& let Err(error) = database.update_message_generation_stats(
|
||||
message.id,
|
||||
stats.duration_ms,
|
||||
stats.input_tokens,
|
||||
stats.cached_tokens,
|
||||
stats.output_tokens,
|
||||
)
|
||||
{
|
||||
self.error =
|
||||
Some(format!("Could not save generation summary: {error}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
match result {
|
||||
Ok(_) if self.stop_requested => {
|
||||
self.generating = false;
|
||||
@@ -1605,6 +1660,7 @@ mod tests {
|
||||
system: false,
|
||||
compaction: false,
|
||||
compaction_tail_start: None,
|
||||
generation_stats: None,
|
||||
reasoning: reasoning.map(str::to_owned),
|
||||
reasoning_complete: false,
|
||||
reasoning_open: false,
|
||||
@@ -1664,6 +1720,7 @@ mod tests {
|
||||
system: false,
|
||||
compaction: false,
|
||||
compaction_tail_start: None,
|
||||
generation_stats: None,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
reasoning_open: false,
|
||||
@@ -1729,6 +1786,7 @@ mod tests {
|
||||
system: compaction,
|
||||
compaction,
|
||||
compaction_tail_start: tail,
|
||||
generation_stats: None,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
reasoning_open: false,
|
||||
@@ -1765,6 +1823,7 @@ mod tests {
|
||||
system,
|
||||
compaction,
|
||||
compaction_tail_start: None,
|
||||
generation_stats: None,
|
||||
reasoning: None,
|
||||
reasoning_complete: true,
|
||||
reasoning_open: false,
|
||||
|
||||
@@ -947,10 +947,7 @@ fn preference_group_style(_: &Theme) -> container::Style {
|
||||
}
|
||||
}
|
||||
|
||||
fn chat_message_style(theme: &Theme, user: bool) -> container::Style {
|
||||
if !user {
|
||||
return overview_style(theme);
|
||||
}
|
||||
fn chat_message_style(_: &Theme) -> container::Style {
|
||||
container::Style {
|
||||
background: Some(Background::Color(Color::from_rgb8(27, 34, 44))),
|
||||
border: Border {
|
||||
@@ -1052,8 +1049,8 @@ mod tests {
|
||||
assert_eq!(palette.background.weak.color, Color::from_rgb8(38, 38, 40));
|
||||
assert_eq!(palette.secondary.base.color, Color::from_rgb8(47, 47, 50));
|
||||
assert_ne!(
|
||||
chat_message_style(&theme, true).background,
|
||||
chat_message_style(&theme, false).background
|
||||
chat_message_style(&theme).background,
|
||||
overview_style(&theme).background
|
||||
);
|
||||
assert_eq!(
|
||||
preference_group_style(&theme).background,
|
||||
|
||||
@@ -180,13 +180,15 @@ impl App {
|
||||
body = body.push(tool_cards(cards));
|
||||
}
|
||||
}
|
||||
let user = message.user;
|
||||
messages = messages.push(
|
||||
container(body)
|
||||
.padding(14)
|
||||
.width(Length::Fill)
|
||||
.style(move |theme| chat_message_style(theme, user)),
|
||||
);
|
||||
let message_body = container(body).padding(14).width(Length::Fill);
|
||||
messages = messages.push(if message.user {
|
||||
message_body.style(chat_message_style)
|
||||
} else {
|
||||
message_body
|
||||
});
|
||||
if let Some(stats) = message.generation_stats {
|
||||
messages = messages.push(generation_summary(stats));
|
||||
}
|
||||
}
|
||||
if let Some(activity) = &self.activity {
|
||||
messages = messages.push(
|
||||
@@ -335,6 +337,26 @@ impl App {
|
||||
}
|
||||
}
|
||||
|
||||
fn generation_summary(stats: crate::app::generation::GenerationStats) -> Element<'static, Message> {
|
||||
column![
|
||||
row![
|
||||
Space::new().width(Length::Fill),
|
||||
text(format!(
|
||||
"{} · {} input · {} cached · {} output",
|
||||
format_duration(stats.duration_ms as f64 / 1_000.0),
|
||||
stats.input_tokens,
|
||||
stats.cached_tokens,
|
||||
stats.output_tokens,
|
||||
))
|
||||
.size(11)
|
||||
.color(muted_text()),
|
||||
],
|
||||
rule::horizontal(1),
|
||||
]
|
||||
.spacing(5)
|
||||
.into()
|
||||
}
|
||||
|
||||
fn tool_cards(cards: Vec<crate::agent::ToolCard>) -> Element<'static, Message> {
|
||||
let mut rows = column![].spacing(0);
|
||||
for (index, card) in cards.into_iter().enumerate() {
|
||||
|
||||
Reference in New Issue
Block a user