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.
This commit is contained in:
Piotr Osiewicz 2024-08-26 18:16:23 +02:00 committed by GitHub
parent 2c541aee24
commit aaddb73b28
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 211 additions and 164 deletions

View File

@ -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<InlineAssistId>,
}
type MessageHeader = MessageMetadata;
pub struct ContextEditor {
context: Model<Context>,
fs: Arc<dyn Fs>,
@ -1709,7 +1712,7 @@ pub struct ContextEditor {
project: Model<Project>,
lsp_adapter_delegate: Option<Arc<dyn LspAdapterDelegate>>,
editor: View<Editor>,
blocks: HashSet<CustomBlockId>,
blocks: HashMap<MessageId, (MessageHeader, CustomBlockId)>,
image_blocks: HashSet<CustomBlockId>,
scroll_position: Option<ScrollPosition>,
remote_id: Option<workspace::ViewId>,
@ -3036,176 +3039,209 @@ impl ContextEditor {
fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
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 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 label = match message.role {
Role::User => {
Label::new("You").color(Color::Default).into_any_element()
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::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),
}
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,
)
})
}
});
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),
)
.into_any_element()
} else {
label.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)),
.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| {
Tooltip::with_meta(
"Error interacting with language model",
None,
"Click for more details",
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| {
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),
.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,
)
.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()
}
}),
disposition: BlockDisposition::Above,
priority: usize::MAX,
})
.into_any_element(),
),
_ => None,
})
.into_any_element()
}
})
.collect::<Vec<_>>();
};
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;
});
}

View File

@ -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<MessageCacheMetadata>,
}
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<usize>) -> bool {
let result = match &self.cache {