From aaddb73b2899896e2a9f39874f20ae2784daf93b Mon Sep 17 00:00:00 2001 From: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com> Date: Mon, 26 Aug 2024 18:16:23 +0200 Subject: [PATCH] assistant: Refesh message headers only for dirty messages (#16881) We've noticed performance issues in long conversations with assistants; the profiles pointed to slowiness in WrapMap (and indeed there were some low hanging fruits that we picked up in https://github.com/zed-industries/zed/pull/16761). That however did not fully resolve the issue, as WrapMap still cracked through in profiles; basically, the speedup I've landed has just moved the post elsewhere. The higher level issue is that we were trying to refresh message headers for all messages, irrespective of whether they've actually needed to be updated. This PR fixes that by using `replace_blocks` API where possible. Release Notes: - Improved performance of Assistant Panel with long conversations. --- crates/assistant/src/assistant_panel.rs | 346 +++++++++++++----------- crates/assistant/src/context.rs | 13 +- 2 files changed, 203 insertions(+), 156 deletions(-) diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index faec9f8ef7b939cebdf9199ae9fbc15862b4ebaa..55916f33453fb2ab2070cd9fd8842b874377c391 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -13,9 +13,10 @@ use crate::{ terminal_inline_assistant::TerminalInlineAssistant, Assist, CacheStatus, ConfirmCommand, Context, ContextEvent, ContextId, ContextStore, CycleMessageRole, DeployHistory, DeployPromptLibrary, InlineAssistId, InlineAssistant, - InsertIntoEditor, MessageStatus, ModelSelector, PendingSlashCommand, PendingSlashCommandStatus, - QuoteSelection, RemoteContextMetadata, SavedContextMetadata, Split, ToggleFocus, - ToggleModelSelector, WorkflowStepResolution, WorkflowStepView, + InsertIntoEditor, Message, MessageId, MessageMetadata, MessageStatus, ModelSelector, + PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection, RemoteContextMetadata, + SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector, WorkflowStepResolution, + WorkflowStepView, }; use crate::{ContextStoreEvent, ModelPickerDelegate}; use anyhow::{anyhow, Result}; @@ -1702,6 +1703,8 @@ struct WorkflowAssist { assist_ids: Vec, } +type MessageHeader = MessageMetadata; + pub struct ContextEditor { context: Model, fs: Arc, @@ -1709,7 +1712,7 @@ pub struct ContextEditor { project: Model, lsp_adapter_delegate: Option>, editor: View, - blocks: HashSet, + blocks: HashMap, image_blocks: HashSet, scroll_position: Option, remote_id: Option, @@ -3036,176 +3039,209 @@ impl ContextEditor { fn update_message_headers(&mut self, cx: &mut ViewContext) { self.editor.update(cx, |editor, cx| { let buffer = editor.buffer().read(cx).snapshot(cx); - let excerpt_id = *buffer.as_singleton().unwrap().0; - let old_blocks = std::mem::take(&mut self.blocks); - let new_blocks = self - .context - .read(cx) - .messages(cx) - .map(|message| BlockProperties { - position: buffer - .anchor_in_excerpt(excerpt_id, message.anchor) - .unwrap(), - height: 2, - style: BlockStyle::Sticky, - render: Box::new({ - let context = self.context.clone(); - move |cx| { - let message_id = message.id; - let show_spinner = message.role == Role::Assistant - && message.status == MessageStatus::Pending; - let label = match message.role { - Role::User => { - Label::new("You").color(Color::Default).into_any_element() - } - Role::Assistant => { - let label = Label::new("Assistant").color(Color::Info); - if show_spinner { - label - .with_animation( - "pulsating-label", - Animation::new(Duration::from_secs(2)) - .repeat() - .with_easing(pulsating_between(0.4, 0.8)), - |label, delta| label.alpha(delta), - ) - .into_any_element() - } else { - label.into_any_element() - } + let excerpt_id = *buffer.as_singleton().unwrap().0; + let mut old_blocks = std::mem::take(&mut self.blocks); + let mut blocks_to_remove: HashMap<_, _> = old_blocks + .iter() + .map(|(message_id, (_, block_id))| (*message_id, *block_id)) + .collect(); + let mut blocks_to_replace: HashMap<_, RenderBlock> = Default::default(); + + let render_block = |message: MessageMetadata| -> RenderBlock { + Box::new({ + let context = self.context.clone(); + move |cx| { + let message_id = MessageId(message.timestamp); + let show_spinner = message.role == Role::Assistant + && message.status == MessageStatus::Pending; + + let label = match message.role { + Role::User => { + Label::new("You").color(Color::Default).into_any_element() + } + Role::Assistant => { + let label = Label::new("Assistant").color(Color::Info); + if show_spinner { + label + .with_animation( + "pulsating-label", + Animation::new(Duration::from_secs(2)) + .repeat() + .with_easing(pulsating_between(0.4, 0.8)), + |label, delta| label.alpha(delta), + ) + .into_any_element() + } else { + label.into_any_element() } + } - Role::System => Label::new("System") - .color(Color::Warning) - .into_any_element(), - }; + Role::System => Label::new("System") + .color(Color::Warning) + .into_any_element(), + }; - let sender = ButtonLike::new("role") - .style(ButtonStyle::Filled) - .child(label) - .tooltip(|cx| { - Tooltip::with_meta( - "Toggle message role", - None, - "Available roles: You (User), Assistant, System", - cx, - ) - }) - .on_click({ - let context = context.clone(); - move |_, cx| { - context.update(cx, |context, cx| { - context.cycle_message_roles( - HashSet::from_iter(Some(message_id)), - cx, - ) - }) - } - }); + let sender = ButtonLike::new("role") + .style(ButtonStyle::Filled) + .child(label) + .tooltip(|cx| { + Tooltip::with_meta( + "Toggle message role", + None, + "Available roles: You (User), Assistant, System", + cx, + ) + }) + .on_click({ + let context = context.clone(); + move |_, cx| { + context.update(cx, |context, cx| { + context.cycle_message_roles( + HashSet::from_iter(Some(message_id)), + cx, + ) + }) + } + }); - h_flex() - .id(("message_header", message_id.as_u64())) - .pl(cx.gutter_dimensions.full_width()) - .h_11() - .w_full() - .relative() - .gap_1() - .child(sender) - .children(match &message.cache { - Some(cache) if cache.is_final_anchor => match cache.status { - CacheStatus::Cached => Some( - div() - .id("cached") - .child( - Icon::new(IconName::DatabaseZap) - .size(IconSize::XSmall) - .color(Color::Hint), - ) - .tooltip(|cx| { - Tooltip::with_meta( - "Context cached", - None, - "Large messages cached to optimize performance", - cx, - ) - }).into_any_element() - ), - CacheStatus::Pending => Some( - div() - .child( - Icon::new(IconName::Ellipsis) - .size(IconSize::XSmall) - .color(Color::Hint), - ).into_any_element() - ), - }, - _ => None, - }) - .children(match &message.status { - MessageStatus::Error(error) => Some( - Button::new("show-error", "Error") - .color(Color::Error) - .selected_label_color(Color::Error) - .selected_icon_color(Color::Error) - .icon(IconName::XCircle) - .icon_color(Color::Error) - .icon_size(IconSize::Small) - .icon_position(IconPosition::Start) - .tooltip(move |cx| { + h_flex() + .id(("message_header", message_id.as_u64())) + .pl(cx.gutter_dimensions.full_width()) + .h_11() + .w_full() + .relative() + .gap_1() + .child(sender) + .children(match &message.cache { + Some(cache) if cache.is_final_anchor => match cache.status { + CacheStatus::Cached => Some( + div() + .id("cached") + .child( + Icon::new(IconName::DatabaseZap) + .size(IconSize::XSmall) + .color(Color::Hint), + ) + .tooltip(|cx| { Tooltip::with_meta( - "Error interacting with language model", + "Context cached", None, - "Click for more details", + "Large messages cached to optimize performance", cx, ) }) - .on_click({ - let context = context.clone(); - let error = error.clone(); - move |_, cx| { - context.update(cx, |_, cx| { - cx.emit(ContextEvent::ShowAssistError( - error.clone(), - )); - }); - } - }) .into_any_element(), ), - MessageStatus::Canceled => Some( - ButtonLike::new("canceled") - .child( - Icon::new(IconName::XCircle).color(Color::Disabled), - ) + CacheStatus::Pending => Some( + div() .child( - Label::new("Canceled") - .size(LabelSize::Small) - .color(Color::Disabled), + Icon::new(IconName::Ellipsis) + .size(IconSize::XSmall) + .color(Color::Hint), ) - .tooltip(move |cx| { - Tooltip::with_meta( - "Canceled", - None, - "Interaction with the assistant was canceled", - cx, - ) - }) .into_any_element(), ), - _ => None, - }) - .into_any_element() - } - }), - disposition: BlockDisposition::Above, - priority: usize::MAX, + }, + _ => None, + }) + .children(match &message.status { + MessageStatus::Error(error) => Some( + Button::new("show-error", "Error") + .color(Color::Error) + .selected_label_color(Color::Error) + .selected_icon_color(Color::Error) + .icon(IconName::XCircle) + .icon_color(Color::Error) + .icon_size(IconSize::Small) + .icon_position(IconPosition::Start) + .tooltip(move |cx| { + Tooltip::with_meta( + "Error interacting with language model", + None, + "Click for more details", + cx, + ) + }) + .on_click({ + let context = context.clone(); + let error = error.clone(); + move |_, cx| { + context.update(cx, |_, cx| { + cx.emit(ContextEvent::ShowAssistError( + error.clone(), + )); + }); + } + }) + .into_any_element(), + ), + MessageStatus::Canceled => Some( + ButtonLike::new("canceled") + .child(Icon::new(IconName::XCircle).color(Color::Disabled)) + .child( + Label::new("Canceled") + .size(LabelSize::Small) + .color(Color::Disabled), + ) + .tooltip(move |cx| { + Tooltip::with_meta( + "Canceled", + None, + "Interaction with the assistant was canceled", + cx, + ) + }) + .into_any_element(), + ), + _ => None, + }) + .into_any_element() + } }) - .collect::>(); + }; + let create_block_properties = |message: &Message| BlockProperties { + position: buffer + .anchor_in_excerpt(excerpt_id, message.anchor) + .unwrap(), + height: 2, + style: BlockStyle::Sticky, + disposition: BlockDisposition::Above, + priority: usize::MAX, + render: render_block(MessageMetadata::from(message)), + }; + let mut new_blocks = vec![]; + let mut block_index_to_message = vec![]; + for message in self.context.read(cx).messages(cx) { + if let Some(_) = blocks_to_remove.remove(&message.id) { + // This is an old message that we might modify. + let Some((meta, block_id)) = old_blocks.get_mut(&message.id) else { + debug_assert!( + false, + "old_blocks should contain a message_id we've just removed." + ); + continue; + }; + // Should we modify it? + let message_meta = MessageMetadata::from(&message); + if meta != &message_meta { + blocks_to_replace.insert(*block_id, render_block(message_meta.clone())); + *meta = message_meta; + } + } else { + // This is a new message. + new_blocks.push(create_block_properties(&message)); + block_index_to_message.push((message.id, MessageMetadata::from(&message))); + } + } + editor.replace_blocks(blocks_to_replace, None, cx); + editor.remove_blocks(blocks_to_remove.into_values().collect(), None, cx); - editor.remove_blocks(old_blocks, None, cx); let ids = editor.insert_blocks(new_blocks, None, cx); - self.blocks = HashSet::from_iter(ids); + old_blocks.extend(ids.into_iter().zip(block_index_to_message).map( + |(block_id, (message_id, message_meta))| (message_id, (message_meta, block_id)), + )); + self.blocks = old_blocks; }); } diff --git a/crates/assistant/src/context.rs b/crates/assistant/src/context.rs index 41f894aeb206e6c6e2d0188bb32c030f4f8cce4e..6a810dd98801f76828313436e2ac6f8249c2cb3d 100644 --- a/crates/assistant/src/context.rs +++ b/crates/assistant/src/context.rs @@ -330,11 +330,22 @@ pub struct MessageCacheMetadata { pub struct MessageMetadata { pub role: Role, pub status: MessageStatus, - timestamp: clock::Lamport, + pub(crate) timestamp: clock::Lamport, #[serde(skip)] pub cache: Option, } +impl From<&Message> for MessageMetadata { + fn from(message: &Message) -> Self { + Self { + role: message.role, + status: message.status.clone(), + timestamp: message.id.0, + cache: message.cache.clone(), + } + } +} + impl MessageMetadata { pub fn is_cache_valid(&self, buffer: &BufferSnapshot, range: &Range) -> bool { let result = match &self.cache {