mirrors/zed
1
0
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2024-11-28 09:54:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 11

Attachment store for assistant2 (#11327)

Browse source 
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>
This commit is contained in:
Kyle Kelley 2024-05-03 14:48:00 -07:00 committed by GitHub
parent 6bdcfad6ad
commit 3e5dcd1bec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 525 additions and 409 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -382,7 +382,6 @@ dependencies = [
"editor",
"env_logger",
"feature_flags",
"fs",
"futures 0.3.28",
"gpui",
"language",

1
crates/assistant2/Cargo.toml
View file

@ -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

378
crates/assistant2/examples/chat_with_functions.rs
View file

@ -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())
}
}

136
crates/assistant2/src/assistant2.rs
View file

@ -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 {

240
crates/assistant2/src/attachments.rs Normal file
View file

@ -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 })
}
}

2
crates/assistant2/src/ui.rs
View file

@ -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::*;

133
crates/assistant2/src/ui/active_file_button.rs Normal file
View file

@ -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();
}))
}
}

8
crates/assistant2/src/ui/chat_message.rs
View file

@ -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()

28
crates/assistant2/src/ui/composer.rs
View file

@ -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)),
),
),

7
crates/assistant2/src/ui/stories/chat_message.rs
View file

@ -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(|_, _| {}),
)),