From 92fb38acb6d5b03f1eb2386b753c5119a2f4bf6a Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Mon, 16 Dec 2024 16:22:16 -0500 Subject: [PATCH] assistant2: Wire up context for terminal inline assist (#22108) This PR updates up the context picker for the terminal's inline assist. Release Notes: - N/A Co-authored-by: Richard Co-authored-by: Agus --- crates/assistant2/src/inline_assistant.rs | 10 +- .../src/terminal_inline_assistant.rs | 156 +++++++++++------- 2 files changed, 98 insertions(+), 68 deletions(-) diff --git a/crates/assistant2/src/inline_assistant.rs b/crates/assistant2/src/inline_assistant.rs index 61d415a2a82dfa1bf80d52c490b347506d70bf2a..4f260518ae44e6546d47a8c35668ac861651d237 100644 --- a/crates/assistant2/src/inline_assistant.rs +++ b/crates/assistant2/src/inline_assistant.rs @@ -221,19 +221,19 @@ impl InlineAssistant { .map_or(false, |provider| provider.is_authenticated(cx)) }; + let thread_store = workspace + .panel::(cx) + .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade()); + let handle_assist = |cx: &mut ViewContext| match inline_assist_target { InlineAssistTarget::Editor(active_editor) => { InlineAssistant::update_global(cx, |assistant, cx| { - let thread_store = workspace - .panel::(cx) - .map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade()); - assistant.assist(&active_editor, cx.view().downgrade(), thread_store, cx) }) } InlineAssistTarget::Terminal(active_terminal) => { TerminalInlineAssistant::update_global(cx, |assistant, cx| { - assistant.assist(&active_terminal, cx.view().downgrade(), cx) + assistant.assist(&active_terminal, cx.view().downgrade(), thread_store, cx) }) } }; diff --git a/crates/assistant2/src/terminal_inline_assistant.rs b/crates/assistant2/src/terminal_inline_assistant.rs index e3b81979a86ebff88c9327abf38b14397dd104a7..48851bf82c61197c5d65488f4c19c0a9c46b30dd 100644 --- a/crates/assistant2/src/terminal_inline_assistant.rs +++ b/crates/assistant2/src/terminal_inline_assistant.rs @@ -1,5 +1,9 @@ use crate::assistant_settings::AssistantSettings; +use crate::context::attach_context_to_message; +use crate::context_store::ContextStore; +use crate::context_strip::ContextStrip; use crate::prompts::PromptBuilder; +use crate::thread_store::ThreadStore; use anyhow::{Context as _, Result}; use client::telemetry::Telemetry; use collections::{HashMap, VecDeque}; @@ -11,7 +15,7 @@ use fs::Fs; use futures::{channel::mpsc, SinkExt, StreamExt}; use gpui::{ AppContext, Context, EventEmitter, FocusHandle, FocusableView, Global, Model, ModelContext, - Subscription, Task, TextStyle, UpdateGlobal, View, WeakView, + Subscription, Task, TextStyle, UpdateGlobal, View, WeakModel, WeakView, }; use language::Buffer; use language_model::{ @@ -83,6 +87,7 @@ impl TerminalInlineAssistant { &mut self, terminal_view: &View, workspace: WeakView, + thread_store: Option>, cx: &mut WindowContext, ) { let terminal = terminal_view.read(cx).terminal().clone(); @@ -90,6 +95,7 @@ impl TerminalInlineAssistant { let prompt_buffer = cx.new_model(|cx| { MultiBuffer::singleton(cx.new_model(|cx| Buffer::local(String::new(), cx)), cx) }); + let context_store = cx.new_model(|_cx| ContextStore::new()); let codegen = cx.new_model(|_| Codegen::new(terminal, self.telemetry.clone())); let prompt_editor = cx.new_view(|cx| { @@ -99,6 +105,9 @@ impl TerminalInlineAssistant { prompt_buffer.clone(), codegen, self.fs.clone(), + context_store.clone(), + workspace.clone(), + thread_store.clone(), cx, ) }); @@ -116,6 +125,7 @@ impl TerminalInlineAssistant { terminal_view, prompt_editor, workspace.clone(), + context_store, cx, ); @@ -246,12 +256,21 @@ impl TerminalInlineAssistant { &latest_output, )?; + let mut request_message = LanguageModelRequestMessage { + role: Role::User, + content: vec![], + cache: false, + }; + + let context = assist + .context_store + .update(cx, |this, _cx| this.context().clone()); + attach_context_to_message(&mut request_message, context); + + request_message.content.push(prompt.into()); + Ok(LanguageModelRequest { - messages: vec![LanguageModelRequestMessage { - role: Role::User, - content: vec![prompt.into()], - cache: false, - }], + messages: vec![request_message], tools: Vec::new(), stop: Vec::new(), temperature: None, @@ -362,6 +381,7 @@ struct TerminalInlineAssist { prompt_editor: Option>, codegen: Model, workspace: WeakView, + context_store: Model, _subscriptions: Vec, } @@ -371,6 +391,7 @@ impl TerminalInlineAssist { terminal: &View, prompt_editor: View, workspace: WeakView, + context_store: Model, cx: &mut WindowContext, ) -> Self { let codegen = prompt_editor.read(cx).codegen.clone(); @@ -379,6 +400,7 @@ impl TerminalInlineAssist { prompt_editor: Some(prompt_editor.clone()), codegen: codegen.clone(), workspace: workspace.clone(), + context_store, _subscriptions: vec![ cx.subscribe(&prompt_editor, |prompt_editor, event, cx| { TerminalInlineAssistant::update_global(cx, |this, cx| { @@ -437,6 +459,7 @@ struct PromptEditor { id: TerminalInlineAssistId, height_in_lines: u8, editor: View, + context_strip: View, language_model_selector: View, edited_since_done: bool, prompt_history: VecDeque, @@ -452,11 +475,7 @@ impl EventEmitter for PromptEditor {} impl Render for PromptEditor { fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { let status = &self.codegen.read(cx).status; - let mut buttons = vec![Button::new("add-context", "Add Context") - .style(ButtonStyle::Filled) - .icon(IconName::Plus) - .icon_position(IconPosition::Start) - .into_any_element()]; + let mut buttons = Vec::new(); buttons.extend(match status { CodegenStatus::Idle => vec![ @@ -554,64 +573,69 @@ impl Render for PromptEditor { } }); - h_flex() - .bg(cx.theme().colors().editor_background) + v_flex() .border_y_1() .border_color(cx.theme().status().info_border) .py_2() - .h_full() - .w_full() - .on_action(cx.listener(Self::confirm)) - .on_action(cx.listener(Self::secondary_confirm)) - .on_action(cx.listener(Self::cancel)) - .on_action(cx.listener(Self::move_up)) - .on_action(cx.listener(Self::move_down)) + .size_full() .child( h_flex() - .w_12() - .justify_center() - .gap_2() - .child(LanguageModelSelectorPopoverMenu::new( - self.language_model_selector.clone(), - IconButton::new("context", IconName::SettingsAlt) - .shape(IconButtonShape::Square) - .icon_size(IconSize::Small) - .icon_color(Color::Muted) - .tooltip(move |cx| { - Tooltip::with_meta( - format!( - "Using {}", - LanguageModelRegistry::read_global(cx) - .active_model() - .map(|model| model.name().0) - .unwrap_or_else(|| "No model selected".into()), - ), - None, - "Change Model", - cx, - ) - }), - )) - .children( - if let CodegenStatus::Error(error) = &self.codegen.read(cx).status { - let error_message = SharedString::from(error.to_string()); - Some( - div() - .id("error") - .tooltip(move |cx| Tooltip::text(error_message.clone(), cx)) - .child( - Icon::new(IconName::XCircle) - .size(IconSize::Small) - .color(Color::Error), - ), - ) - } else { - None - }, - ), + .bg(cx.theme().colors().editor_background) + .on_action(cx.listener(Self::confirm)) + .on_action(cx.listener(Self::secondary_confirm)) + .on_action(cx.listener(Self::cancel)) + .on_action(cx.listener(Self::move_up)) + .on_action(cx.listener(Self::move_down)) + .child( + h_flex() + .w_12() + .justify_center() + .gap_2() + .child(LanguageModelSelectorPopoverMenu::new( + self.language_model_selector.clone(), + IconButton::new("context", IconName::SettingsAlt) + .shape(IconButtonShape::Square) + .icon_size(IconSize::Small) + .icon_color(Color::Muted) + .tooltip(move |cx| { + Tooltip::with_meta( + format!( + "Using {}", + LanguageModelRegistry::read_global(cx) + .active_model() + .map(|model| model.name().0) + .unwrap_or_else(|| "No model selected".into()), + ), + None, + "Change Model", + cx, + ) + }), + )) + .children( + if let CodegenStatus::Error(error) = &self.codegen.read(cx).status { + let error_message = SharedString::from(error.to_string()); + Some( + div() + .id("error") + .tooltip(move |cx| { + Tooltip::text(error_message.clone(), cx) + }) + .child( + Icon::new(IconName::XCircle) + .size(IconSize::Small) + .color(Color::Error), + ), + ) + } else { + None + }, + ), + ) + .child(div().flex_1().child(self.render_prompt_editor(cx))) + .child(h_flex().gap_1().pr_4().children(buttons)), ) - .child(div().flex_1().child(self.render_prompt_editor(cx))) - .child(h_flex().gap_1().pr_4().children(buttons)) + .child(h_flex().child(self.context_strip.clone())) } } @@ -631,6 +655,9 @@ impl PromptEditor { prompt_buffer: Model, codegen: Model, fs: Arc, + context_store: Model, + workspace: WeakView, + thread_store: Option>, cx: &mut ViewContext, ) -> Self { let prompt_editor = cx.new_view(|cx| { @@ -652,6 +679,9 @@ impl PromptEditor { id, height_in_lines: 1, editor: prompt_editor, + context_strip: cx.new_view(|cx| { + ContextStrip::new(context_store, workspace.clone(), thread_store.clone(), cx) + }), language_model_selector: cx.new_view(|cx| { let fs = fs.clone(); LanguageModelSelector::new(