From b3ffbea37699b44faa0191aace1cac56f0ca58e4 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Wed, 11 Dec 2024 12:51:05 -0500 Subject: [PATCH] assistant2: Allow removing individual context (#21868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR adds the ability to remove individual pieces of context from the message editor: Screenshot 2024-12-11 at 12 38 45 PM Release Notes: - N/A --- crates/assistant2/src/context.rs | 12 +++++++ crates/assistant2/src/message_editor.rs | 41 +++++++++++++++--------- crates/assistant2/src/ui/context_pill.rs | 30 +++++++++++++++-- 3 files changed, 65 insertions(+), 18 deletions(-) diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs index c865139ce3731..9d095a10d87e1 100644 --- a/crates/assistant2/src/context.rs +++ b/crates/assistant2/src/context.rs @@ -1,8 +1,20 @@ use gpui::SharedString; +use serde::{Deserialize, Serialize}; +use util::post_inc; + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)] +pub struct ContextId(pub(crate) usize); + +impl ContextId { + pub fn post_inc(&mut self) -> Self { + Self(post_inc(&mut self.0)) + } +} /// Some context attached to a message in a thread. #[derive(Debug, Clone)] pub struct Context { + pub id: ContextId, pub name: SharedString, pub kind: ContextKind, pub text: SharedString, diff --git a/crates/assistant2/src/message_editor.rs b/crates/assistant2/src/message_editor.rs index d5e778f5103d3..c237955e2df99 100644 --- a/crates/assistant2/src/message_editor.rs +++ b/crates/assistant2/src/message_editor.rs @@ -1,3 +1,5 @@ +use std::rc::Rc; + use editor::{Editor, EditorElement, EditorStyle}; use gpui::{AppContext, FocusableView, Model, TextStyle, View}; use language_model::{LanguageModelRegistry, LanguageModelRequestTool}; @@ -10,7 +12,7 @@ use ui::{ PopoverMenuHandle, Tooltip, }; -use crate::context::{Context, ContextKind}; +use crate::context::{Context, ContextId, ContextKind}; use crate::context_picker::{ContextPicker, ContextPickerDelegate}; use crate::thread::{RequestKind, Thread}; use crate::ui::ContextPill; @@ -20,19 +22,14 @@ pub struct MessageEditor { thread: Model, editor: View, context: Vec, + next_context_id: ContextId, pub(crate) context_picker_handle: PopoverMenuHandle>, use_tools: bool, } impl MessageEditor { pub fn new(thread: Model, cx: &mut ViewContext) -> Self { - let mocked_context = vec![Context { - name: "shape.rs".into(), - kind: ContextKind::File, - text: "```rs\npub enum Shape {\n Circle,\n Square,\n Triangle,\n}".into(), - }]; - - Self { + let mut this = Self { thread, editor: cx.new_view(|cx| { let mut editor = Editor::auto_height(80, cx); @@ -40,10 +37,20 @@ impl MessageEditor { editor }), - context: mocked_context, + context: Vec::new(), + next_context_id: ContextId(0), context_picker_handle: PopoverMenuHandle::default(), use_tools: false, - } + }; + + this.context.push(Context { + id: this.next_context_id.post_inc(), + name: "shape.rs".into(), + kind: ContextKind::File, + text: "```rs\npub enum Shape {\n Circle,\n Square,\n Triangle,\n}".into(), + }); + + this } fn chat(&mut self, _: &Chat, cx: &mut ViewContext) { @@ -178,11 +185,15 @@ impl Render for MessageEditor { .shape(IconButtonShape::Square) .icon_size(IconSize::Small), )) - .children( - self.context - .iter() - .map(|context| ContextPill::new(context.clone())), - ) + .children(self.context.iter().map(|context| { + ContextPill::new(context.clone()).on_remove({ + let context = context.clone(); + Rc::new(cx.listener(move |this, _event, cx| { + this.context.retain(|other| other.id != context.id); + cx.notify(); + })) + }) + })) .when(!self.context.is_empty(), |parent| { parent.child( IconButton::new("remove-all-context", IconName::Eraser) diff --git a/crates/assistant2/src/ui/context_pill.rs b/crates/assistant2/src/ui/context_pill.rs index 35fc6ab0e22ee..dd74465ad0adb 100644 --- a/crates/assistant2/src/ui/context_pill.rs +++ b/crates/assistant2/src/ui/context_pill.rs @@ -1,25 +1,49 @@ -use ui::prelude::*; +use std::rc::Rc; + +use gpui::ClickEvent; +use ui::{prelude::*, IconButtonShape}; use crate::context::Context; #[derive(IntoElement)] pub struct ContextPill { context: Context, + on_remove: Option>, } impl ContextPill { pub fn new(context: Context) -> Self { - Self { context } + Self { + context, + on_remove: None, + } + } + + pub fn on_remove(mut self, on_remove: Rc) -> Self { + self.on_remove = Some(on_remove); + self } } impl RenderOnce for ContextPill { fn render(self, cx: &mut WindowContext) -> impl IntoElement { - div() + h_flex() + .gap_1() .px_1() .border_1() .border_color(cx.theme().colors().border) .rounded_md() .child(Label::new(self.context.name.clone()).size(LabelSize::Small)) + .when_some(self.on_remove, |parent, on_remove| { + parent.child( + IconButton::new("remove", IconName::Close) + .shape(IconButtonShape::Square) + .icon_size(IconSize::XSmall) + .on_click({ + let on_remove = on_remove.clone(); + move |event, cx| on_remove(event, cx) + }), + ) + }) } }