From 374c298bd528f13bd729727368d39c45aec98eb9 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Thu, 2 Jan 2025 17:36:57 -0300 Subject: [PATCH] assistant2: Suggest current thread in inline assistant (#22586) Release Notes: - N/A --------- Co-authored-by: Marshall --- crates/assistant2/src/active_thread.rs | 2 +- crates/assistant2/src/assistant_panel.rs | 6 +- crates/assistant2/src/context.rs | 8 +- .../context_picker/thread_context_picker.rs | 24 +-- crates/assistant2/src/context_store.rs | 16 +- crates/assistant2/src/context_strip.rs | 196 ++++++++++-------- crates/assistant2/src/thread.rs | 21 ++ crates/assistant2/src/ui/context_pill.rs | 2 +- 8 files changed, 161 insertions(+), 114 deletions(-) diff --git a/crates/assistant2/src/active_thread.rs b/crates/assistant2/src/active_thread.rs index 4b49adb3293f6906f27d77a670faf393f42866c8..448a156deaad18a57f6220dfe41b4dcf8c75d95e 100644 --- a/crates/assistant2/src/active_thread.rs +++ b/crates/assistant2/src/active_thread.rs @@ -22,7 +22,7 @@ pub struct ActiveThread { workspace: WeakView, language_registry: Arc, tools: Arc, - thread: Model, + pub(crate) thread: Model, messages: Vec, list_state: ListState, rendered_messages_by_id: HashMap>, diff --git a/crates/assistant2/src/assistant_panel.rs b/crates/assistant2/src/assistant_panel.rs index 522d2b4eea60de3e449547ad379724b72c129f79..7e229251ef3ce48f490dfab8628c7caa2dc92d45 100644 --- a/crates/assistant2/src/assistant_panel.rs +++ b/crates/assistant2/src/assistant_panel.rs @@ -19,7 +19,7 @@ use workspace::Workspace; use crate::active_thread::ActiveThread; use crate::assistant_settings::{AssistantDockPosition, AssistantSettings}; use crate::message_editor::MessageEditor; -use crate::thread::{ThreadError, ThreadId}; +use crate::thread::{Thread, ThreadError, ThreadId}; use crate::thread_history::{PastThread, ThreadHistory}; use crate::thread_store::ThreadStore; use crate::{NewThread, OpenHistory, ToggleFocus}; @@ -206,6 +206,10 @@ impl AssistantPanel { self.message_editor.focus_handle(cx).focus(cx); } + pub(crate) fn active_thread(&self, cx: &AppContext) -> Model { + self.thread.read(cx).thread.clone() + } + pub(crate) fn delete_thread(&mut self, thread_id: &ThreadId, cx: &mut ViewContext) { self.thread_store .update(cx, |this, cx| this.delete_thread(thread_id, cx)); diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs index d08c96b0eef75529e7db2227d916cfd832b7149e..074ae18924357887ec6765ef7bda066419024964 100644 --- a/crates/assistant2/src/context.rs +++ b/crates/assistant2/src/context.rs @@ -4,6 +4,8 @@ use project::ProjectEntryId; use serde::{Deserialize, Serialize}; use util::post_inc; +use crate::thread::ThreadId; + #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, Serialize, Deserialize)] pub struct ContextId(pub(crate) usize); @@ -22,12 +24,12 @@ pub struct Context { pub text: SharedString, } -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum ContextKind { File(ProjectEntryId), Directory, FetchedUrl, - Thread, + Thread(ThreadId), } pub fn attach_context_to_message( @@ -55,7 +57,7 @@ pub fn attach_context_to_message( fetch_context.push_str(&context.text); fetch_context.push('\n'); } - ContextKind::Thread => { + ContextKind::Thread(_) => { thread_context.push_str(&context.name); thread_context.push('\n'); thread_context.push_str(&context.text); diff --git a/crates/assistant2/src/context_picker/thread_context_picker.rs b/crates/assistant2/src/context_picker/thread_context_picker.rs index e2bd1cff229146637d86e0e450c794e990cb931a..42a5eebccd5d394654c068a5fe7e8a528a22995b 100644 --- a/crates/assistant2/src/context_picker/thread_context_picker.rs +++ b/crates/assistant2/src/context_picker/thread_context_picker.rs @@ -169,25 +169,11 @@ impl PickerDelegate for ThreadContextPickerDelegate { self.context_store .update(cx, |context_store, cx| { - let text = thread.update(cx, |thread, _cx| { - let mut text = String::new(); - - for message in thread.messages() { - text.push_str(match message.role { - language_model::Role::User => "User:", - language_model::Role::Assistant => "Assistant:", - language_model::Role::System => "System:", - }); - text.push('\n'); - - text.push_str(&message.text); - text.push('\n'); - } - - text - }); - - context_store.insert_context(ContextKind::Thread, entry.summary.clone(), text); + context_store.insert_context( + ContextKind::Thread(thread.read(cx).id().clone()), + entry.summary.clone(), + thread.read(cx).text(), + ); }) .ok(); diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs index 8db5df0087469a65da196b12737cd4d64eec1c49..76fafdf797406392794f21b2760dcc90d751d91a 100644 --- a/crates/assistant2/src/context_store.rs +++ b/crates/assistant2/src/context_store.rs @@ -1,7 +1,10 @@ use gpui::SharedString; use project::ProjectEntryId; -use crate::context::{Context, ContextId, ContextKind}; +use crate::{ + context::{Context, ContextId, ContextKind}, + thread::ThreadId, +}; pub struct ContextStore { context: Vec, @@ -49,9 +52,14 @@ impl ContextStore { pub fn contains_project_entry(&self, entry_id: ProjectEntryId) -> bool { self.context.iter().any(|probe| match probe.kind { ContextKind::File(probe_entry_id) => probe_entry_id == entry_id, - ContextKind::Directory => false, - ContextKind::FetchedUrl => false, - ContextKind::Thread => false, + ContextKind::Directory | ContextKind::FetchedUrl | ContextKind::Thread(_) => false, + }) + } + + pub fn contains_thread(&self, thread_id: &ThreadId) -> bool { + self.context.iter().any(|probe| match probe.kind { + ContextKind::Thread(ref probe_thread_id) => probe_thread_id == thread_id, + ContextKind::File(_) | ContextKind::Directory | ContextKind::FetchedUrl => false, }) } } diff --git a/crates/assistant2/src/context_strip.rs b/crates/assistant2/src/context_strip.rs index a4c304fcbc6bc700ce98bd5c7554210306209dd0..b297ff323496ffcb0da551956c2f6c999cfa6862 100644 --- a/crates/assistant2/src/context_strip.rs +++ b/crates/assistant2/src/context_strip.rs @@ -1,18 +1,19 @@ use std::rc::Rc; use editor::Editor; -use gpui::{EntityId, FocusHandle, Model, Subscription, View, WeakModel, WeakView}; +use gpui::{AppContext, FocusHandle, Model, View, WeakModel, WeakView}; use language::Buffer; use project::ProjectEntryId; use ui::{prelude::*, PopoverMenu, PopoverMenuHandle, Tooltip}; -use workspace::{ItemHandle, Workspace}; +use workspace::Workspace; use crate::context::ContextKind; use crate::context_picker::{ConfirmBehavior, ContextPicker}; use crate::context_store::ContextStore; +use crate::thread::{Thread, ThreadId}; use crate::thread_store::ThreadStore; use crate::ui::ContextPill; -use crate::ToggleContextPicker; +use crate::{AssistantPanel, ToggleContextPicker}; use settings::Settings; pub struct ContextStrip { @@ -20,21 +21,8 @@ pub struct ContextStrip { context_picker: View, context_picker_menu_handle: PopoverMenuHandle, focus_handle: FocusHandle, - workspace_active_pane_id: Option, - suggested_context: Option, - _subscription: Option, -} - -pub enum SuggestContextKind { - File, - Thread, -} - -#[derive(Clone)] -pub struct SuggestedContext { - entry_id: ProjectEntryId, - title: SharedString, - buffer: WeakModel, + suggest_context_kind: SuggestContextKind, + workspace: WeakView, } impl ContextStrip { @@ -47,20 +35,6 @@ impl ContextStrip { suggest_context_kind: SuggestContextKind, cx: &mut ViewContext, ) -> Self { - let subscription = match suggest_context_kind { - SuggestContextKind::File => { - if let Some(workspace) = workspace.upgrade() { - Some(cx.subscribe(&workspace, Self::handle_workspace_event)) - } else { - None - } - } - SuggestContextKind::Thread => { - // TODO: Suggest current thread - None - } - }; - Self { context_store: context_store.clone(), context_picker: cx.new_view(|cx| { @@ -74,56 +48,65 @@ impl ContextStrip { }), context_picker_menu_handle, focus_handle, - workspace_active_pane_id: None, - suggested_context: None, - _subscription: subscription, + suggest_context_kind, + workspace, } } - fn handle_workspace_event( - &mut self, - workspace: View, - event: &workspace::Event, - cx: &mut ViewContext, - ) { - match event { - workspace::Event::WorkspaceCreated(_) | workspace::Event::ActiveItemChanged => { - let workspace = workspace.read(cx); - - if let Some(active_item) = workspace.active_item(cx) { - let new_active_item_id = Some(active_item.item_id()); - - if self.workspace_active_pane_id != new_active_item_id { - self.suggested_context = Self::suggested_file(active_item, cx); - self.workspace_active_pane_id = new_active_item_id; - } - } else { - self.suggested_context = None; - self.workspace_active_pane_id = None; - } - } - _ => {} + fn suggested_context(&self, cx: &ViewContext) -> Option { + match self.suggest_context_kind { + SuggestContextKind::File => self.suggested_file(cx), + SuggestContextKind::Thread => self.suggested_thread(cx), } } - fn suggested_file( - active_item: Box, - cx: &WindowContext, - ) -> Option { + fn suggested_file(&self, cx: &ViewContext) -> Option { + let workspace = self.workspace.upgrade()?; + let active_item = workspace.read(cx).active_item(cx)?; let entry_id = *active_item.project_entry_ids(cx).first()?; + if self.context_store.read(cx).contains_project_entry(entry_id) { + return None; + } + let editor = active_item.to_any().downcast::().ok()?.read(cx); let active_buffer = editor.buffer().read(cx).as_singleton()?; let file = active_buffer.read(cx).file()?; let title = file.path().to_string_lossy().into_owned().into(); - Some(SuggestedContext { + Some(SuggestedContext::File { entry_id, title, buffer: active_buffer.downgrade(), }) } + + fn suggested_thread(&self, cx: &ViewContext) -> Option { + let workspace = self.workspace.upgrade()?; + let active_thread = workspace + .read(cx) + .panel::(cx)? + .read(cx) + .active_thread(cx); + let weak_active_thread = active_thread.downgrade(); + + let active_thread = active_thread.read(cx); + + if self + .context_store + .read(cx) + .contains_thread(active_thread.id()) + { + return None; + } + + Some(SuggestedContext::Thread { + id: active_thread.id().clone(), + title: active_thread.summary().unwrap_or("Active Thread".into()), + thread: weak_active_thread, + }) + } } impl Render for ContextStrip { @@ -133,13 +116,7 @@ impl Render for ContextStrip { let context_picker = self.context_picker.clone(); let focus_handle = self.focus_handle.clone(); - let suggested_context = self.suggested_context.as_ref().and_then(|suggested| { - if context_store.contains_project_entry(suggested.entry_id) { - None - } else { - Some(suggested.clone()) - } - }); + let suggested_context = self.suggested_context(cx); h_flex() .flex_wrap() @@ -172,7 +149,7 @@ impl Render for ContextStrip { }) .with_handle(self.context_picker_menu_handle.clone()), ) - .when(context.is_empty() && self.suggested_context.is_none(), { + .when(context.is_empty() && suggested_context.is_none(), { |parent| { parent.child( h_flex() @@ -209,24 +186,13 @@ impl Render for ContextStrip { })) .when_some(suggested_context, |el, suggested| { el.child( - Button::new("add-suggested-context", suggested.title.clone()) + Button::new("add-suggested-context", suggested.title().clone()) .on_click({ let context_store = self.context_store.clone(); cx.listener(move |_this, _event, cx| { - let Some(buffer) = suggested.buffer.upgrade() else { - return; - }; - - let title = suggested.title.clone(); - let text = buffer.read(cx).text(); - - context_store.update(cx, move |context_store, _cx| { - context_store.insert_context( - ContextKind::File(suggested.entry_id), - title, - text, - ); + context_store.update(cx, |context_store, cx| { + suggested.accept(context_store, cx); }); cx.notify(); }) @@ -260,3 +226,63 @@ impl Render for ContextStrip { }) } } + +pub enum SuggestContextKind { + File, + Thread, +} + +#[derive(Clone)] +pub enum SuggestedContext { + File { + entry_id: ProjectEntryId, + title: SharedString, + buffer: WeakModel, + }, + Thread { + id: ThreadId, + title: SharedString, + thread: WeakModel, + }, +} + +impl SuggestedContext { + pub fn title(&self) -> &SharedString { + match self { + Self::File { title, .. } => title, + Self::Thread { title, .. } => title, + } + } + + pub fn accept(&self, context_store: &mut ContextStore, cx: &mut AppContext) { + match self { + Self::File { + entry_id, + title, + buffer, + } => { + let Some(buffer) = buffer.upgrade() else { + return; + }; + let text = buffer.read(cx).text(); + + context_store.insert_context( + ContextKind::File(*entry_id), + title.clone(), + text.clone(), + ); + } + Self::Thread { id, title, thread } => { + let Some(thread) = thread.upgrade() else { + return; + }; + + context_store.insert_context( + ContextKind::Thread(id.clone()), + title.clone(), + thread.read(cx).text(), + ); + } + } + } +} diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs index ae5a564c1ad1dc762c5040806b978464a893499e..f31a2b04c6fd3903424fb6d6f3ec01d76344d430 100644 --- a/crates/assistant2/src/thread.rs +++ b/crates/assistant2/src/thread.rs @@ -164,6 +164,27 @@ impl Thread { id } + /// Returns the representation of this [`Thread`] in a textual form. + /// + /// This is the representation we use when attaching a thread as context to another thread. + pub fn text(&self) -> String { + let mut text = String::new(); + + for message in &self.messages { + text.push_str(match message.role { + language_model::Role::User => "User:", + language_model::Role::Assistant => "Assistant:", + language_model::Role::System => "System:", + }); + text.push('\n'); + + text.push_str(&message.text); + text.push('\n'); + } + + text + } + pub fn to_completion_request( &self, _request_kind: RequestKind, diff --git a/crates/assistant2/src/ui/context_pill.rs b/crates/assistant2/src/ui/context_pill.rs index 8aadc9d6a6f205c569a9c2c13539594ac110693d..b5c79d3473efb4375489940ba4abe77a5df331a4 100644 --- a/crates/assistant2/src/ui/context_pill.rs +++ b/crates/assistant2/src/ui/context_pill.rs @@ -36,7 +36,7 @@ impl RenderOnce for ContextPill { ContextKind::File(_) => IconName::File, ContextKind::Directory => IconName::Folder, ContextKind::FetchedUrl => IconName::Globe, - ContextKind::Thread => IconName::MessageCircle, + ContextKind::Thread(_) => IconName::MessageCircle, }; h_flex()