diff --git a/assets/icons/eraser.svg b/assets/icons/eraser.svg new file mode 100644 index 0000000000..edb893a8c6 --- /dev/null +++ b/assets/icons/eraser.svg @@ -0,0 +1,4 @@ + + + + diff --git a/crates/assistant2/src/active_thread.rs b/crates/assistant2/src/active_thread.rs index d9cd8fcc46..bf0086ee11 100644 --- 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, @@ -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() } diff --git a/crates/assistant2/src/assistant.rs b/crates/assistant2/src/assistant.rs index 3c8520680e..d83a4f1818 100644 --- 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}; diff --git a/crates/assistant2/src/context.rs b/crates/assistant2/src/context.rs new file mode 100644 index 0000000000..c865139ce3 --- /dev/null +++ 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, +} diff --git a/crates/assistant2/src/message_editor.rs b/crates/assistant2/src/message_editor.rs index e0d620c2c1..d5e778f510 100644 --- 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, editor: View, + context: Vec, pub(crate) context_picker_handle: PopoverMenuHandle>, use_tools: bool, } impl MessageEditor { pub fn new(thread: Model, cx: &mut ViewContext) -> 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::>(); 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); diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs index 833f8c9b03..77c0cd9836 100644 --- 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>, messages: Vec, next_message_id: MessageId, + context_by_message: HashMap>, completion_count: usize, pending_completions: Vec, tools: Arc, @@ -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> { + 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, cx: &mut ModelContext) { - self.insert_message(Role::User, text, cx) + pub fn insert_user_message( + &mut self, + text: impl Into, + context: Vec, + cx: &mut ModelContext, + ) { + 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, cx: &mut ModelContext, - ) { + ) -> 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 diff --git a/crates/assistant2/src/thread_store.rs b/crates/assistant2/src/thread_store.rs index 94cb72ce43..cd9a1e3d6b 100644 --- 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: diff --git a/crates/assistant2/src/ui.rs b/crates/assistant2/src/ui.rs new file mode 100644 index 0000000000..b10c09b300 --- /dev/null +++ b/crates/assistant2/src/ui.rs @@ -0,0 +1,3 @@ +mod context_pill; + +pub use context_pill::*; diff --git a/crates/assistant2/src/ui/context_pill.rs b/crates/assistant2/src/ui/context_pill.rs new file mode 100644 index 0000000000..35fc6ab0e2 --- /dev/null +++ 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)) + } +} diff --git a/crates/ui/src/components/icon.rs b/crates/ui/src/components/icon.rs index ed7c294db9..1885c92736 100644 --- 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,