Cargo.lock 🔗
@@ -469,6 +469,7 @@ dependencies = [
"feature_flags",
"fs",
"futures 0.3.31",
+ "fuzzy",
"gpui",
"handlebars 4.5.0",
"indoc",
assistant2: Add initial support for attaching file context (#21934)
Marshall Bowers created
This PR adds the initial support for attaching files as context to a
thread in Assistant2.
Release Notes:
- N/AChange summary
Cargo.lock | 1
crates/assistant2/Cargo.toml | 1
crates/assistant2/src/assistant_panel.rs | 10
crates/assistant2/src/context_picker.rs | 226 ++--
crates/assistant2/src/context_picker/file_context_picker.rs | 289 +++++++
crates/assistant2/src/message_editor.rs | 65 +
6 files changed, 460 insertions(+), 132 deletions(-)
Detailed changes
crates/assistant2/Cargo.toml 🔗
@@ -28,6 +28,7 @@ editor.workspace = true
feature_flags.workspace = true
fs.workspace = true
futures.workspace = true
+fuzzy.workspace = true
gpui.workspace = true
handlebars.workspace = true
language.workspace = true
crates/assistant2/src/assistant_panel.rs 🔗
@@ -88,13 +88,13 @@ impl AssistantPanel {
thread: cx.new_view(|cx| {
ActiveThread::new(
thread.clone(),
- workspace,
+ workspace.clone(),
language_registry,
tools.clone(),
cx,
)
}),
- message_editor: cx.new_view(|cx| MessageEditor::new(thread.clone(), cx)),
+ message_editor: cx.new_view(|cx| MessageEditor::new(workspace, thread.clone(), cx)),
tools,
local_timezone: UtcOffset::from_whole_seconds(
chrono::Local::now().offset().local_minus_utc(),
@@ -123,7 +123,8 @@ impl AssistantPanel {
cx,
)
});
- self.message_editor = cx.new_view(|cx| MessageEditor::new(thread, cx));
+ self.message_editor =
+ cx.new_view(|cx| MessageEditor::new(self.workspace.clone(), thread, cx));
self.message_editor.focus_handle(cx).focus(cx);
}
@@ -145,7 +146,8 @@ impl AssistantPanel {
cx,
)
});
- self.message_editor = cx.new_view(|cx| MessageEditor::new(thread, cx));
+ self.message_editor =
+ cx.new_view(|cx| MessageEditor::new(self.workspace.clone(), thread, cx));
self.message_editor.focus_handle(cx).focus(cx);
}
crates/assistant2/src/context_picker.rs 🔗
@@ -1,15 +1,93 @@
+mod file_context_picker;
+
use std::sync::Arc;
-use gpui::{DismissEvent, SharedString, Task, WeakView};
-use picker::{Picker, PickerDelegate, PickerEditorPosition};
-use ui::{prelude::*, ListItem, ListItemSpacing, PopoverMenu, PopoverTrigger, Tooltip};
+use gpui::{
+ AppContext, DismissEvent, EventEmitter, FocusHandle, FocusableView, SharedString, Task, View,
+ WeakView,
+};
+use picker::{Picker, PickerDelegate};
+use ui::{prelude::*, ListItem, ListItemSpacing, Tooltip};
+use util::ResultExt;
+use workspace::Workspace;
+use crate::context_picker::file_context_picker::FileContextPicker;
use crate::message_editor::MessageEditor;
-#[derive(IntoElement)]
-pub(super) struct ContextPicker<T: PopoverTrigger> {
- message_editor: WeakView<MessageEditor>,
- trigger: T,
+#[derive(Debug, Clone)]
+enum ContextPickerMode {
+ Default,
+ File(View<FileContextPicker>),
+}
+
+pub(super) struct ContextPicker {
+ mode: ContextPickerMode,
+ picker: View<Picker<ContextPickerDelegate>>,
+}
+
+impl ContextPicker {
+ pub fn new(
+ workspace: WeakView<Workspace>,
+ message_editor: WeakView<MessageEditor>,
+ cx: &mut ViewContext<Self>,
+ ) -> Self {
+ let delegate = ContextPickerDelegate {
+ context_picker: cx.view().downgrade(),
+ workspace: workspace.clone(),
+ message_editor: message_editor.clone(),
+ entries: vec![
+ ContextPickerEntry {
+ name: "directory".into(),
+ description: "Insert any directory".into(),
+ icon: IconName::Folder,
+ },
+ ContextPickerEntry {
+ name: "file".into(),
+ description: "Insert any file".into(),
+ icon: IconName::File,
+ },
+ ContextPickerEntry {
+ name: "web".into(),
+ description: "Fetch content from URL".into(),
+ icon: IconName::Globe,
+ },
+ ],
+ selected_ix: 0,
+ };
+
+ let picker = cx.new_view(|cx| {
+ Picker::nonsearchable_uniform_list(delegate, cx).max_height(Some(rems(20.).into()))
+ });
+
+ ContextPicker {
+ mode: ContextPickerMode::Default,
+ picker,
+ }
+ }
+
+ pub fn reset_mode(&mut self) {
+ self.mode = ContextPickerMode::Default;
+ }
+}
+
+impl EventEmitter<DismissEvent> for ContextPicker {}
+
+impl FocusableView for ContextPicker {
+ fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
+ match &self.mode {
+ ContextPickerMode::Default => self.picker.focus_handle(cx),
+ ContextPickerMode::File(file_picker) => file_picker.focus_handle(cx),
+ }
+ }
+}
+
+impl Render for ContextPicker {
+ fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+ v_flex().min_w(px(400.)).map(|parent| match &self.mode {
+ ContextPickerMode::Default => parent.child(self.picker.clone()),
+ ContextPickerMode::File(file_picker) => parent.child(file_picker.clone()),
+ })
+ }
}
#[derive(Clone)]
@@ -20,26 +98,18 @@ struct ContextPickerEntry {
}
pub(crate) struct ContextPickerDelegate {
- all_entries: Vec<ContextPickerEntry>,
- filtered_entries: Vec<ContextPickerEntry>,
+ context_picker: WeakView<ContextPicker>,
+ workspace: WeakView<Workspace>,
message_editor: WeakView<MessageEditor>,
+ entries: Vec<ContextPickerEntry>,
selected_ix: usize,
}
-impl<T: PopoverTrigger> ContextPicker<T> {
- pub(crate) fn new(message_editor: WeakView<MessageEditor>, trigger: T) -> Self {
- ContextPicker {
- message_editor,
- trigger,
- }
- }
-}
-
impl PickerDelegate for ContextPickerDelegate {
type ListItem = ListItem;
fn match_count(&self) -> usize {
- self.filtered_entries.len()
+ self.entries.len()
}
fn selected_index(&self) -> usize {
@@ -47,7 +117,7 @@ impl PickerDelegate for ContextPickerDelegate {
}
fn set_selected_index(&mut self, ix: usize, cx: &mut ViewContext<Picker<Self>>) {
- self.selected_ix = ix.min(self.filtered_entries.len().saturating_sub(1));
+ self.selected_ix = ix.min(self.entries.len().saturating_sub(1));
cx.notify();
}
@@ -55,52 +125,41 @@ impl PickerDelegate for ContextPickerDelegate {
"Select a context source…".into()
}
- fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
- let all_commands = self.all_entries.clone();
- cx.spawn(|this, mut cx| async move {
- let filtered_commands = cx
- .background_executor()
- .spawn(async move {
- if query.is_empty() {
- all_commands
- } else {
- all_commands
- .into_iter()
- .filter(|model_info| {
- model_info
- .name
- .to_lowercase()
- .contains(&query.to_lowercase())
- })
- .collect()
- }
- })
- .await;
-
- this.update(&mut cx, |this, cx| {
- this.delegate.filtered_entries = filtered_commands;
- this.delegate.set_selected_index(0, cx);
- cx.notify();
- })
- .ok();
- })
+ fn update_matches(&mut self, _query: String, _cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
+ Task::ready(())
}
fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
- if let Some(entry) = self.filtered_entries.get(self.selected_ix) {
- self.message_editor
- .update(cx, |_message_editor, _cx| {
- println!("Insert context from {}", entry.name);
+ if let Some(entry) = self.entries.get(self.selected_ix) {
+ self.context_picker
+ .update(cx, |this, cx| {
+ match entry.name.to_string().as_str() {
+ "file" => {
+ this.mode = ContextPickerMode::File(cx.new_view(|cx| {
+ FileContextPicker::new(
+ self.context_picker.clone(),
+ self.workspace.clone(),
+ self.message_editor.clone(),
+ cx,
+ )
+ }));
+ }
+ _ => {}
+ }
+
+ cx.focus_self();
})
- .ok();
- cx.emit(DismissEvent);
+ .log_err();
}
}
- fn dismissed(&mut self, _cx: &mut ViewContext<Picker<Self>>) {}
-
- fn editor_position(&self) -> PickerEditorPosition {
- PickerEditorPosition::End
+ fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
+ self.context_picker
+ .update(cx, |this, cx| match this.mode {
+ ContextPickerMode::Default => cx.emit(DismissEvent),
+ ContextPickerMode::File(_) => {}
+ })
+ .log_err();
}
fn render_match(
@@ -109,7 +168,7 @@ impl PickerDelegate for ContextPickerDelegate {
selected: bool,
_cx: &mut ViewContext<Picker<Self>>,
) -> Option<Self::ListItem> {
- let entry = self.filtered_entries.get(ix)?;
+ let entry = &self.entries[ix];
Some(
ListItem::new(ix)
@@ -148,50 +207,3 @@ impl PickerDelegate for ContextPickerDelegate {
)
}
}
-
-impl<T: PopoverTrigger> RenderOnce for ContextPicker<T> {
- fn render(self, cx: &mut WindowContext) -> impl IntoElement {
- let entries = vec![
- ContextPickerEntry {
- name: "directory".into(),
- description: "Insert any directory".into(),
- icon: IconName::Folder,
- },
- ContextPickerEntry {
- name: "file".into(),
- description: "Insert any file".into(),
- icon: IconName::File,
- },
- ContextPickerEntry {
- name: "web".into(),
- description: "Fetch content from URL".into(),
- icon: IconName::Globe,
- },
- ];
-
- let delegate = ContextPickerDelegate {
- all_entries: entries.clone(),
- message_editor: self.message_editor.clone(),
- filtered_entries: entries,
- selected_ix: 0,
- };
-
- let picker =
- cx.new_view(|cx| Picker::uniform_list(delegate, cx).max_height(Some(rems(20.).into())));
-
- let handle = self
- .message_editor
- .update(cx, |this, _| this.context_picker_handle.clone())
- .ok();
- PopoverMenu::new("context-picker")
- .menu(move |_cx| Some(picker.clone()))
- .trigger(self.trigger)
- .attach(gpui::AnchorCorner::TopLeft)
- .anchor(gpui::AnchorCorner::BottomLeft)
- .offset(gpui::Point {
- x: px(0.0),
- y: px(-16.0),
- })
- .when_some(handle, |this, handle| this.with_handle(handle))
- }
-}
crates/assistant2/src/context_picker/file_context_picker.rs 🔗
@@ -0,0 +1,289 @@
+use std::fmt::Write as _;
+use std::ops::RangeInclusive;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicBool;
+use std::sync::Arc;
+
+use fuzzy::PathMatch;
+use gpui::{AppContext, DismissEvent, FocusHandle, FocusableView, Task, View, WeakView};
+use picker::{Picker, PickerDelegate};
+use project::{PathMatchCandidateSet, WorktreeId};
+use ui::{prelude::*, ListItem, ListItemSpacing};
+use util::ResultExt as _;
+use workspace::Workspace;
+
+use crate::context::ContextKind;
+use crate::context_picker::ContextPicker;
+use crate::message_editor::MessageEditor;
+
+pub struct FileContextPicker {
+ picker: View<Picker<FileContextPickerDelegate>>,
+}
+
+impl FileContextPicker {
+ pub fn new(
+ context_picker: WeakView<ContextPicker>,
+ workspace: WeakView<Workspace>,
+ message_editor: WeakView<MessageEditor>,
+ cx: &mut ViewContext<Self>,
+ ) -> Self {
+ let delegate = FileContextPickerDelegate::new(context_picker, workspace, message_editor);
+ let picker = cx.new_view(|cx| Picker::uniform_list(delegate, cx));
+
+ Self { picker }
+ }
+}
+
+impl FocusableView for FileContextPicker {
+ fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
+ self.picker.focus_handle(cx)
+ }
+}
+
+impl Render for FileContextPicker {
+ fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+ self.picker.clone()
+ }
+}
+
+pub struct FileContextPickerDelegate {
+ context_picker: WeakView<ContextPicker>,
+ workspace: WeakView<Workspace>,
+ message_editor: WeakView<MessageEditor>,
+ matches: Vec<PathMatch>,
+ selected_index: usize,
+}
+
+impl FileContextPickerDelegate {
+ pub fn new(
+ context_picker: WeakView<ContextPicker>,
+ workspace: WeakView<Workspace>,
+ message_editor: WeakView<MessageEditor>,
+ ) -> Self {
+ Self {
+ context_picker,
+ workspace,
+ message_editor,
+ matches: Vec::new(),
+ selected_index: 0,
+ }
+ }
+
+ fn search(
+ &mut self,
+ query: String,
+ cancellation_flag: Arc<AtomicBool>,
+ workspace: &View<Workspace>,
+ cx: &mut ViewContext<Picker<Self>>,
+ ) -> Task<Vec<PathMatch>> {
+ if query.is_empty() {
+ let workspace = workspace.read(cx);
+ let project = workspace.project().read(cx);
+ let entries = workspace.recent_navigation_history(Some(10), cx);
+
+ let entries = entries
+ .into_iter()
+ .map(|entries| entries.0)
+ .chain(project.worktrees(cx).flat_map(|worktree| {
+ let worktree = worktree.read(cx);
+ let id = worktree.id();
+ worktree
+ .child_entries(Path::new(""))
+ .filter(|entry| entry.kind.is_file())
+ .map(move |entry| project::ProjectPath {
+ worktree_id: id,
+ path: entry.path.clone(),
+ })
+ }))
+ .collect::<Vec<_>>();
+
+ let path_prefix: Arc<str> = Arc::default();
+ Task::ready(
+ entries
+ .into_iter()
+ .filter_map(|entry| {
+ let worktree = project.worktree_for_id(entry.worktree_id, cx)?;
+ let mut full_path = PathBuf::from(worktree.read(cx).root_name());
+ full_path.push(&entry.path);
+ Some(PathMatch {
+ score: 0.,
+ positions: Vec::new(),
+ worktree_id: entry.worktree_id.to_usize(),
+ path: full_path.into(),
+ path_prefix: path_prefix.clone(),
+ distance_to_relative_ancestor: 0,
+ is_dir: false,
+ })
+ })
+ .collect(),
+ )
+ } else {
+ let worktrees = workspace.read(cx).visible_worktrees(cx).collect::<Vec<_>>();
+ let candidate_sets = worktrees
+ .into_iter()
+ .map(|worktree| {
+ let worktree = worktree.read(cx);
+
+ PathMatchCandidateSet {
+ snapshot: worktree.snapshot(),
+ include_ignored: worktree
+ .root_entry()
+ .map_or(false, |entry| entry.is_ignored),
+ include_root_name: true,
+ candidates: project::Candidates::Files,
+ }
+ })
+ .collect::<Vec<_>>();
+
+ let executor = cx.background_executor().clone();
+ cx.foreground_executor().spawn(async move {
+ fuzzy::match_path_sets(
+ candidate_sets.as_slice(),
+ query.as_str(),
+ None,
+ false,
+ 100,
+ &cancellation_flag,
+ executor,
+ )
+ .await
+ })
+ }
+ }
+}
+
+impl PickerDelegate for FileContextPickerDelegate {
+ type ListItem = ListItem;
+
+ fn match_count(&self) -> usize {
+ self.matches.len()
+ }
+
+ fn selected_index(&self) -> usize {
+ self.selected_index
+ }
+
+ fn set_selected_index(&mut self, ix: usize, _cx: &mut ViewContext<Picker<Self>>) {
+ self.selected_index = ix;
+ }
+
+ fn placeholder_text(&self, _cx: &mut WindowContext) -> Arc<str> {
+ "Search files…".into()
+ }
+
+ fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
+ let Some(workspace) = self.workspace.upgrade() else {
+ return Task::ready(());
+ };
+
+ let search_task = self.search(query, Arc::<AtomicBool>::default(), &workspace, cx);
+
+ cx.spawn(|this, mut cx| async move {
+ // TODO: This should be probably be run in the background.
+ let paths = search_task.await;
+
+ this.update(&mut cx, |this, _cx| {
+ this.delegate.matches = paths;
+ })
+ .log_err();
+ })
+ }
+
+ fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
+ let mat = &self.matches[self.selected_index];
+
+ let workspace = self.workspace.clone();
+ let Some(project) = workspace
+ .upgrade()
+ .map(|workspace| workspace.read(cx).project().clone())
+ else {
+ return;
+ };
+ let path = mat.path.clone();
+ let worktree_id = WorktreeId::from_usize(mat.worktree_id);
+ cx.spawn(|this, mut cx| async move {
+ let Some(open_buffer_task) = project
+ .update(&mut cx, |project, cx| {
+ project.open_buffer((worktree_id, path.clone()), cx)
+ })
+ .ok()
+ else {
+ return anyhow::Ok(());
+ };
+
+ let buffer = open_buffer_task.await?;
+
+ this.update(&mut cx, |this, cx| {
+ this.delegate
+ .message_editor
+ .update(cx, |message_editor, cx| {
+ let mut text = String::new();
+ text.push_str(&codeblock_fence_for_path(Some(&path), None));
+ text.push_str(&buffer.read(cx).text());
+ if !text.ends_with('\n') {
+ text.push('\n');
+ }
+
+ text.push_str("```\n");
+
+ message_editor.insert_context(
+ ContextKind::File,
+ path.to_string_lossy().to_string(),
+ text,
+ );
+ })
+ })??;
+
+ anyhow::Ok(())
+ })
+ .detach_and_log_err(cx);
+ }
+
+ fn dismissed(&mut self, cx: &mut ViewContext<Picker<Self>>) {
+ self.context_picker
+ .update(cx, |this, cx| {
+ this.reset_mode();
+ cx.emit(DismissEvent);
+ })
+ .log_err();
+ }
+
+ fn render_match(
+ &self,
+ ix: usize,
+ selected: bool,
+ _cx: &mut ViewContext<Picker<Self>>,
+ ) -> Option<Self::ListItem> {
+ let mat = &self.matches[ix];
+
+ Some(
+ ListItem::new(ix)
+ .inset(true)
+ .spacing(ListItemSpacing::Sparse)
+ .selected(selected)
+ .child(mat.path.to_string_lossy().to_string()),
+ )
+ }
+}
+
+fn codeblock_fence_for_path(path: Option<&Path>, row_range: Option<RangeInclusive<u32>>) -> String {
+ let mut text = String::new();
+ write!(text, "```").unwrap();
+
+ if let Some(path) = path {
+ if let Some(extension) = path.extension().and_then(|ext| ext.to_str()) {
+ write!(text, "{} ", extension).unwrap();
+ }
+
+ write!(text, "{}", path.display()).unwrap();
+ } else {
+ write!(text, "untitled").unwrap();
+ }
+
+ if let Some(row_range) = row_range {
+ write!(text, ":{}-{}", row_range.start() + 1, row_range.end() + 1).unwrap();
+ }
+
+ text.push('\n');
+ text
+}
crates/assistant2/src/message_editor.rs 🔗
@@ -1,19 +1,19 @@
use std::rc::Rc;
use editor::{Editor, EditorElement, EditorStyle};
-use gpui::{AppContext, FocusableView, Model, TextStyle, View};
+use gpui::{AppContext, FocusableView, Model, TextStyle, View, WeakView};
use language_model::{LanguageModelRegistry, LanguageModelRequestTool};
use language_model_selector::LanguageModelSelector;
-use picker::Picker;
use settings::Settings;
use theme::ThemeSettings;
use ui::{
prelude::*, ButtonLike, CheckboxWithLabel, ElevationIndex, IconButtonShape, KeyBinding,
- PopoverMenuHandle, Tooltip,
+ PopoverMenu, PopoverMenuHandle, Tooltip,
};
+use workspace::Workspace;
use crate::context::{Context, ContextId, ContextKind};
-use crate::context_picker::{ContextPicker, ContextPickerDelegate};
+use crate::context_picker::ContextPicker;
use crate::thread::{RequestKind, Thread};
use crate::ui::ContextPill;
use crate::{Chat, ToggleModelSelector};
@@ -23,13 +23,19 @@ pub struct MessageEditor {
editor: View<Editor>,
context: Vec<Context>,
next_context_id: ContextId,
- pub(crate) context_picker_handle: PopoverMenuHandle<Picker<ContextPickerDelegate>>,
+ context_picker: View<ContextPicker>,
+ pub(crate) context_picker_handle: PopoverMenuHandle<ContextPicker>,
use_tools: bool,
}
impl MessageEditor {
- pub fn new(thread: Model<Thread>, cx: &mut ViewContext<Self>) -> Self {
- let mut this = Self {
+ pub fn new(
+ workspace: WeakView<Workspace>,
+ thread: Model<Thread>,
+ cx: &mut ViewContext<Self>,
+ ) -> Self {
+ let weak_self = cx.view().downgrade();
+ Self {
thread,
editor: cx.new_view(|cx| {
let mut editor = Editor::auto_height(80, cx);
@@ -39,18 +45,24 @@ impl MessageEditor {
}),
context: Vec::new(),
next_context_id: ContextId(0),
+ context_picker: cx.new_view(|cx| ContextPicker::new(workspace.clone(), weak_self, cx)),
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(),
+ pub fn insert_context(
+ &mut self,
+ kind: ContextKind,
+ name: impl Into<SharedString>,
+ text: impl Into<SharedString>,
+ ) {
+ self.context.push(Context {
+ id: self.next_context_id.post_inc(),
+ name: name.into(),
+ kind,
+ text: text.into(),
});
-
- this
}
fn chat(&mut self, _: &Chat, cx: &mut ViewContext<Self>) {
@@ -167,6 +179,7 @@ impl Render for MessageEditor {
let font_size = TextSize::Default.rems(cx);
let line_height = font_size.to_pixels(cx.rem_size()) * 1.3;
let focus_handle = self.editor.focus_handle(cx);
+ let context_picker = self.context_picker.clone();
v_flex()
.key_context("MessageEditor")
@@ -179,12 +192,22 @@ impl Render for MessageEditor {
h_flex()
.flex_wrap()
.gap_2()
- .child(ContextPicker::new(
- cx.view().downgrade(),
- IconButton::new("add-context", IconName::Plus)
- .shape(IconButtonShape::Square)
- .icon_size(IconSize::Small),
- ))
+ .child(
+ PopoverMenu::new("context-picker")
+ .menu(move |_cx| Some(context_picker.clone()))
+ .trigger(
+ IconButton::new("add-context", IconName::Plus)
+ .shape(IconButtonShape::Square)
+ .icon_size(IconSize::Small),
+ )
+ .attach(gpui::AnchorCorner::TopLeft)
+ .anchor(gpui::AnchorCorner::BottomLeft)
+ .offset(gpui::Point {
+ x: px(0.0),
+ y: px(-16.0),
+ })
+ .with_handle(self.context_picker_handle.clone()),
+ )
.children(self.context.iter().map(|context| {
ContextPill::new(context.clone()).on_remove({
let context = context.clone();