From 3e5dcd1bec4f22023c81ee6048c3305cfb61fcda Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Fri, 3 May 2024 14:48:00 -0700 Subject: [PATCH] Attachment store for assistant2 (#11327) This sets up a way for the user (or Zed) to _push_ context instead of having the model retrieve it with a function. Our first use is the contents of the current file. image image I heard the asst2 example was deleted in another branch so I deleted that here too since we wanted the workspace access. Release Notes: - N/A --------- Co-authored-by: Marshall --- Cargo.lock | 1 - crates/assistant2/Cargo.toml | 1 - .../examples/chat_with_functions.rs | 378 ------------------ crates/assistant2/src/assistant2.rs | 136 +++++-- crates/assistant2/src/attachments.rs | 240 +++++++++++ crates/assistant2/src/ui.rs | 2 + .../assistant2/src/ui/active_file_button.rs | 133 ++++++ crates/assistant2/src/ui/chat_message.rs | 8 + crates/assistant2/src/ui/composer.rs | 28 +- .../assistant2/src/ui/stories/chat_message.rs | 7 + 10 files changed, 525 insertions(+), 409 deletions(-) delete mode 100644 crates/assistant2/examples/chat_with_functions.rs create mode 100644 crates/assistant2/src/attachments.rs create mode 100644 crates/assistant2/src/ui/active_file_button.rs diff --git a/Cargo.lock b/Cargo.lock index 894b827c4f..9034263ab8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,7 +382,6 @@ dependencies = [ "editor", "env_logger", "feature_flags", - "fs", "futures 0.3.28", "gpui", "language", diff --git a/crates/assistant2/Cargo.toml b/crates/assistant2/Cargo.toml index 0e5b090cf5..e006db6749 100644 --- a/crates/assistant2/Cargo.toml +++ b/crates/assistant2/Cargo.toml @@ -22,7 +22,6 @@ client.workspace = true collections.workspace = true editor.workspace = true feature_flags.workspace = true -fs.workspace = true futures.workspace = true gpui.workspace = true language.workspace = true diff --git a/crates/assistant2/examples/chat_with_functions.rs b/crates/assistant2/examples/chat_with_functions.rs deleted file mode 100644 index 1b8afa1973..0000000000 --- a/crates/assistant2/examples/chat_with_functions.rs +++ /dev/null @@ -1,378 +0,0 @@ -//! This example creates a basic Chat UI with a function for rolling a die. - -use anyhow::{Context as _, Result}; -use assets::Assets; -use assistant2::AssistantPanel; -use assistant_tooling::{LanguageModelTool, ToolRegistry}; -use client::{Client, UserStore}; -use fs::Fs; -use futures::StreamExt as _; -use gpui::{actions, AnyElement, App, AppContext, KeyBinding, Model, Task, View, WindowOptions}; -use language::LanguageRegistry; -use project::Project; -use rand::Rng; -use schemars::JsonSchema; -use serde::{Deserialize, Serialize}; -use settings::{KeymapFile, DEFAULT_KEYMAP_PATH}; -use std::{path::PathBuf, sync::Arc}; -use theme::LoadThemes; -use ui::{div, prelude::*, Render}; -use util::ResultExt as _; - -actions!(example, [Quit]); - -struct RollDiceTool {} - -impl RollDiceTool { - fn new() -> Self { - Self {} - } -} - -#[derive(Serialize, Deserialize, JsonSchema, Clone)] -#[serde(rename_all = "snake_case")] -enum Die { - D6 = 6, - D20 = 20, -} - -impl Die { - fn into_str(&self) -> &'static str { - match self { - Die::D6 => "d6", - Die::D20 => "d20", - } - } -} - -#[derive(Serialize, Deserialize, JsonSchema, Clone)] -struct DiceParams { - /// The number of dice to roll. - num_dice: u8, - /// Which die to roll. Defaults to a d6 if not provided. - die_type: Option, -} - -#[derive(Serialize, Deserialize)] -struct DieRoll { - die: Die, - roll: u8, -} - -impl DieRoll { - fn render(&self) -> AnyElement { - match self.die { - Die::D6 => { - let face = match self.roll { - 6 => div().child("⚅"), - 5 => div().child("⚄"), - 4 => div().child("⚃"), - 3 => div().child("⚂"), - 2 => div().child("⚁"), - 1 => div().child("⚀"), - _ => div().child("😅"), - }; - face.text_3xl().into_any_element() - } - _ => div() - .child(format!("{}", self.roll)) - .text_3xl() - .into_any_element(), - } - } -} - -#[derive(Serialize, Deserialize)] -struct DiceRoll { - rolls: Vec, -} - -pub struct DiceView { - result: Result, -} - -impl Render for DiceView { - fn render(&mut self, _cx: &mut ViewContext) -> impl IntoElement { - let output = match &self.result { - Ok(output) => output, - Err(_) => return "Somehow dice failed 🎲".into_any_element(), - }; - - h_flex() - .children( - output - .rolls - .iter() - .map(|roll| div().p_2().child(roll.render())), - ) - .into_any_element() - } -} - -impl LanguageModelTool for RollDiceTool { - type Input = DiceParams; - type Output = DiceRoll; - type View = DiceView; - - fn name(&self) -> String { - "roll_dice".to_string() - } - - fn description(&self) -> String { - "Rolls N many dice and returns the results.".to_string() - } - - fn execute( - &self, - input: &Self::Input, - _cx: &mut WindowContext, - ) -> Task> { - let rolls = (0..input.num_dice) - .map(|_| { - let die_type = input.die_type.as_ref().unwrap_or(&Die::D6).clone(); - - DieRoll { - die: die_type.clone(), - roll: rand::thread_rng().gen_range(1..=die_type as u8), - } - }) - .collect(); - - return Task::ready(Ok(DiceRoll { rolls })); - } - - fn output_view( - _tool_call_id: String, - _input: Self::Input, - result: Result, - cx: &mut WindowContext, - ) -> gpui::View { - cx.new_view(|_cx| DiceView { result }) - } - - fn format(_: &Self::Input, output: &Result) -> String { - let output = match output { - Ok(output) => output, - Err(_) => return "Somehow dice failed 🎲".to_string(), - }; - - let mut result = String::new(); - for roll in &output.rolls { - let die = &roll.die; - result.push_str(&format!("{}: {}\n", die.into_str(), roll.roll)); - } - result - } -} - -struct FileBrowserTool { - fs: Arc, - root_dir: PathBuf, -} - -impl FileBrowserTool { - fn new(fs: Arc, root_dir: PathBuf) -> Self { - Self { fs, root_dir } - } -} - -#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] -struct FileBrowserParams { - command: FileBrowserCommand, -} - -#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] -enum FileBrowserCommand { - Ls { path: PathBuf }, - Cat { path: PathBuf }, -} - -#[derive(Serialize, Deserialize)] -enum FileBrowserOutput { - Ls { entries: Vec }, - Cat { content: String }, -} - -pub struct FileBrowserView { - result: Result, -} - -impl Render for FileBrowserView { - fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { - let Ok(output) = self.result.as_ref() else { - return h_flex().child("Failed to perform operation"); - }; - - match output { - FileBrowserOutput::Ls { entries } => v_flex().children( - entries - .into_iter() - .map(|entry| h_flex().text_ui(cx).child(entry.clone())), - ), - FileBrowserOutput::Cat { content } => h_flex().child(content.clone()), - } - } -} - -impl LanguageModelTool for FileBrowserTool { - type Input = FileBrowserParams; - type Output = FileBrowserOutput; - type View = FileBrowserView; - - fn name(&self) -> String { - "file_browser".to_string() - } - - fn description(&self) -> String { - "A tool for browsing the filesystem.".to_string() - } - - fn execute( - &self, - input: &Self::Input, - cx: &mut WindowContext, - ) -> Task> { - cx.spawn({ - let fs = self.fs.clone(); - let root_dir = self.root_dir.clone(); - let input = input.clone(); - |_cx| async move { - match input.command { - FileBrowserCommand::Ls { path } => { - let path = root_dir.join(path); - - let mut output = fs.read_dir(&path).await?; - - let mut entries = Vec::new(); - while let Some(entry) = output.next().await { - let entry = entry?; - entries.push(entry.display().to_string()); - } - - Ok(FileBrowserOutput::Ls { entries }) - } - FileBrowserCommand::Cat { path } => { - let path = root_dir.join(path); - - let output = fs.load(&path).await?; - - Ok(FileBrowserOutput::Cat { content: output }) - } - } - } - }) - } - - fn output_view( - _tool_call_id: String, - _input: Self::Input, - result: Result, - cx: &mut WindowContext, - ) -> gpui::View { - cx.new_view(|_cx| FileBrowserView { result }) - } - - fn format(_input: &Self::Input, output: &Result) -> String { - let Ok(output) = output else { - return "Failed to perform command: {input:?}".to_string(); - }; - - match output { - FileBrowserOutput::Ls { entries } => entries.join("\n"), - FileBrowserOutput::Cat { content } => content.to_owned(), - } - } -} - -fn main() { - env_logger::init(); - App::new().with_assets(Assets).run(|cx| { - cx.bind_keys(Some(KeyBinding::new("cmd-q", Quit, None))); - cx.on_action(|_: &Quit, cx: &mut AppContext| { - cx.quit(); - }); - - settings::init(cx); - language::init(cx); - Project::init_settings(cx); - editor::init(cx); - theme::init(LoadThemes::JustBase, cx); - Assets.load_fonts(cx).unwrap(); - KeymapFile::load_asset(DEFAULT_KEYMAP_PATH, cx).unwrap(); - client::init_settings(cx); - release_channel::init("0.130.0", cx); - - let client = Client::production(cx); - { - let client = client.clone(); - cx.spawn(|cx| async move { client.authenticate_and_connect(false, &cx).await }) - .detach_and_log_err(cx); - } - assistant2::init(client.clone(), cx); - - let language_registry = Arc::new(LanguageRegistry::new( - Task::ready(()), - cx.background_executor().clone(), - )); - - let user_store = cx.new_model(|cx| UserStore::new(client.clone(), cx)); - let node_runtime = node_runtime::RealNodeRuntime::new(client.http_client()); - languages::init(language_registry.clone(), node_runtime, cx); - - cx.spawn(|cx| async move { - cx.update(|cx| { - let fs = Arc::new(fs::RealFs::new(None)); - let cwd = std::env::current_dir().expect("Failed to get current working directory"); - - cx.open_window(WindowOptions::default(), |cx| { - let mut tool_registry = ToolRegistry::new(); - tool_registry - .register(RollDiceTool::new(), cx) - .context("failed to register DummyTool") - .log_err(); - - tool_registry - .register(FileBrowserTool::new(fs, cwd), cx) - .context("failed to register FileBrowserTool") - .log_err(); - - let tool_registry = Arc::new(tool_registry); - - println!("Tools registered"); - for definition in tool_registry.definitions() { - println!("{}", definition); - } - - cx.new_view(|cx| Example::new(language_registry, tool_registry, user_store, cx)) - }); - cx.activate(true); - }) - }) - .detach_and_log_err(cx); - }) -} - -struct Example { - assistant_panel: View, -} - -impl Example { - fn new( - language_registry: Arc, - tool_registry: Arc, - user_store: Model, - cx: &mut ViewContext, - ) -> Self { - Self { - assistant_panel: cx.new_view(|cx| { - AssistantPanel::new(language_registry, tool_registry, user_store, None, cx) - }), - } - } -} - -impl Render for Example { - fn render(&mut self, _cx: &mut ViewContext) -> impl ui::prelude::IntoElement { - div().size_full().child(self.assistant_panel.clone()) - } -} diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs index ae5d9d6feb..dc777e42b0 100644 --- a/crates/assistant2/src/assistant2.rs +++ b/crates/assistant2/src/assistant2.rs @@ -1,4 +1,5 @@ mod assistant_settings; +mod attachments; mod completion_provider; mod tools; pub mod ui; @@ -6,6 +7,7 @@ pub mod ui; use ::ui::{div, prelude::*, Color, ViewContext}; use anyhow::{Context, Result}; use assistant_tooling::{ToolFunctionCall, ToolRegistry}; +use attachments::{ActiveEditorAttachmentTool, UserAttachment, UserAttachmentStore}; use client::{proto, Client, UserStore}; use collections::HashMap; use completion_provider::*; @@ -23,8 +25,8 @@ use semantic_index::{CloudEmbeddingProvider, ProjectIndex, SemanticIndex}; use serde::Deserialize; use settings::Settings; use std::sync::Arc; -use ui::{Composer, ProjectIndexButton}; -use util::{paths::EMBEDDINGS_DIR, ResultExt}; +use ui::{ActiveFileButton, Composer, ProjectIndexButton}; +use util::{maybe, paths::EMBEDDINGS_DIR, ResultExt}; use workspace::{ dock::{DockPosition, Panel, PanelEvent}, Workspace, @@ -129,13 +131,16 @@ impl AssistantPanel { .context("failed to register CreateBufferTool") .log_err(); - let tool_registry = Arc::new(tool_registry); + let mut attachment_store = UserAttachmentStore::new(); + attachment_store.register(ActiveEditorAttachmentTool::new(workspace.clone(), cx)); Self::new( app_state.languages.clone(), - tool_registry, + Arc::new(attachment_store), + Arc::new(tool_registry), user_store, Some(project_index), + workspace, cx, ) }) @@ -144,17 +149,21 @@ impl AssistantPanel { pub fn new( language_registry: Arc, + attachment_store: Arc, tool_registry: Arc, user_store: Model, project_index: Option>, + workspace: WeakView, cx: &mut ViewContext, ) -> Self { let chat = cx.new_view(|cx| { AssistantChat::new( language_registry.clone(), + attachment_store.clone(), tool_registry.clone(), user_store, project_index, + workspace, cx, ) }); @@ -229,11 +238,13 @@ pub struct AssistantChat { language_registry: Arc, composer_editor: View, project_index_button: Option>, + active_file_button: Option>, user_store: Model, next_message_id: MessageId, collapsed_messages: HashMap, editing_message: Option, pending_completion: Option>, + attachment_store: Arc, tool_registry: Arc, project_index: Option>, } @@ -247,9 +258,11 @@ struct EditingMessage { impl AssistantChat { fn new( language_registry: Arc, + attachment_store: Arc, tool_registry: Arc, user_store: Model, project_index: Option>, + workspace: WeakView, cx: &mut ViewContext, ) -> Self { let model = CompletionProvider::get(cx).default_model(); @@ -268,6 +281,15 @@ impl AssistantChat { cx.new_view(|cx| ProjectIndexButton::new(project_index, tool_registry.clone(), cx)) }); + let active_file_button = match workspace.upgrade() { + Some(workspace) => { + Some(cx.new_view( + |cx| ActiveFileButton::new(attachment_store.clone(), workspace, cx), // + )) + } + _ => None, + }; + Self { model, messages: Vec::new(), @@ -281,11 +303,13 @@ impl AssistantChat { user_store, language_registry, project_index_button, + active_file_button, project_index, next_message_id: MessageId(0), editing_message: None, collapsed_messages: HashMap::default(), pending_completion: None, + attachment_store, tool_registry, } } @@ -351,7 +375,12 @@ impl AssistantChat { editor }); composer_editor.clear(cx); - ChatMessage::User(UserMessage { id, body }) + + ChatMessage::User(UserMessage { + id, + body, + attachments: Vec::new(), + }) }); self.push_message(message, cx); } else { @@ -361,6 +390,29 @@ impl AssistantChat { let mode = *mode; self.pending_completion = Some(cx.spawn(move |this, mut cx| async move { + let attachments_task = this.update(&mut cx, |this, cx| { + let attachment_store = this.attachment_store.clone(); + attachment_store.call_all_attachment_tools(cx) + }); + + let attachments = maybe!(async { + let attachments_task = attachments_task?; + let attachments = attachments_task.await?; + + anyhow::Ok(attachments) + }) + .await + .log_err() + .unwrap_or_default(); + + // Set the attachments to the _last_ user message + this.update(&mut cx, |this, _cx| { + if let Some(ChatMessage::User(message)) = this.messages.last_mut() { + message.attachments = attachments; + } + }) + .log_err(); + Self::request_completion( this.clone(), mode, @@ -588,7 +640,11 @@ impl AssistantChat { let is_last = ix == self.messages.len() - 1; match &self.messages[ix] { - ChatMessage::User(UserMessage { id, body }) => div() + ChatMessage::User(UserMessage { + id, + body, + attachments, + }) => div() .id(SharedString::from(format!("message-{}-container", id.0))) .when(!is_last, |element| element.mb_2()) .map(|element| { @@ -596,6 +652,7 @@ impl AssistantChat { element.child(Composer::new( body.clone(), self.project_index_button.clone(), + self.active_file_button.clone(), crate::ui::ModelSelector::new( cx.view().downgrade(), self.model.clone(), @@ -629,6 +686,16 @@ impl AssistantChat { ) .element(ElementId::from(id.0), cx), ), + Some( + h_flex() + .gap_2() + .children( + attachments + .iter() + .map(|attachment| attachment.view.clone()), + ) + .into_any_element(), + ), self.is_message_collapsed(id), Box::new(cx.listener({ let id = *id; @@ -658,12 +725,38 @@ impl AssistantChat { ) }; + let tools = tool_calls + .iter() + .map(|tool_call| { + let result = &tool_call.result; + let name = tool_call.name.clone(); + match result { + Some(result) => div() + .p_2() + .child(result.into_any_element(&name)) + .into_any_element(), + None => div() + .p_2() + .child(Label::new(name).color(Color::Modified)) + .child("Running...") + .into_any_element(), + } + }) + .collect::>(); + + let tools_body = if tools.is_empty() { + None + } else { + Some(div().children(tools).into_any_element()) + }; + div() .when(!is_last, |element| element.mb_2()) .child(crate::ui::ChatMessage::new( *id, UserOrAssistant::Assistant, assistant_body, + tools_body, self.is_message_collapsed(id), Box::new(cx.listener({ let id = *id; @@ -672,22 +765,7 @@ impl AssistantChat { } })), )) - // TODO: Should the errors and tool calls get passed into `ChatMessage`? .child(self.render_error(error.clone(), ix, cx)) - .children(tool_calls.iter().map(|tool_call| { - let result = &tool_call.result; - let name = tool_call.name.clone(); - match result { - Some(result) => { - div().p_2().child(result.into_any_element(&name)).into_any() - } - None => div() - .p_2() - .child(Label::new(name).color(Color::Modified)) - .child("Running...") - .into_any(), - } - })) .into_any() } } @@ -698,11 +776,15 @@ impl AssistantChat { for message in &self.messages { match message { - ChatMessage::User(UserMessage { body, .. }) => { - // When we re-introduce contexts like active file, we'll inject them here instead of relying on the model to request them - // contexts.iter().for_each(|context| { - // completion_messages.extend(context.completion_messages(cx)) - // }); + ChatMessage::User(UserMessage { + body, attachments, .. + }) => { + completion_messages.extend( + attachments + .into_iter() + .filter_map(|attachment| attachment.message.clone()) + .map(|content| CompletionMessage::System { content }), + ); // Show user's message last so that the assistant is grounded in the user's request completion_messages.push(CompletionMessage::User { @@ -773,6 +855,7 @@ impl Render for AssistantChat { .child(Composer::new( self.composer_editor.clone(), self.project_index_button.clone(), + self.active_file_button.clone(), crate::ui::ModelSelector::new(cx.view().downgrade(), self.model.clone()) .into_any_element(), )) @@ -807,6 +890,7 @@ impl ChatMessage { struct UserMessage { id: MessageId, body: View, + attachments: Vec, } struct AssistantMessage { diff --git a/crates/assistant2/src/attachments.rs b/crates/assistant2/src/attachments.rs new file mode 100644 index 0000000000..b5e73d6177 --- /dev/null +++ b/crates/assistant2/src/attachments.rs @@ -0,0 +1,240 @@ +use std::{ + any::TypeId, + sync::{ + atomic::{AtomicBool, Ordering::SeqCst}, + Arc, + }, +}; + +use anyhow::{anyhow, Result}; +use collections::HashMap; +use editor::Editor; +use futures::future::join_all; +use gpui::{AnyView, Render, Task, View, WeakView}; +use ui::{prelude::*, ButtonLike, Tooltip, WindowContext}; +use util::{maybe, ResultExt}; +use workspace::Workspace; + +/// A collected attachment from running an attachment tool +pub struct UserAttachment { + pub message: Option, + pub view: AnyView, +} + +pub struct UserAttachmentStore { + attachment_tools: HashMap, +} + +/// Internal representation of an attachment tool to allow us to treat them dynamically +struct DynamicAttachment { + enabled: AtomicBool, + call: Box Task>>, +} + +impl UserAttachmentStore { + pub fn new() -> Self { + Self { + attachment_tools: HashMap::default(), + } + } + + pub fn register(&mut self, attachment: A) { + let call = Box::new(move |cx: &mut WindowContext| { + let result = attachment.run(cx); + + cx.spawn(move |mut cx| async move { + let result: Result = result.await; + let message = A::format(&result); + let view = cx.update(|cx| A::view(result, cx))?; + + Ok(UserAttachment { + message, + view: view.into(), + }) + }) + }); + + self.attachment_tools.insert( + TypeId::of::(), + DynamicAttachment { + call, + enabled: AtomicBool::new(true), + }, + ); + } + + pub fn set_attachment_tool_enabled(&self, is_enabled: bool) { + if let Some(attachment) = self.attachment_tools.get(&TypeId::of::()) { + attachment.enabled.store(is_enabled, SeqCst); + } + } + + pub fn is_attachment_tool_enabled(&self) -> bool { + if let Some(attachment) = self.attachment_tools.get(&TypeId::of::()) { + attachment.enabled.load(SeqCst) + } else { + false + } + } + + pub fn call( + &self, + cx: &mut WindowContext, + ) -> Task> { + let Some(attachment) = self.attachment_tools.get(&TypeId::of::()) else { + return Task::ready(Err(anyhow!("no attachment tool"))); + }; + + (attachment.call)(cx) + } + + pub fn call_all_attachment_tools( + self: Arc, + cx: &mut WindowContext<'_>, + ) -> Task>> { + let this = self.clone(); + cx.spawn(|mut cx| async move { + let attachment_tasks = cx.update(|cx| { + let mut tasks = Vec::new(); + for attachment in this + .attachment_tools + .values() + .filter(|attachment| attachment.enabled.load(SeqCst)) + { + tasks.push((attachment.call)(cx)) + } + + tasks + })?; + + let attachments = join_all(attachment_tasks.into_iter()).await; + + Ok(attachments + .into_iter() + .filter_map(|attachment| attachment.log_err()) + .collect()) + }) + } +} + +/// +pub trait AttachmentTool { + type Output: 'static; + type View: Render; + + fn run(&self, cx: &mut WindowContext) -> Task>; + + fn format(output: &Result) -> Option; + + fn view(output: Result, cx: &mut WindowContext) -> View; +} + +pub struct ActiveEditorAttachment { + filename: Arc, + language: Arc, + text: Arc, +} + +pub struct FileAttachmentView { + output: Result, +} + +impl Render for FileAttachmentView { + fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { + match &self.output { + Ok(attachment) => { + let filename = attachment.filename.clone(); + + // todo!(): make the button link to the actual file to open + ButtonLike::new("file-attachment") + .child( + h_flex() + .gap_1() + .bg(cx.theme().colors().editor_background) + .rounded_md() + .child(ui::Icon::new(IconName::File)) + .child(filename.to_string()), + ) + .tooltip({ + move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx) + }) + .into_any_element() + } + // todo!(): show a better error view when the file attaching didn't work + Err(err) => div().child(err.to_string()).into_any_element(), + } + } +} + +pub struct ActiveEditorAttachmentTool { + workspace: WeakView, +} + +impl ActiveEditorAttachmentTool { + pub fn new(workspace: WeakView, _cx: &mut WindowContext) -> Self { + Self { workspace } + } +} + +impl AttachmentTool for ActiveEditorAttachmentTool { + type Output = ActiveEditorAttachment; + type View = FileAttachmentView; + + fn run(&self, cx: &mut WindowContext) -> Task> { + Task::ready(maybe!({ + let active_buffer = self + .workspace + .update(cx, |workspace, cx| { + workspace + .active_item(cx) + .and_then(|item| Some(item.act_as::(cx)?.read(cx).buffer().clone())) + })? + .ok_or_else(|| anyhow!("no active buffer"))?; + + let buffer = active_buffer.read(cx); + + if let Some(singleton) = buffer.as_singleton() { + let singleton = singleton.read(cx); + + let filename = singleton + .file() + .map(|file| file.path().to_string_lossy()) + .unwrap_or("Untitled".into()); + + let text = singleton.text(); + + let language = singleton + .language() + .map(|l| { + let name = l.code_fence_block_name(); + name.to_string() + }) + .unwrap_or_default(); + + return Ok(ActiveEditorAttachment { + filename: filename.into(), + language: language.into(), + text: text.into(), + }); + } + + Err(anyhow!("no active buffer")) + })) + } + + fn format(output: &Result) -> Option { + let output = output.as_ref().ok()?; + + let filename = &output.filename; + let language = &output.language; + let text = &output.text; + + Some(format!( + "User's active file `{filename}`:\n\n```{language}\n{text}```\n\n" + )) + } + + fn view(output: Result, cx: &mut WindowContext) -> View { + cx.new_view(|_cx| FileAttachmentView { output }) + } +} diff --git a/crates/assistant2/src/ui.rs b/crates/assistant2/src/ui.rs index 4dffffef91..3333620a47 100644 --- a/crates/assistant2/src/ui.rs +++ b/crates/assistant2/src/ui.rs @@ -1,3 +1,4 @@ +mod active_file_button; mod chat_message; mod chat_notice; mod composer; @@ -6,6 +7,7 @@ mod project_index_button; #[cfg(feature = "stories")] mod stories; +pub use active_file_button::*; pub use chat_message::*; pub use chat_notice::*; pub use composer::*; diff --git a/crates/assistant2/src/ui/active_file_button.rs b/crates/assistant2/src/ui/active_file_button.rs new file mode 100644 index 0000000000..d6381b6b04 --- /dev/null +++ b/crates/assistant2/src/ui/active_file_button.rs @@ -0,0 +1,133 @@ +use crate::attachments::{ActiveEditorAttachmentTool, UserAttachmentStore}; +use editor::Editor; +use gpui::{prelude::*, Subscription, View}; +use std::sync::Arc; +use ui::{prelude::*, ButtonLike, Color, Icon, IconName, Tooltip}; +use workspace::Workspace; + +#[derive(Clone)] +enum Status { + ActiveFile(String), + #[allow(dead_code)] + NoFile, +} + +pub struct ActiveFileButton { + attachment_store: Arc, + status: Status, + #[allow(dead_code)] + workspace_subscription: Subscription, +} + +impl ActiveFileButton { + pub fn new( + attachment_store: Arc, + workspace: View, + cx: &mut ViewContext, + ) -> Self { + let workspace_subscription = cx.subscribe(&workspace, Self::handle_workspace_event); + + cx.defer(move |this, cx| this.update_active_buffer(workspace.clone(), cx)); + + Self { + attachment_store, + status: Status::NoFile, + workspace_subscription, + } + } + + pub fn set_enabled(&mut self, enabled: bool) { + self.attachment_store + .set_attachment_tool_enabled::(enabled); + } + + pub fn update_active_buffer(&mut self, workspace: View, cx: &mut ViewContext) { + let active_buffer = workspace + .read(cx) + .active_item(cx) + .and_then(|item| Some(item.act_as::(cx)?.read(cx).buffer().clone())); + + if let Some(buffer) = active_buffer { + let buffer = buffer.read(cx); + + if let Some(singleton) = buffer.as_singleton() { + let singleton = singleton.read(cx); + + let filename: String = singleton + .file() + .map(|file| file.path().to_string_lossy()) + .unwrap_or("Untitled".into()) + .into(); + + self.status = Status::ActiveFile(filename); + } + } + } + + fn handle_workspace_event( + &mut self, + workspace: View, + event: &workspace::Event, + cx: &mut ViewContext, + ) { + if let workspace::Event::ActiveItemChanged = event { + self.update_active_buffer(workspace, cx); + } + } +} + +impl Render for ActiveFileButton { + fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { + let is_enabled = self + .attachment_store + .is_attachment_tool_enabled::(); + + let icon = if is_enabled { + Icon::new(IconName::File) + .size(IconSize::XSmall) + .color(Color::Default) + } else { + Icon::new(IconName::File) + .size(IconSize::XSmall) + .color(Color::Disabled) + }; + + let indicator = None; + + let status = self.status.clone(); + + ButtonLike::new("active-file-button") + .child( + ui::IconWithIndicator::new(icon, indicator) + .indicator_border_color(Some(gpui::transparent_black())), + ) + .tooltip({ + move |cx| { + let status = status.clone(); + let (tooltip, meta) = match (is_enabled, status) { + (false, _) => ( + "Active file disabled".to_string(), + Some("Click to enable".to_string()), + ), + (true, Status::ActiveFile(filename)) => ( + format!("Active file {filename} enabled"), + Some("Click to disable".to_string()), + ), + (true, Status::NoFile) => { + ("No file active for conversation".to_string(), None) + } + }; + + if let Some(meta) = meta { + Tooltip::with_meta(tooltip, None, meta, cx) + } else { + Tooltip::text(tooltip, cx) + } + } + }) + .on_click(cx.listener(move |this, _, cx| { + this.set_enabled(!is_enabled); + cx.notify(); + })) + } +} diff --git a/crates/assistant2/src/ui/chat_message.rs b/crates/assistant2/src/ui/chat_message.rs index ab26979897..cfbac5523e 100644 --- a/crates/assistant2/src/ui/chat_message.rs +++ b/crates/assistant2/src/ui/chat_message.rs @@ -16,6 +16,7 @@ pub struct ChatMessage { id: MessageId, player: UserOrAssistant, message: Option, + tools_used: Option, collapsed: bool, on_collapse_handle_click: Box, } @@ -25,6 +26,7 @@ impl ChatMessage { id: MessageId, player: UserOrAssistant, message: Option, + tools_used: Option, collapsed: bool, on_collapse_handle_click: Box, ) -> Self { @@ -32,6 +34,7 @@ impl ChatMessage { id, player, message, + tools_used, collapsed, on_collapse_handle_click, } @@ -66,6 +69,10 @@ impl RenderOnce for ChatMessage { // Clamp the message height to exactly 1.5 lines when collapsed. let collapsed_height = content_padding.to_pixels(cx.rem_size()) + cx.line_height() * 1.5; + let tools_used = self + .tools_used + .map(|attachment| div().mt_3().child(attachment)); + let content = self.message.map(|message| { div() .overflow_hidden() @@ -75,6 +82,7 @@ impl RenderOnce for ChatMessage { .when(self.collapsed, |this| this.h(collapsed_height)) .bg(cx.theme().colors().surface_background) .child(message) + .children(tools_used) }); v_flex() diff --git a/crates/assistant2/src/ui/composer.rs b/crates/assistant2/src/ui/composer.rs index 4bd75ecb68..b6515a9548 100644 --- a/crates/assistant2/src/ui/composer.rs +++ b/crates/assistant2/src/ui/composer.rs @@ -1,14 +1,18 @@ -use crate::{ui::ProjectIndexButton, AssistantChat, CompletionProvider}; +use crate::{ + ui::{ActiveFileButton, ProjectIndexButton}, + AssistantChat, CompletionProvider, +}; use editor::{Editor, EditorElement, EditorStyle}; use gpui::{AnyElement, FontStyle, FontWeight, TextStyle, View, WeakView, WhiteSpace}; use settings::Settings; use theme::ThemeSettings; -use ui::{popover_menu, prelude::*, ButtonLike, ContextMenu, Tooltip}; +use ui::{popover_menu, prelude::*, ButtonLike, ContextMenu, Divider, Tooltip}; #[derive(IntoElement)] pub struct Composer { editor: View, project_index_button: Option>, + active_file_button: Option>, model_selector: AnyElement, } @@ -16,11 +20,13 @@ impl Composer { pub fn new( editor: View, project_index_button: Option>, + active_file_button: Option>, model_selector: AnyElement, ) -> Self { Self { editor, project_index_button, + active_file_button, model_selector, } } @@ -32,6 +38,14 @@ impl Composer { .map(|view| view.into_any_element()), ) } + + fn render_attachment_tools(&mut self, _cx: &mut WindowContext) -> impl IntoElement { + h_flex().children( + self.active_file_button + .clone() + .map(|view| view.into_any_element()), + ) + } } impl RenderOnce for Composer { @@ -83,7 +97,15 @@ impl RenderOnce for Composer { .gap_2() .justify_between() .w_full() - .child(h_flex().gap_1().child(self.render_tools(cx))) + .child( + h_flex().gap_1().child( + h_flex() + .gap_2() + .child(self.render_tools(cx)) + .child(Divider::vertical()) + .child(self.render_attachment_tools(cx)), + ), + ) .child(h_flex().gap_1().child(self.model_selector)), ), ), diff --git a/crates/assistant2/src/ui/stories/chat_message.rs b/crates/assistant2/src/ui/stories/chat_message.rs index a20a47c6c3..3058d0cdea 100644 --- a/crates/assistant2/src/ui/stories/chat_message.rs +++ b/crates/assistant2/src/ui/stories/chat_message.rs @@ -29,6 +29,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::User(Some(user_1.clone())), Some(div().child("What can I do here?").into_any_element()), + None, false, Box::new(|_, _| {}), ), @@ -39,6 +40,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::User(Some(user_1.clone())), Some(div().child("What can I do here?").into_any_element()), + None, true, Box::new(|_, _| {}), ), @@ -52,6 +54,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::Assistant, Some(div().child("You can talk to me!").into_any_element()), + None, false, Box::new(|_, _| {}), ), @@ -62,6 +65,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::Assistant, Some(div().child(MULTI_LINE_MESSAGE).into_any_element()), + None, true, Box::new(|_, _| {}), ), @@ -76,6 +80,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::User(Some(user_1.clone())), Some(div().child("What is Rust??").into_any_element()), + None, false, Box::new(|_, _| {}), )) @@ -83,6 +88,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::Assistant, Some(div().child("Rust is a multi-paradigm programming language focused on performance and safety").into_any_element()), + None, false, Box::new(|_, _| {}), )) @@ -90,6 +96,7 @@ impl Render for ChatMessageStory { MessageId(0), UserOrAssistant::User(Some(user_1)), Some(div().child("Sounds pretty cool!").into_any_element()), + None, false, Box::new(|_, _| {}), )),