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. <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
GitHub
parent
6bdcfad6ad
commit
3e5dcd1bec
10 changed files with 525 additions and 409 deletions
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue