From 59b5b9af905d7640270784eac9e53cad21a01353 Mon Sep 17 00:00:00 2001 From: Agus Zubiaga Date: Thu, 2 Jan 2025 10:33:47 -0300 Subject: [PATCH] assistant2: Suggest current file as context (#22526) Suggest adding the current file as context in the new assistant panel. https://github.com/user-attachments/assets/62bc267b-3dfe-4a3b-a6af-c89af2c779a8 Note: This doesn't include suggesting the current thread in the inline assistant. Release Notes: - N/A --- crates/assistant2/src/context.rs | 5 +- crates/assistant2/src/context_picker.rs | 27 +-- .../src/context_picker/file_context_picker.rs | 17 +- crates/assistant2/src/context_store.rs | 10 ++ crates/assistant2/src/context_strip.rs | 154 ++++++++++++++++-- crates/assistant2/src/inline_prompt_editor.rs | 4 +- crates/assistant2/src/message_editor.rs | 3 +- crates/assistant2/src/ui/context_pill.rs | 2 +- 8 files changed, 191 insertions(+), 31 deletions(-) diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs index 1a5a2eed95c7b490227d2ff4f05493b659321f7b..d08c96b0eef75529e7db2227d916cfd832b7149e 100644 --- a/crates/assistant2/src/context.rs +++ b/crates/assistant2/src/context.rs @@ -1,5 +1,6 @@ use gpui::SharedString; use language_model::{LanguageModelRequestMessage, MessageContent}; +use project::ProjectEntryId; use serde::{Deserialize, Serialize}; use util::post_inc; @@ -23,7 +24,7 @@ pub struct Context { #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum ContextKind { - File, + File(ProjectEntryId), Directory, FetchedUrl, Thread, @@ -40,7 +41,7 @@ pub fn attach_context_to_message( for context in context.into_iter() { match context.kind { - ContextKind::File => { + ContextKind::File(_) => { file_context.push_str(&context.text); file_context.push('\n'); } diff --git a/crates/assistant2/src/context_picker.rs b/crates/assistant2/src/context_picker.rs index 5b590ac47ac4ad084ad5f87054f545bf9ec968b0..ec3421d63f7321759cbd40fbe629f8087fd3dca9 100644 --- a/crates/assistant2/src/context_picker.rs +++ b/crates/assistant2/src/context_picker.rs @@ -15,7 +15,6 @@ use ui::{prelude::*, ListItem, ListItemSpacing}; use util::ResultExt; use workspace::Workspace; -use crate::context::ContextKind; use crate::context_picker::directory_context_picker::DirectoryContextPicker; use crate::context_picker::fetch_context_picker::FetchContextPicker; use crate::context_picker::file_context_picker::FileContextPicker; @@ -54,7 +53,7 @@ impl ContextPicker { let mut entries = Vec::new(); entries.push(ContextPickerEntry { name: "File".into(), - kind: ContextKind::File, + kind: ContextPickerEntryKind::File, icon: IconName::File, }); let release_channel = ReleaseChannel::global(cx); @@ -63,20 +62,20 @@ impl ContextPicker { if release_channel == ReleaseChannel::Dev { entries.push(ContextPickerEntry { name: "Folder".into(), - kind: ContextKind::Directory, + kind: ContextPickerEntryKind::Directory, icon: IconName::Folder, }); } entries.push(ContextPickerEntry { name: "Fetch".into(), - kind: ContextKind::FetchedUrl, + kind: ContextPickerEntryKind::FetchedUrl, icon: IconName::Globe, }); if thread_store.is_some() { entries.push(ContextPickerEntry { name: "Thread".into(), - kind: ContextKind::Thread, + kind: ContextPickerEntryKind::Thread, icon: IconName::MessageCircle, }); } @@ -140,10 +139,18 @@ impl Render for ContextPicker { #[derive(Clone)] struct ContextPickerEntry { name: SharedString, - kind: ContextKind, + kind: ContextPickerEntryKind, icon: IconName, } +#[derive(Debug, Clone)] +enum ContextPickerEntryKind { + File, + Directory, + FetchedUrl, + Thread, +} + pub(crate) struct ContextPickerDelegate { context_picker: WeakView, workspace: WeakView, @@ -183,7 +190,7 @@ impl PickerDelegate for ContextPickerDelegate { self.context_picker .update(cx, |this, cx| { match entry.kind { - ContextKind::File => { + ContextPickerEntryKind::File => { this.mode = ContextPickerMode::File(cx.new_view(|cx| { FileContextPicker::new( self.context_picker.clone(), @@ -194,7 +201,7 @@ impl PickerDelegate for ContextPickerDelegate { ) })); } - ContextKind::Directory => { + ContextPickerEntryKind::Directory => { this.mode = ContextPickerMode::Directory(cx.new_view(|cx| { DirectoryContextPicker::new( self.context_picker.clone(), @@ -205,7 +212,7 @@ impl PickerDelegate for ContextPickerDelegate { ) })); } - ContextKind::FetchedUrl => { + ContextPickerEntryKind::FetchedUrl => { this.mode = ContextPickerMode::Fetch(cx.new_view(|cx| { FetchContextPicker::new( self.context_picker.clone(), @@ -216,7 +223,7 @@ impl PickerDelegate for ContextPickerDelegate { ) })); } - ContextKind::Thread => { + ContextPickerEntryKind::Thread => { if let Some(thread_store) = self.thread_store.as_ref() { this.mode = ContextPickerMode::Thread(cx.new_view(|cx| { ThreadContextPicker::new( diff --git a/crates/assistant2/src/context_picker/file_context_picker.rs b/crates/assistant2/src/context_picker/file_context_picker.rs index db497b43a04fe2c4f3d515b712702ebe43db7307..c8258246e50016f773dc8e97d1f38fe803ad0cc6 100644 --- a/crates/assistant2/src/context_picker/file_context_picker.rs +++ b/crates/assistant2/src/context_picker/file_context_picker.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use fuzzy::PathMatch; use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakModel, WeakView}; use picker::{Picker, PickerDelegate}; -use project::{PathMatchCandidateSet, WorktreeId}; +use project::{PathMatchCandidateSet, ProjectPath, WorktreeId}; use ui::{prelude::*, ListItem}; use util::ResultExt as _; use workspace::Workspace; @@ -207,11 +207,20 @@ impl PickerDelegate for FileContextPickerDelegate { let worktree_id = WorktreeId::from_usize(mat.worktree_id); let confirm_behavior = self.confirm_behavior; cx.spawn(|this, mut cx| async move { - let Some(open_buffer_task) = project + let Some((entry_id, open_buffer_task)) = project .update(&mut cx, |project, cx| { - project.open_buffer((worktree_id, path.clone()), cx) + let project_path = ProjectPath { + worktree_id, + path: path.clone(), + }; + + let entry_id = project.entry_for_path(&project_path, cx)?.id; + let task = project.open_buffer(project_path, cx); + + Some((entry_id, task)) }) .ok() + .flatten() else { return anyhow::Ok(()); }; @@ -232,7 +241,7 @@ impl PickerDelegate for FileContextPickerDelegate { text.push_str("```\n"); context_store.insert_context( - ContextKind::File, + ContextKind::File(entry_id), path.to_string_lossy().to_string(), text, ); diff --git a/crates/assistant2/src/context_store.rs b/crates/assistant2/src/context_store.rs index febd1f597dbfad52e2f4f5209e0efe1a7de9ef24..8db5df0087469a65da196b12737cd4d64eec1c49 100644 --- a/crates/assistant2/src/context_store.rs +++ b/crates/assistant2/src/context_store.rs @@ -1,4 +1,5 @@ use gpui::SharedString; +use project::ProjectEntryId; use crate::context::{Context, ContextId, ContextKind}; @@ -44,4 +45,13 @@ impl ContextStore { pub fn remove_context(&mut self, id: &ContextId) { self.context.retain(|context| context.id != *id); } + + 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, + }) + } } diff --git a/crates/assistant2/src/context_strip.rs b/crates/assistant2/src/context_strip.rs index 7d7bda8b3570f322d1d42e363b2c708a6f634ec4..a4c304fcbc6bc700ce98bd5c7554210306209dd0 100644 --- a/crates/assistant2/src/context_strip.rs +++ b/crates/assistant2/src/context_strip.rs @@ -1,9 +1,13 @@ use std::rc::Rc; -use gpui::{FocusHandle, Model, View, WeakModel, WeakView}; +use editor::Editor; +use gpui::{EntityId, FocusHandle, Model, Subscription, View, WeakModel, WeakView}; +use language::Buffer; +use project::ProjectEntryId; use ui::{prelude::*, PopoverMenu, PopoverMenuHandle, Tooltip}; -use workspace::Workspace; +use workspace::{ItemHandle, Workspace}; +use crate::context::ContextKind; use crate::context_picker::{ConfirmBehavior, ContextPicker}; use crate::context_store::ContextStore; use crate::thread_store::ThreadStore; @@ -16,6 +20,21 @@ 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, } impl ContextStrip { @@ -25,8 +44,23 @@ impl ContextStrip { thread_store: Option>, focus_handle: FocusHandle, context_picker_menu_handle: PopoverMenuHandle, + 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| { @@ -40,16 +74,73 @@ impl ContextStrip { }), context_picker_menu_handle, focus_handle, + workspace_active_pane_id: None, + suggested_context: None, + _subscription: subscription, } } + + 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_file( + active_item: Box, + cx: &WindowContext, + ) -> Option { + let entry_id = *active_item.project_entry_ids(cx).first()?; + + 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 { + entry_id, + title, + buffer: active_buffer.downgrade(), + }) + } } impl Render for ContextStrip { fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { - let context = self.context_store.read(cx).context().clone(); + let context_store = self.context_store.read(cx); + let context = context_store.context().clone(); 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()) + } + }); + h_flex() .flex_wrap() .gap_1() @@ -60,13 +151,17 @@ impl Render for ContextStrip { IconButton::new("add-context", IconName::Plus) .icon_size(IconSize::Small) .style(ui::ButtonStyle::Filled) - .tooltip(move |cx| { - Tooltip::for_action_in( - "Add Context", - &ToggleContextPicker, - &focus_handle, - cx, - ) + .tooltip({ + let focus_handle = focus_handle.clone(); + + move |cx| { + Tooltip::for_action_in( + "Add Context", + &ToggleContextPicker, + &focus_handle, + cx, + ) + } }), ) .attach(gpui::Corner::TopLeft) @@ -77,7 +172,7 @@ impl Render for ContextStrip { }) .with_handle(self.context_picker_menu_handle.clone()), ) - .when(context.is_empty(), { + .when(context.is_empty() && self.suggested_context.is_none(), { |parent| { parent.child( h_flex() @@ -91,7 +186,7 @@ impl Render for ContextStrip { .children( ui::KeyBinding::for_action_in( &ToggleContextPicker, - &self.focus_handle, + &focus_handle, cx, ) .map(|binding| binding.into_any_element()), @@ -112,6 +207,41 @@ impl Render for ContextStrip { })) }) })) + .when_some(suggested_context, |el, suggested| { + el.child( + 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, + ); + }); + cx.notify(); + }) + }) + .icon(IconName::Plus) + .icon_position(IconPosition::Start) + .icon_size(IconSize::XSmall) + .icon_color(Color::Muted) + .label_size(LabelSize::Small) + .style(ButtonStyle::Filled) + .tooltip(|cx| { + Tooltip::with_meta("Suggested Context", None, "Click to add it", cx) + }), + ) + }) .when(!context.is_empty(), { move |parent| { parent.child( diff --git a/crates/assistant2/src/inline_prompt_editor.rs b/crates/assistant2/src/inline_prompt_editor.rs index 4df5fe33b55bf9ddb9b4768160bebccde3416920..f14ba83c175e343e1b2beaf00b0682d85653a935 100644 --- a/crates/assistant2/src/inline_prompt_editor.rs +++ b/crates/assistant2/src/inline_prompt_editor.rs @@ -2,7 +2,7 @@ use crate::assistant_model_selector::AssistantModelSelector; use crate::buffer_codegen::BufferCodegen; use crate::context_picker::ContextPicker; use crate::context_store::ContextStore; -use crate::context_strip::ContextStrip; +use crate::context_strip::{ContextStrip, SuggestContextKind}; use crate::terminal_codegen::TerminalCodegen; use crate::thread_store::ThreadStore; use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist}; @@ -793,6 +793,7 @@ impl PromptEditor { thread_store.clone(), prompt_editor.focus_handle(cx), context_picker_menu_handle.clone(), + SuggestContextKind::Thread, cx, ) }), @@ -932,6 +933,7 @@ impl PromptEditor { thread_store.clone(), prompt_editor.focus_handle(cx), context_picker_menu_handle.clone(), + SuggestContextKind::Thread, cx, ) }), diff --git a/crates/assistant2/src/message_editor.rs b/crates/assistant2/src/message_editor.rs index 3f31552e463130d0f680c02822bdca55145401bb..7b67d93091d51b2d24e66a4af3b7886728a9828b 100644 --- a/crates/assistant2/src/message_editor.rs +++ b/crates/assistant2/src/message_editor.rs @@ -20,7 +20,7 @@ use workspace::Workspace; use crate::assistant_model_selector::AssistantModelSelector; use crate::context_picker::{ConfirmBehavior, ContextPicker}; use crate::context_store::ContextStore; -use crate::context_strip::ContextStrip; +use crate::context_strip::{ContextStrip, SuggestContextKind}; use crate::thread::{RequestKind, Thread}; use crate::thread_store::ThreadStore; use crate::{Chat, ToggleContextPicker, ToggleModelSelector}; @@ -87,6 +87,7 @@ impl MessageEditor { Some(thread_store.clone()), editor.focus_handle(cx), context_picker_menu_handle.clone(), + SuggestContextKind::File, cx, ) }), diff --git a/crates/assistant2/src/ui/context_pill.rs b/crates/assistant2/src/ui/context_pill.rs index fb926386e2ac899a4cf9fa6d3a68269107b7d623..8aadc9d6a6f205c569a9c2c13539594ac110693d 100644 --- a/crates/assistant2/src/ui/context_pill.rs +++ b/crates/assistant2/src/ui/context_pill.rs @@ -33,7 +33,7 @@ impl RenderOnce for ContextPill { px(4.) }; let icon = match self.context.kind { - ContextKind::File => IconName::File, + ContextKind::File(_) => IconName::File, ContextKind::Directory => IconName::Folder, ContextKind::FetchedUrl => IconName::Globe, ContextKind::Thread => IconName::MessageCircle,