Sketch in assistant edit button (#19705)

Add an edit button to the assistant. This is totally hacked in for now, just to see how this would feel rendered simply in the UI. ![CleanShot 2024-10-24 at 16 26 14@2x](https://github.com/user-attachments/assets/e630d078-78b7-42d7-93f1-cf61c00bd20e) cc @as-cii @danilo-leal Release Notes: - N/A --------- Co-authored-by: Danilo Leal <67129314+danilo-leal@users.noreply.github.com> Co-authored-by: Richard Feldman <oss@rtfeldman.com>
2024-10-29 11:21:10 -06:00 · 2024-10-29 11:21:10 -06:00 · cfa20ff221
commit cfa20ff221
parent 759d136fe6
8 changed files with 155 additions and 35 deletions
--- a/assets/keymaps/default-linux.json
+++ b/assets/keymaps/default-linux.json
@ -532,6 +532,7 @@
    "context": "ContextEditor > Editor",
    "bindings": {
      "ctrl-enter": "assistant::Assist",
+      "ctrl-shift-enter": "assistant::Edit",
      "ctrl-s": "workspace::Save",
      "ctrl->": "assistant::QuoteSelection",
      "ctrl-<": "assistant::InsertIntoEditor",
--- a/assets/keymaps/default-macos.json
+++ b/assets/keymaps/default-macos.json
@ -201,6 +201,7 @@
    "context": "ContextEditor > Editor",
    "bindings": {
      "cmd-enter": "assistant::Assist",
+      "cmd-shift-enter": "assistant::Edit",
      "cmd-s": "workspace::Save",
      "cmd->": "assistant::QuoteSelection",
      "cmd-<": "assistant::InsertIntoEditor",
--- a/crates/assistant/src/assistant.rs
+++ b/crates/assistant/src/assistant.rs
@ -59,6 +59,7 @@ actions!(
    assistant,
    [
        Assist,
+        Edit,
        Split,
        CopyCode,
        CycleMessageRole,
--- a/crates/assistant/src/assistant_panel.rs
+++ b/crates/assistant/src/assistant_panel.rs
@ -13,10 +13,11 @@ use crate::{
    terminal_inline_assistant::TerminalInlineAssistant,
    Assist, AssistantPatch, AssistantPatchStatus, CacheStatus, ConfirmCommand, Content, Context,
    ContextEvent, ContextId, ContextStore, ContextStoreEvent, CopyCode, CycleMessageRole,
-    DeployHistory, DeployPromptLibrary, InlineAssistant, InsertDraggedFiles, InsertIntoEditor,
-    Message, MessageId, MessageMetadata, MessageStatus, ModelPickerDelegate, ModelSelector,
-    NewContext, PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection,
-    RemoteContextMetadata, SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector,
+    DeployHistory, DeployPromptLibrary, Edit, InlineAssistant, InsertDraggedFiles,
+    InsertIntoEditor, Message, MessageId, MessageMetadata, MessageStatus, ModelPickerDelegate,
+    ModelSelector, NewContext, PendingSlashCommand, PendingSlashCommandStatus, QuoteSelection,
+    RemoteContextMetadata, RequestType, SavedContextMetadata, Split, ToggleFocus,
+    ToggleModelSelector,
 };
 use anyhow::Result;
 use assistant_slash_command::{SlashCommand, SlashCommandOutputSection};
@ -1588,23 +1589,11 @@ impl ContextEditor {
    }

    fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
-        let provider = LanguageModelRegistry::read_global(cx).active_provider();
-        if provider
-            .as_ref()
-            .map_or(false, |provider| provider.must_accept_terms(cx))
-        {
-            self.show_accept_terms = true;
-            cx.notify();
-            return;
-        }
+        self.send_to_model(RequestType::Chat, cx);
+    }

-        if self.focus_active_patch(cx) {
-            return;
-        }
-
-        self.last_error = None;
-        self.send_to_model(cx);
-        cx.notify();
+    fn edit(&mut self, _: &Edit, cx: &mut ViewContext<Self>) {
+        self.send_to_model(RequestType::SuggestEdits, cx);
    }

    fn focus_active_patch(&mut self, cx: &mut ViewContext<Self>) -> bool {
@ -1622,8 +1611,27 @@ impl ContextEditor {
        false
    }

-    fn send_to_model(&mut self, cx: &mut ViewContext<Self>) {
-        if let Some(user_message) = self.context.update(cx, |context, cx| context.assist(cx)) {
+    fn send_to_model(&mut self, request_type: RequestType, cx: &mut ViewContext<Self>) {
+        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+        if provider
+            .as_ref()
+            .map_or(false, |provider| provider.must_accept_terms(cx))
+        {
+            self.show_accept_terms = true;
+            cx.notify();
+            return;
+        }
+
+        if self.focus_active_patch(cx) {
+            return;
+        }
+
+        self.last_error = None;
+
+        if let Some(user_message) = self
+            .context
+            .update(cx, |context, cx| context.assist(request_type, cx))
+        {
            let new_selection = {
                let cursor = user_message
                    .start
@ -1640,6 +1648,8 @@ impl ContextEditor {
            // Avoid scrolling to the new cursor position so the assistant's output is stable.
            cx.defer(|this, _| this.scroll_position = None);
        }
+
+        cx.notify();
    }

    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
@ -3644,7 +3654,13 @@ impl ContextEditor {
                button.tooltip(move |_| tooltip.clone())
            })
            .layer(ElevationIndex::ModalSurface)
-            .child(Label::new("Send"))
+            .child(Label::new(
+                if AssistantSettings::get_global(cx).are_live_diffs_enabled(cx) {
+                    "Chat"
+                } else {
+                    "Send"
+                },
+            ))
            .children(
                KeyBinding::for_action_in(&Assist, &focus_handle, cx)
                    .map(|binding| binding.into_any_element()),
@ -3654,6 +3670,57 @@ impl ContextEditor {
            })
    }

+    fn render_edit_button(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+        let focus_handle = self.focus_handle(cx).clone();
+
+        let (style, tooltip) = match token_state(&self.context, cx) {
+            Some(TokenState::NoTokensLeft { .. }) => (
+                ButtonStyle::Tinted(TintColor::Negative),
+                Some(Tooltip::text("Token limit reached", cx)),
+            ),
+            Some(TokenState::HasMoreTokens {
+                over_warn_threshold,
+                ..
+            }) => {
+                let (style, tooltip) = if over_warn_threshold {
+                    (
+                        ButtonStyle::Tinted(TintColor::Warning),
+                        Some(Tooltip::text("Token limit is close to exhaustion", cx)),
+                    )
+                } else {
+                    (ButtonStyle::Filled, None)
+                };
+                (style, tooltip)
+            }
+            None => (ButtonStyle::Filled, None),
+        };
+
+        let provider = LanguageModelRegistry::read_global(cx).active_provider();
+
+        let has_configuration_error = configuration_error(cx).is_some();
+        let needs_to_accept_terms = self.show_accept_terms
+            && provider
+                .as_ref()
+                .map_or(false, |provider| provider.must_accept_terms(cx));
+        let disabled = has_configuration_error || needs_to_accept_terms;
+
+        ButtonLike::new("edit_button")
+            .disabled(disabled)
+            .style(style)
+            .when_some(tooltip, |button, tooltip| {
+                button.tooltip(move |_| tooltip.clone())
+            })
+            .layer(ElevationIndex::ModalSurface)
+            .child(Label::new("Suggest Edits"))
+            .children(
+                KeyBinding::for_action_in(&Edit, &focus_handle, cx)
+                    .map(|binding| binding.into_any_element()),
+            )
+            .on_click(move |_event, cx| {
+                focus_handle.dispatch_action(&Edit, cx);
+            })
+    }
+
    fn render_last_error(&self, cx: &mut ViewContext<Self>) -> Option<AnyElement> {
        let last_error = self.last_error.as_ref()?;

@ -3910,6 +3977,7 @@ impl Render for ContextEditor {
            .capture_action(cx.listener(ContextEditor::paste))
            .capture_action(cx.listener(ContextEditor::cycle_message_role))
            .capture_action(cx.listener(ContextEditor::confirm_command))
+            .on_action(cx.listener(ContextEditor::edit))
            .on_action(cx.listener(ContextEditor::assist))
            .on_action(cx.listener(ContextEditor::split))
            .size_full()
@ -3974,7 +4042,21 @@ impl Render for ContextEditor {
                            h_flex()
                                .w_full()
                                .justify_end()
-                                .child(div().child(self.render_send_button(cx))),
+                                .when(
+                                    AssistantSettings::get_global(cx).are_live_diffs_enabled(cx),
+                                    |buttons| {
+                                        buttons
+                                            .items_center()
+                                            .gap_1p5()
+                                            .child(self.render_edit_button(cx))
+                                            .child(
+                                                Label::new("or")
+                                                    .size(LabelSize::Small)
+                                                    .color(Color::Muted),
+                                            )
+                                    },
+                                )
+                                .child(self.render_send_button(cx)),
                        ),
                ),
            )
--- a/crates/assistant/src/context.rs
+++ b/crates/assistant/src/context.rs
@ -66,6 +66,14 @@ impl ContextId {
    }
 }

+#[derive(Clone, Copy, Debug)]
+pub enum RequestType {
+    /// Request a normal chat response from the model.
+    Chat,
+    /// Add a preamble to the message, which tells the model to return a structured response that suggests edits.
+    SuggestEdits,
+}
+
 #[derive(Clone, Debug)]
 pub enum ContextOperation {
    InsertMessage {
@ -1028,7 +1036,7 @@ impl Context {
    }

    pub(crate) fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
-        let request = self.to_completion_request(cx);
+        let request = self.to_completion_request(RequestType::SuggestEdits, cx); // Conservatively assume SuggestEdits, since it takes more tokens.
        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
            return;
        };
@ -1171,7 +1179,7 @@ impl Context {
        }

        let request = {
-            let mut req = self.to_completion_request(cx);
+            let mut req = self.to_completion_request(RequestType::Chat, cx);
            // Skip the last message because it's likely to change and
            // therefore would be a waste to cache.
            req.messages.pop();
@ -1859,7 +1867,11 @@ impl Context {
        })
    }

-    pub fn assist(&mut self, cx: &mut ModelContext<Self>) -> Option<MessageAnchor> {
+    pub fn assist(
+        &mut self,
+        request_type: RequestType,
+        cx: &mut ModelContext<Self>,
+    ) -> Option<MessageAnchor> {
        let model_registry = LanguageModelRegistry::read_global(cx);
        let provider = model_registry.active_provider()?;
        let model = model_registry.active_model()?;
@ -1872,7 +1884,7 @@ impl Context {
        // Compute which messages to cache, including the last one.
        self.mark_cache_anchors(&model.cache_configuration(), false, cx);

-        let mut request = self.to_completion_request(cx);
+        let mut request = self.to_completion_request(request_type, cx);

        if cx.has_flag::<ToolUseFeatureFlag>() {
            let tool_registry = ToolRegistry::global(cx);
@ -2074,7 +2086,11 @@ impl Context {
        Some(user_message)
    }

-    pub fn to_completion_request(&self, cx: &AppContext) -> LanguageModelRequest {
+    pub fn to_completion_request(
+        &self,
+        request_type: RequestType,
+        cx: &AppContext,
+    ) -> LanguageModelRequest {
        let buffer = self.buffer.read(cx);

        let mut contents = self.contents(cx).peekable();
@ -2163,6 +2179,25 @@ impl Context {
            completion_request.messages.push(request_message);
        }

+        if let RequestType::SuggestEdits = request_type {
+            if let Ok(preamble) = self.prompt_builder.generate_workflow_prompt() {
+                let last_elem_index = completion_request.messages.len();
+
+                completion_request
+                    .messages
+                    .push(LanguageModelRequestMessage {
+                        role: Role::User,
+                        content: vec![MessageContent::Text(preamble)],
+                        cache: false,
+                    });
+
+                // The preamble message should be sent right before the last actual user message.
+                completion_request
+                    .messages
+                    .swap(last_elem_index, last_elem_index.saturating_sub(1));
+            }
+        }
+
        completion_request
    }

@ -2477,7 +2512,7 @@ impl Context {
                return;
            }

-            let mut request = self.to_completion_request(cx);
+            let mut request = self.to_completion_request(RequestType::Chat, cx);
            request.messages.push(LanguageModelRequestMessage {
                role: Role::User,
                content: vec![
--- a/crates/assistant/src/inline_assistant.rs
+++ b/crates/assistant/src/inline_assistant.rs
@ -1,7 +1,7 @@
 use crate::{
    assistant_settings::AssistantSettings, humanize_token_count, prompts::PromptBuilder,
    AssistantPanel, AssistantPanelEvent, CharOperation, CycleNextInlineAssist,
-    CyclePreviousInlineAssist, LineDiff, LineOperation, ModelSelector, StreamingDiff,
+    CyclePreviousInlineAssist, LineDiff, LineOperation, ModelSelector, RequestType, StreamingDiff,
 };
 use anyhow::{anyhow, Context as _, Result};
 use client::{telemetry::Telemetry, ErrorExt};
@ -2234,7 +2234,7 @@ impl InlineAssist {
                    .read(cx)
                    .active_context(cx)?
                    .read(cx)
-                    .to_completion_request(cx),
+                    .to_completion_request(RequestType::Chat, cx),
            )
        } else {
            None
--- a/crates/assistant/src/terminal_inline_assistant.rs
+++ b/crates/assistant/src/terminal_inline_assistant.rs
@ -1,6 +1,6 @@
 use crate::{
    humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent,
-    ModelSelector, DEFAULT_CONTEXT_LINES,
+    ModelSelector, RequestType, DEFAULT_CONTEXT_LINES,
 };
 use anyhow::{Context as _, Result};
 use client::telemetry::Telemetry;
@ -251,7 +251,7 @@ impl TerminalInlineAssistant {
                        .read(cx)
                        .active_context(cx)?
                        .read(cx)
-                        .to_completion_request(cx),
+                        .to_completion_request(RequestType::Chat, cx),
                )
            })
        } else {
--- a/crates/ui/src/components/keybinding.rs
+++ b/crates/ui/src/components/keybinding.rs
@ -184,7 +184,7 @@ pub struct KeyIcon {
 impl RenderOnce for KeyIcon {
    fn render(self, _cx: &mut WindowContext) -> impl IntoElement {
        Icon::new(self.icon)
-            .size(IconSize::Small)
+            .size(IconSize::XSmall)
            .color(Color::Muted)
    }
 }