Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

assistant2: Wire up context for terminal inline assist (#22108)

Browse source 
This PR updates up the context picker for the terminal's inline assist.

Release Notes:

- N/A

Co-authored-by: Richard <richard@zed.dev>
Co-authored-by: Agus <agus@zed.dev>
This commit is contained in:
Marshall Bowers 2024-12-16 16:22:16 -05:00 • committed by GitHub
parent 84392fbc2f
commit 92fb38acb6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 98 additions and 68 deletions
Download patch file Download diff file Expand all files Collapse all files

10
crates/assistant2/src/inline_assistant.rs
View file

@ -221,19 +221,19 @@ impl InlineAssistant {
.map_or(false, |provider| provider.is_authenticated(cx)) .map_or(false, |provider| provider.is_authenticated(cx))
}; };
let thread_store = workspace
.panel::<AssistantPanel>(cx)
.map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
let handle_assist = |cx: &mut ViewContext<Workspace>| match inline_assist_target { let handle_assist = |cx: &mut ViewContext<Workspace>| match inline_assist_target {
InlineAssistTarget::Editor(active_editor) => { InlineAssistTarget::Editor(active_editor) => {
InlineAssistant::update_global(cx, |assistant, cx| { InlineAssistant::update_global(cx, |assistant, cx| {
let thread_store = workspace
.panel::<AssistantPanel>(cx)
.map(|assistant_panel| assistant_panel.read(cx).thread_store().downgrade());
assistant.assist(&active_editor, cx.view().downgrade(), thread_store, cx) assistant.assist(&active_editor, cx.view().downgrade(), thread_store, cx)
}) })
} }
InlineAssistTarget::Terminal(active_terminal) => { InlineAssistTarget::Terminal(active_terminal) => {
TerminalInlineAssistant::update_global(cx, |assistant, cx| { TerminalInlineAssistant::update_global(cx, |assistant, cx| {
assistant.assist(&active_terminal, cx.view().downgrade(), cx) assistant.assist(&active_terminal, cx.view().downgrade(), thread_store, cx)
}) })
} }
}; };

156
crates/assistant2/src/terminal_inline_assistant.rs
View file

@ -1,5 +1,9 @@
use crate::assistant_settings::AssistantSettings; 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::prompts::PromptBuilder;
use crate::thread_store::ThreadStore;
use anyhow::{Context as _, Result}; use anyhow::{Context as _, Result};
use client::telemetry::Telemetry; use client::telemetry::Telemetry;
use collections::{HashMap, VecDeque}; use collections::{HashMap, VecDeque};
@ -11,7 +15,7 @@ use fs::Fs;
use futures::{channel::mpsc, SinkExt, StreamExt}; use futures::{channel::mpsc, SinkExt, StreamExt};
use gpui::{ use gpui::{
AppContext, Context, EventEmitter, FocusHandle, FocusableView, Global, Model, ModelContext, 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::Buffer;
use language_model::{ use language_model::{
@ -83,6 +87,7 @@ impl TerminalInlineAssistant {
&mut self, &mut self,
terminal_view: &View<TerminalView>, terminal_view: &View<TerminalView>,
workspace: WeakView<Workspace>, workspace: WeakView<Workspace>,
thread_store: Option<WeakModel<ThreadStore>>,
cx: &mut WindowContext, cx: &mut WindowContext,
) { ) {
let terminal = terminal_view.read(cx).terminal().clone(); let terminal = terminal_view.read(cx).terminal().clone();
@ -90,6 +95,7 @@ impl TerminalInlineAssistant {
let prompt_buffer = cx.new_model(|cx| { let prompt_buffer = cx.new_model(|cx| {
MultiBuffer::singleton(cx.new_model(|cx| Buffer::local(String::new(), cx)), 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 codegen = cx.new_model(|_| Codegen::new(terminal, self.telemetry.clone()));
let prompt_editor = cx.new_view(|cx| { let prompt_editor = cx.new_view(|cx| {
@ -99,6 +105,9 @@ impl TerminalInlineAssistant {
prompt_buffer.clone(), prompt_buffer.clone(),
codegen, codegen,
self.fs.clone(), self.fs.clone(),
context_store.clone(),
workspace.clone(),
thread_store.clone(),
cx, cx,
) )
}); });
@ -116,6 +125,7 @@ impl TerminalInlineAssistant {
terminal_view, terminal_view,
prompt_editor, prompt_editor,
workspace.clone(), workspace.clone(),
context_store,
cx, cx,
); );
@ -246,12 +256,21 @@ impl TerminalInlineAssistant {
&latest_output, &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 { Ok(LanguageModelRequest {
messages: vec![LanguageModelRequestMessage { messages: vec![request_message],
role: Role::User,
content: vec![prompt.into()],
cache: false,
}],
tools: Vec::new(), tools: Vec::new(),
stop: Vec::new(), stop: Vec::new(),
temperature: None, temperature: None,
@ -362,6 +381,7 @@ struct TerminalInlineAssist {
prompt_editor: Option<View<PromptEditor>>, prompt_editor: Option<View<PromptEditor>>,
codegen: Model<Codegen>, codegen: Model<Codegen>,
workspace: WeakView<Workspace>, workspace: WeakView<Workspace>,
context_store: Model<ContextStore>,
_subscriptions: Vec<Subscription>, _subscriptions: Vec<Subscription>,
} }
@ -371,6 +391,7 @@ impl TerminalInlineAssist {
terminal: &View<TerminalView>, terminal: &View<TerminalView>,
prompt_editor: View<PromptEditor>, prompt_editor: View<PromptEditor>,
workspace: WeakView<Workspace>, workspace: WeakView<Workspace>,
context_store: Model<ContextStore>,
cx: &mut WindowContext, cx: &mut WindowContext,
) -> Self { ) -> Self {
let codegen = prompt_editor.read(cx).codegen.clone(); let codegen = prompt_editor.read(cx).codegen.clone();
@ -379,6 +400,7 @@ impl TerminalInlineAssist {
prompt_editor: Some(prompt_editor.clone()), prompt_editor: Some(prompt_editor.clone()),
codegen: codegen.clone(), codegen: codegen.clone(),
workspace: workspace.clone(), workspace: workspace.clone(),
context_store,
_subscriptions: vec![ _subscriptions: vec![
cx.subscribe(&prompt_editor, |prompt_editor, event, cx| { cx.subscribe(&prompt_editor, |prompt_editor, event, cx| {
TerminalInlineAssistant::update_global(cx, |this, cx| { TerminalInlineAssistant::update_global(cx, |this, cx| {
@ -437,6 +459,7 @@ struct PromptEditor {
id: TerminalInlineAssistId, id: TerminalInlineAssistId,
height_in_lines: u8, height_in_lines: u8,
editor: View<Editor>, editor: View<Editor>,
context_strip: View<ContextStrip>,
language_model_selector: View<LanguageModelSelector>, language_model_selector: View<LanguageModelSelector>,
edited_since_done: bool, edited_since_done: bool,
prompt_history: VecDeque<String>, prompt_history: VecDeque<String>,
@ -452,11 +475,7 @@ impl EventEmitter<PromptEditorEvent> for PromptEditor {}
impl Render for PromptEditor { impl Render for PromptEditor {
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement { fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
let status = &self.codegen.read(cx).status; let status = &self.codegen.read(cx).status;
let mut buttons = vec![Button::new("add-context", "Add Context") let mut buttons = Vec::new();
.style(ButtonStyle::Filled)
.icon(IconName::Plus)
.icon_position(IconPosition::Start)
.into_any_element()];
buttons.extend(match status { buttons.extend(match status {
CodegenStatus::Idle => vec![ CodegenStatus::Idle => vec![
@ -554,64 +573,69 @@ impl Render for PromptEditor {
} }
}); });
h_flex() v_flex()
.bg(cx.theme().colors().editor_background)
.border_y_1() .border_y_1()
.border_color(cx.theme().status().info_border) .border_color(cx.theme().status().info_border)
.py_2() .py_2()
.h_full() .size_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))
.child( .child(
h_flex() h_flex()
.w_12() .bg(cx.theme().colors().editor_background)
.justify_center() .on_action(cx.listener(Self::confirm))
.gap_2() .on_action(cx.listener(Self::secondary_confirm))
.child(LanguageModelSelectorPopoverMenu::new( .on_action(cx.listener(Self::cancel))
self.language_model_selector.clone(), .on_action(cx.listener(Self::move_up))
IconButton::new("context", IconName::SettingsAlt) .on_action(cx.listener(Self::move_down))
.shape(IconButtonShape::Square) .child(
.icon_size(IconSize::Small) h_flex()
.icon_color(Color::Muted) .w_12()
.tooltip(move |cx| { .justify_center()
Tooltip::with_meta( .gap_2()
format!( .child(LanguageModelSelectorPopoverMenu::new(
"Using {}", self.language_model_selector.clone(),
LanguageModelRegistry::read_global(cx) IconButton::new("context", IconName::SettingsAlt)
.active_model() .shape(IconButtonShape::Square)
.map(|model| model.name().0) .icon_size(IconSize::Small)
.unwrap_or_else(|| "No model selected".into()), .icon_color(Color::Muted)
), .tooltip(move |cx| {
None, Tooltip::with_meta(
"Change Model", format!(
cx, "Using {}",
) LanguageModelRegistry::read_global(cx)
}), .active_model()
)) .map(|model| model.name().0)
.children( .unwrap_or_else(|| "No model selected".into()),
if let CodegenStatus::Error(error) = &self.codegen.read(cx).status { ),
let error_message = SharedString::from(error.to_string()); None,
Some( "Change Model",
div() cx,
.id("error") )
.tooltip(move |cx| Tooltip::text(error_message.clone(), cx)) }),
.child( ))
Icon::new(IconName::XCircle) .children(
.size(IconSize::Small) if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
.color(Color::Error), let error_message = SharedString::from(error.to_string());
), Some(
) div()
} else { .id("error")
None .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().child(self.context_strip.clone()))
.child(h_flex().gap_1().pr_4().children(buttons))
} }
} }
@ -631,6 +655,9 @@ impl PromptEditor {
prompt_buffer: Model<MultiBuffer>, prompt_buffer: Model<MultiBuffer>,
codegen: Model<Codegen>, codegen: Model<Codegen>,
fs: Arc<dyn Fs>, fs: Arc<dyn Fs>,
context_store: Model<ContextStore>,
workspace: WeakView<Workspace>,
thread_store: Option<WeakModel<ThreadStore>>,
cx: &mut ViewContext<Self>, cx: &mut ViewContext<Self>,
) -> Self { ) -> Self {
let prompt_editor = cx.new_view(|cx| { let prompt_editor = cx.new_view(|cx| {
@ -652,6 +679,9 @@ impl PromptEditor {
id, id,
height_in_lines: 1, height_in_lines: 1,
editor: prompt_editor, 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| { language_model_selector: cx.new_view(|cx| {
let fs = fs.clone(); let fs = fs.clone();
LanguageModelSelector::new( LanguageModelSelector::new(