Cargo.lock 🔗
@@ -382,7 +382,6 @@ dependencies = [
"editor",
"env_logger",
"feature_flags",
- "fs",
"futures 0.3.28",
"gpui",
"language",
Attachment store for assistant2 (#11327)
Kyle Kelley and Marshall created
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.
<img width="399" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/198429a5-82af-4b82-86f6-cb961f10de5c">
<img width="393" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/cfb52444-723b-4fc1-bddc-57e1810c512b">
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 <marshall@zed.dev>Change summary
Cargo.lock | 1
crates/assistant2/Cargo.toml | 1
crates/assistant2/examples/chat_with_functions.rs | 378 -----------------
crates/assistant2/src/assistant2.rs | 136 ++++-
crates/assistant2/src/attachments.rs | 240 ++++++++++
crates/assistant2/src/ui.rs | 2
crates/assistant2/src/ui/active_file_button.rs | 133 +++++
crates/assistant2/src/ui/chat_message.rs | 8
crates/assistant2/src/ui/composer.rs | 28 +
crates/assistant2/src/ui/stories/chat_message.rs | 7
10 files changed, 525 insertions(+), 409 deletions(-)
Detailed changes
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
crates/assistant2/examples/chat_with_functions.rs 🔗
@@ -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<Die>,
-}
-
-#[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<DieRoll>,
-}
-
-pub struct DiceView {
- result: Result<DiceRoll>,
-}
-
-impl Render for DiceView {
- fn render(&mut self, _cx: &mut ViewContext<Self>) -> 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<gpui::Result<Self::Output>> {
- 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<Self::Output>,
- cx: &mut WindowContext,
- ) -> gpui::View<Self::View> {
- cx.new_view(|_cx| DiceView { result })
- }
-
- fn format(_: &Self::Input, output: &Result<Self::Output>) -> 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<dyn Fs>,
- root_dir: PathBuf,
-}
-
-impl FileBrowserTool {
- fn new(fs: Arc<dyn Fs>, 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<String> },
- Cat { content: String },
-}
-
-pub struct FileBrowserView {
- result: Result<FileBrowserOutput>,
-}
-
-impl Render for FileBrowserView {
- fn render(&mut self, cx: &mut ViewContext<Self>) -> 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<gpui::Result<Self::Output>> {
- 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<Self::Output>,
- cx: &mut WindowContext,
- ) -> gpui::View<Self::View> {
- cx.new_view(|_cx| FileBrowserView { result })
- }
-
- fn format(_input: &Self::Input, output: &Result<Self::Output>) -> 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<AssistantPanel>,
-}
-
-impl Example {
- fn new(
- language_registry: Arc<LanguageRegistry>,
- tool_registry: Arc<ToolRegistry>,
- user_store: Model<UserStore>,
- cx: &mut ViewContext<Self>,
- ) -> 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<Self>) -> impl ui::prelude::IntoElement {
- div().size_full().child(self.assistant_panel.clone())
- }
-}
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<LanguageRegistry>,
+ attachment_store: Arc<UserAttachmentStore>,
tool_registry: Arc<ToolRegistry>,
user_store: Model<UserStore>,
project_index: Option<Model<ProjectIndex>>,
+ workspace: WeakView<Workspace>,
cx: &mut ViewContext<Self>,
) -> 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<LanguageRegistry>,
composer_editor: View<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
+ active_file_button: Option<View<ActiveFileButton>>,
user_store: Model<UserStore>,
next_message_id: MessageId,
collapsed_messages: HashMap<MessageId, bool>,
editing_message: Option<EditingMessage>,
pending_completion: Option<Task<()>>,
+ attachment_store: Arc<UserAttachmentStore>,
tool_registry: Arc<ToolRegistry>,
project_index: Option<Model<ProjectIndex>>,
}
@@ -247,9 +258,11 @@ struct EditingMessage {
impl AssistantChat {
fn new(
language_registry: Arc<LanguageRegistry>,
+ attachment_store: Arc<UserAttachmentStore>,
tool_registry: Arc<ToolRegistry>,
user_store: Model<UserStore>,
project_index: Option<Model<ProjectIndex>>,
+ workspace: WeakView<Workspace>,
cx: &mut ViewContext<Self>,
) -> 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::<Vec<AnyElement>>();
+
+ 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<Editor>,
+ attachments: Vec<UserAttachment>,
}
struct AssistantMessage {
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<String>,
+ pub view: AnyView,
+}
+
+pub struct UserAttachmentStore {
+ attachment_tools: HashMap<TypeId, DynamicAttachment>,
+}
+
+/// Internal representation of an attachment tool to allow us to treat them dynamically
+struct DynamicAttachment {
+ enabled: AtomicBool,
+ call: Box<dyn Fn(&mut WindowContext) -> Task<Result<UserAttachment>>>,
+}
+
+impl UserAttachmentStore {
+ pub fn new() -> Self {
+ Self {
+ attachment_tools: HashMap::default(),
+ }
+ }
+
+ pub fn register<A: AttachmentTool + 'static>(&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<A::Output> = 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::<A>(),
+ DynamicAttachment {
+ call,
+ enabled: AtomicBool::new(true),
+ },
+ );
+ }
+
+ pub fn set_attachment_tool_enabled<A: AttachmentTool + 'static>(&self, is_enabled: bool) {
+ if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
+ attachment.enabled.store(is_enabled, SeqCst);
+ }
+ }
+
+ pub fn is_attachment_tool_enabled<A: AttachmentTool + 'static>(&self) -> bool {
+ if let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) {
+ attachment.enabled.load(SeqCst)
+ } else {
+ false
+ }
+ }
+
+ pub fn call<A: AttachmentTool + 'static>(
+ &self,
+ cx: &mut WindowContext,
+ ) -> Task<Result<UserAttachment>> {
+ let Some(attachment) = self.attachment_tools.get(&TypeId::of::<A>()) else {
+ return Task::ready(Err(anyhow!("no attachment tool")));
+ };
+
+ (attachment.call)(cx)
+ }
+
+ pub fn call_all_attachment_tools(
+ self: Arc<Self>,
+ cx: &mut WindowContext<'_>,
+ ) -> Task<Result<Vec<UserAttachment>>> {
+ 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<Result<Self::Output>>;
+
+ fn format(output: &Result<Self::Output>) -> Option<String>;
+
+ fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View>;
+}
+
+pub struct ActiveEditorAttachment {
+ filename: Arc<str>,
+ language: Arc<str>,
+ text: Arc<str>,
+}
+
+pub struct FileAttachmentView {
+ output: Result<ActiveEditorAttachment>,
+}
+
+impl Render for FileAttachmentView {
+ fn render(&mut self, cx: &mut ViewContext<Self>) -> 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<Workspace>,
+}
+
+impl ActiveEditorAttachmentTool {
+ pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
+ Self { workspace }
+ }
+}
+
+impl AttachmentTool for ActiveEditorAttachmentTool {
+ type Output = ActiveEditorAttachment;
+ type View = FileAttachmentView;
+
+ fn run(&self, cx: &mut WindowContext) -> Task<Result<ActiveEditorAttachment>> {
+ Task::ready(maybe!({
+ let active_buffer = self
+ .workspace
+ .update(cx, |workspace, cx| {
+ workspace
+ .active_item(cx)
+ .and_then(|item| Some(item.act_as::<Editor>(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<Self::Output>) -> Option<String> {
+ 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<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
+ cx.new_view(|_cx| FileAttachmentView { output })
+ }
+}
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::*;
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<UserAttachmentStore>,
+ status: Status,
+ #[allow(dead_code)]
+ workspace_subscription: Subscription,
+}
+
+impl ActiveFileButton {
+ pub fn new(
+ attachment_store: Arc<UserAttachmentStore>,
+ workspace: View<Workspace>,
+ cx: &mut ViewContext<Self>,
+ ) -> 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::<ActiveEditorAttachmentTool>(enabled);
+ }
+
+ pub fn update_active_buffer(&mut self, workspace: View<Workspace>, cx: &mut ViewContext<Self>) {
+ let active_buffer = workspace
+ .read(cx)
+ .active_item(cx)
+ .and_then(|item| Some(item.act_as::<Editor>(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<Workspace>,
+ event: &workspace::Event,
+ cx: &mut ViewContext<Self>,
+ ) {
+ if let workspace::Event::ActiveItemChanged = event {
+ self.update_active_buffer(workspace, cx);
+ }
+ }
+}
+
+impl Render for ActiveFileButton {
+ fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
+ let is_enabled = self
+ .attachment_store
+ .is_attachment_tool_enabled::<ActiveEditorAttachmentTool>();
+
+ 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();
+ }))
+ }
+}
crates/assistant2/src/ui/chat_message.rs 🔗
@@ -16,6 +16,7 @@ pub struct ChatMessage {
id: MessageId,
player: UserOrAssistant,
message: Option<AnyElement>,
+ tools_used: Option<AnyElement>,
collapsed: bool,
on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
}
@@ -25,6 +26,7 @@ impl ChatMessage {
id: MessageId,
player: UserOrAssistant,
message: Option<AnyElement>,
+ tools_used: Option<AnyElement>,
collapsed: bool,
on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
) -> 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()
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<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
+ active_file_button: Option<View<ActiveFileButton>>,
model_selector: AnyElement,
}
@@ -16,11 +20,13 @@ impl Composer {
pub fn new(
editor: View<Editor>,
project_index_button: Option<View<ProjectIndexButton>>,
+ active_file_button: Option<View<ActiveFileButton>>,
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)),
),
),
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(|_, _| {}),
)),