Skip to content

Commit

Permalink
assistant2: Wire up context for terminal inline assist (zed-industrie…
Browse files Browse the repository at this point in the history
…s#22108)

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

Release Notes:

- N/A

Co-authored-by: Richard <[email protected]>
Co-authored-by: Agus <[email protected]>
  • Loading branch information
3 people authored Dec 16, 2024
1 parent 84392fb commit 92fb38a
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 68 deletions.
10 changes: 5 additions & 5 deletions crates/assistant2/src/inline_assistant.rs
Original file line number Diff line number Diff line change
Expand Up @@ -221,19 +221,19 @@ impl InlineAssistant {
.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 {
InlineAssistTarget::Editor(active_editor) => {
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)
})
}
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)
})
}
};
Expand Down
156 changes: 93 additions & 63 deletions crates/assistant2/src/terminal_inline_assistant.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand All @@ -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::{
Expand Down Expand Up @@ -83,13 +87,15 @@ impl TerminalInlineAssistant {
&mut self,
terminal_view: &View<TerminalView>,
workspace: WeakView<Workspace>,
thread_store: Option<WeakModel<ThreadStore>>,
cx: &mut WindowContext,
) {
let terminal = terminal_view.read(cx).terminal().clone();
let assist_id = self.next_assist_id.post_inc();
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| {
Expand All @@ -99,6 +105,9 @@ impl TerminalInlineAssistant {
prompt_buffer.clone(),
codegen,
self.fs.clone(),
context_store.clone(),
workspace.clone(),
thread_store.clone(),
cx,
)
});
Expand All @@ -116,6 +125,7 @@ impl TerminalInlineAssistant {
terminal_view,
prompt_editor,
workspace.clone(),
context_store,
cx,
);

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -362,6 +381,7 @@ struct TerminalInlineAssist {
prompt_editor: Option<View<PromptEditor>>,
codegen: Model<Codegen>,
workspace: WeakView<Workspace>,
context_store: Model<ContextStore>,
_subscriptions: Vec<Subscription>,
}

Expand All @@ -371,6 +391,7 @@ impl TerminalInlineAssist {
terminal: &View<TerminalView>,
prompt_editor: View<PromptEditor>,
workspace: WeakView<Workspace>,
context_store: Model<ContextStore>,
cx: &mut WindowContext,
) -> Self {
let codegen = prompt_editor.read(cx).codegen.clone();
Expand All @@ -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| {
Expand Down Expand Up @@ -437,6 +459,7 @@ struct PromptEditor {
id: TerminalInlineAssistId,
height_in_lines: u8,
editor: View<Editor>,
context_strip: View<ContextStrip>,
language_model_selector: View<LanguageModelSelector>,
edited_since_done: bool,
prompt_history: VecDeque<String>,
Expand All @@ -452,11 +475,7 @@ impl EventEmitter<PromptEditorEvent> for PromptEditor {}
impl Render for PromptEditor {
fn render(&mut self, cx: &mut ViewContext<Self>) -> 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![
Expand Down Expand Up @@ -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()))
}
}

Expand All @@ -631,6 +655,9 @@ impl PromptEditor {
prompt_buffer: Model<MultiBuffer>,
codegen: Model<Codegen>,
fs: Arc<dyn Fs>,
context_store: Model<ContextStore>,
workspace: WeakView<Workspace>,
thread_store: Option<WeakModel<ThreadStore>>,
cx: &mut ViewContext<Self>,
) -> Self {
let prompt_editor = cx.new_view(|cx| {
Expand All @@ -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(
Expand Down

0 comments on commit 92fb38a

Please sign in to comment.