assistant2: Sketch in sending file context to model (#21829)

This PR sketches in support for sending file context attached to a message to the model. Right now the context is just mocked. <img width="1159" alt="Screenshot 2024-12-10 at 4 18 41 PM" src="https://github.com/user-attachments/assets/3ee4e86a-7893-42dc-98f9-982aa202d310"> <img width="1159" alt="Screenshot 2024-12-10 at 4 18 53 PM" src="https://github.com/user-attachments/assets/8a3c2dd7-a466-4dbf-83ec-1c7d969c1a4b"> Release Notes: - N/A
2024-12-10 16:35:53 -05:00 · 2024-12-10 16:35:53 -05:00 · c255e55599
commit c255e55599
parent f80eb73213
10 changed files with 145 additions and 15 deletions
--- a/assets/icons/eraser.svg
+++ b/assets/icons/eraser.svg
@ -0,0 +1,4 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-eraser">
+    <path d="m7 21-4.3-4.3c-1-1-1-2.5 0-3.4l9.6-9.6c1-1 2.5-1 3.4 0l5.6 5.6c1 1 1 2.5 0 3.4L13 21"/>
+    <path d="M22 21H7"/><path d="m5 11 9 9"/>
+</svg>
--- a/crates/assistant2/src/active_thread.rs
+++ b/crates/assistant2/src/active_thread.rs
@ -15,6 +15,7 @@ use ui::prelude::*;
 use workspace::Workspace;

 use crate::thread::{MessageId, Thread, ThreadError, ThreadEvent};
+use crate::ui::ContextPill;

 pub struct ActiveThread {
    workspace: WeakView<Workspace>,
@ -202,6 +203,8 @@ impl ActiveThread {
            return Empty.into_any();
        };

+        let context = self.thread.read(cx).context_for_message(message_id);
+
        let (role_icon, role_name) = match message.role {
            Role::User => (IconName::Person, "You"),
            Role::Assistant => (IconName::ZedAssistant, "Assistant"),
@ -229,7 +232,16 @@ impl ActiveThread {
                                    .child(Label::new(role_name).size(LabelSize::Small)),
                            ),
                    )
-                    .child(v_flex().p_1p5().text_ui(cx).child(markdown.clone())),
+                    .child(v_flex().p_1p5().text_ui(cx).child(markdown.clone()))
+                    .when_some(context, |parent, context| {
+                        parent.child(
+                            h_flex().flex_wrap().gap_2().p_1p5().children(
+                                context
+                                    .iter()
+                                    .map(|context| ContextPill::new(context.clone())),
+                            ),
+                        )
+                    }),
            )
            .into_any()
    }
--- a/crates/assistant2/src/assistant.rs
+++ b/crates/assistant2/src/assistant.rs
@ -1,10 +1,12 @@
 mod active_thread;
 mod assistant_panel;
+mod context;
 mod context_picker;
 mod message_editor;
 mod thread;
 mod thread_history;
 mod thread_store;
+mod ui;

 use command_palette_hooks::CommandPaletteFilter;
 use feature_flags::{Assistant2FeatureFlag, FeatureFlagAppExt};
--- a/crates/assistant2/src/context.rs
+++ b/crates/assistant2/src/context.rs
@ -0,0 +1,14 @@
+use gpui::SharedString;
+
+/// Some context attached to a message in a thread.
+#[derive(Debug, Clone)]
+pub struct Context {
+    pub name: SharedString,
+    pub kind: ContextKind,
+    pub text: SharedString,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum ContextKind {
+    File,
+}
--- a/crates/assistant2/src/message_editor.rs
+++ b/crates/assistant2/src/message_editor.rs
@ -10,19 +10,28 @@ use ui::{
    PopoverMenuHandle, Tooltip,
 };

+use crate::context::{Context, ContextKind};
 use crate::context_picker::{ContextPicker, ContextPickerDelegate};
 use crate::thread::{RequestKind, Thread};
+use crate::ui::ContextPill;
 use crate::{Chat, ToggleModelSelector};

 pub struct MessageEditor {
    thread: Model<Thread>,
    editor: View<Editor>,
+    context: Vec<Context>,
    pub(crate) context_picker_handle: PopoverMenuHandle<Picker<ContextPickerDelegate>>,
    use_tools: bool,
 }

 impl MessageEditor {
    pub fn new(thread: Model<Thread>, cx: &mut ViewContext<Self>) -> Self {
+        let mocked_context = vec![Context {
+            name: "shape.rs".into(),
+            kind: ContextKind::File,
+            text: "```rs\npub enum Shape {\n    Circle,\n    Square,\n    Triangle,\n}".into(),
+        }];
+
        Self {
            thread,
            editor: cx.new_view(|cx| {
@ -31,6 +40,7 @@ impl MessageEditor {

                editor
            }),
+            context: mocked_context,
            context_picker_handle: PopoverMenuHandle::default(),
            use_tools: false,
        }
@ -62,9 +72,10 @@ impl MessageEditor {
            editor.clear(cx);
            text
        });
+        let context = self.context.drain(..).collect::<Vec<_>>();

        self.thread.update(cx, |thread, cx| {
-            thread.insert_user_message(user_message, cx);
+            thread.insert_user_message(user_message, context, cx);
            let mut request = thread.to_completion_request(request_kind, cx);

            if self.use_tools {
@ -158,12 +169,32 @@ impl Render for MessageEditor {
            .p_2()
            .bg(cx.theme().colors().editor_background)
            .child(
-                h_flex().gap_2().child(ContextPicker::new(
-                    cx.view().downgrade(),
-                    IconButton::new("add-context", IconName::Plus)
-                        .shape(IconButtonShape::Square)
-                        .icon_size(IconSize::Small),
-                )),
+                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),
+                    ))
+                    .children(
+                        self.context
+                            .iter()
+                            .map(|context| ContextPill::new(context.clone())),
+                    )
+                    .when(!self.context.is_empty(), |parent| {
+                        parent.child(
+                            IconButton::new("remove-all-context", IconName::Eraser)
+                                .shape(IconButtonShape::Square)
+                                .icon_size(IconSize::Small)
+                                .tooltip(move |cx| Tooltip::text("Remove All Context", cx))
+                                .on_click(cx.listener(|this, _event, cx| {
+                                    this.context.clear();
+                                    cx.notify();
+                                })),
+                        )
+                    }),
            )
            .child({
                let settings = ThemeSettings::get_global(cx);
--- a/crates/assistant2/src/thread.rs
+++ b/crates/assistant2/src/thread.rs
@ -17,6 +17,8 @@ use serde::{Deserialize, Serialize};
 use util::{post_inc, TryFutureExt as _};
 use uuid::Uuid;

+use crate::context::{Context, ContextKind};
+
 #[derive(Debug, Clone, Copy)]
 pub enum RequestKind {
    Chat,
@ -62,6 +64,7 @@ pub struct Thread {
    pending_summary: Task<Option<()>>,
    messages: Vec<Message>,
    next_message_id: MessageId,
+    context_by_message: HashMap<MessageId, Vec<Context>>,
    completion_count: usize,
    pending_completions: Vec<PendingCompletion>,
    tools: Arc<ToolWorkingSet>,
@ -79,6 +82,7 @@ impl Thread {
            pending_summary: Task::ready(None),
            messages: Vec::new(),
            next_message_id: MessageId(0),
+            context_by_message: HashMap::default(),
            completion_count: 0,
            pending_completions: Vec::new(),
            tools,
@ -125,12 +129,22 @@ impl Thread {
        &self.tools
    }

+    pub fn context_for_message(&self, id: MessageId) -> Option<&Vec<Context>> {
+        self.context_by_message.get(&id)
+    }
+
    pub fn pending_tool_uses(&self) -> Vec<&PendingToolUse> {
        self.pending_tool_uses_by_id.values().collect()
    }

-    pub fn insert_user_message(&mut self, text: impl Into<String>, cx: &mut ModelContext<Self>) {
-        self.insert_message(Role::User, text, cx)
+    pub fn insert_user_message(
+        &mut self,
+        text: impl Into<String>,
+        context: Vec<Context>,
+        cx: &mut ModelContext<Self>,
+    ) {
+        let message_id = self.insert_message(Role::User, text, cx);
+        self.context_by_message.insert(message_id, context);
    }

    pub fn insert_message(
@ -138,7 +152,7 @@ impl Thread {
        role: Role,
        text: impl Into<String>,
        cx: &mut ModelContext<Self>,
-    ) {
+    ) -> MessageId {
        let id = self.next_message_id.post_inc();
        self.messages.push(Message {
            id,
@ -147,6 +161,7 @@ impl Thread {
        });
        self.touch_updated_at();
        cx.emit(ThreadEvent::MessageAdded(id));
+        id
    }

    pub fn to_completion_request(
@ -176,6 +191,29 @@ impl Thread {
                }
            }

+            if let Some(context) = self.context_for_message(message.id) {
+                let mut file_context = String::new();
+
+                for context in context.iter() {
+                    match context.kind {
+                        ContextKind::File => {
+                            file_context.push_str(&context.text);
+                            file_context.push_str("\n");
+                        }
+                    }
+                }
+
+                let mut context_text = String::new();
+                if !file_context.is_empty() {
+                    context_text.push_str("The following files are available:\n");
+                    context_text.push_str(&file_context);
+                }
+
+                request_message
+                    .content
+                    .push(MessageContent::Text(context_text))
+            }
+
            if !message.text.is_empty() {
                request_message
                    .content
--- a/crates/assistant2/src/thread_store.rs
+++ b/crates/assistant2/src/thread_store.rs
@ -159,9 +159,9 @@ impl ThreadStore {
        self.threads.push(cx.new_model(|cx| {
            let mut thread = Thread::new(self.tools.clone(), cx);
            thread.set_summary("Introduction to quantum computing", cx);
-            thread.insert_user_message("Hello! Can you help me understand quantum computing?", cx);
+            thread.insert_user_message("Hello! Can you help me understand quantum computing?", Vec::new(), cx);
            thread.insert_message(Role::Assistant, "Of course! I'd be happy to help you understand quantum computing. Quantum computing is a fascinating field that uses the principles of quantum mechanics to process information. Unlike classical computers that use bits (0s and 1s), quantum computers use quantum bits or 'qubits'. These qubits can exist in multiple states simultaneously, a property called superposition. This allows quantum computers to perform certain calculations much faster than classical computers. What specific aspect of quantum computing would you like to know more about?", cx);
-            thread.insert_user_message("That's interesting! Can you explain how quantum entanglement is used in quantum computing?", cx);
+            thread.insert_user_message("That's interesting! Can you explain how quantum entanglement is used in quantum computing?", Vec::new(), cx);
            thread.insert_message(Role::Assistant, "Certainly! Quantum entanglement is a key principle used in quantum computing. When two qubits become entangled, the state of one qubit is directly related to the state of the other, regardless of the distance between them. This property is used in quantum computing to create complex quantum states and to perform operations on multiple qubits simultaneously. Entanglement allows quantum computers to process information in ways that classical computers cannot, potentially solving certain problems much more efficiently. For example, it's crucial in quantum error correction and in algorithms like quantum teleportation, which is important for quantum communication.", cx);
            thread
        }));
@ -169,7 +169,7 @@ impl ThreadStore {
        self.threads.push(cx.new_model(|cx| {
            let mut thread = Thread::new(self.tools.clone(), cx);
            thread.set_summary("Rust web development and async programming", cx);
-            thread.insert_user_message("Can you show me an example of Rust code for a simple web server?", cx);
+            thread.insert_user_message("Can you show me an example of Rust code for a simple web server?", Vec::new(), cx);
            thread.insert_message(Role::Assistant, "Certainly! Here's an example of a simple web server in Rust using the `actix-web` framework:

        ```rust
@ -206,7 +206,7 @@ impl ThreadStore {
        ```

        Then you can run the server with `cargo run` and access it at `http://localhost:8080`.".unindent(), cx);
-            thread.insert_user_message("That's great! Can you explain more about async functions in Rust?", cx);
+            thread.insert_user_message("That's great! Can you explain more about async functions in Rust?", Vec::new(), cx);
            thread.insert_message(Role::Assistant, "Certainly! Async functions are a key feature in Rust for writing efficient, non-blocking code, especially for I/O-bound operations. Here's an overview:

        1. **Syntax**: Async functions are declared using the `async` keyword:
--- a/crates/assistant2/src/ui.rs
+++ b/crates/assistant2/src/ui.rs
@ -0,0 +1,3 @@
+mod context_pill;
+
+pub use context_pill::*;
--- a/crates/assistant2/src/ui/context_pill.rs
+++ b/crates/assistant2/src/ui/context_pill.rs
@ -0,0 +1,25 @@
+use ui::prelude::*;
+
+use crate::context::Context;
+
+#[derive(IntoElement)]
+pub struct ContextPill {
+    context: Context,
+}
+
+impl ContextPill {
+    pub fn new(context: Context) -> Self {
+        Self { context }
+    }
+}
+
+impl RenderOnce for ContextPill {
+    fn render(self, cx: &mut WindowContext) -> impl IntoElement {
+        div()
+            .px_1()
+            .border_1()
+            .border_color(cx.theme().colors().border)
+            .rounded_md()
+            .child(Label::new(self.context.name.clone()).size(LabelSize::Small))
+    }
+}
--- a/crates/ui/src/components/icon.rs
+++ b/crates/ui/src/components/icon.rs
@ -169,6 +169,7 @@ pub enum IconName {
    Ellipsis,
    EllipsisVertical,
    Envelope,
+    Eraser,
    Escape,
    ExpandVertical,
    Exit,