Centralize project context provided to the assistant (#11471)
This PR restructures the way that tools and attachments add information about the current project to a conversation with the assistant. Rather than each tool call or attachment generating a new tool or system message containing information about the project, they can all collectively mutate a new type called a `ProjectContext`, which stores all of the project data that should be sent to the assistant. That data is then formatted in a single place, and passed to the assistant in one system message. This prevents multiple tools/attachments from including redundant context. Release Notes: - N/A --------- Co-authored-by: Kyle <kylek@zed.dev>
This commit is contained in:
Max Brunsfeld
GitHub
parent
f2a415135b
commit
a64e20ed96
15 changed files with 841 additions and 518 deletions
|
|
@ -1,137 +1,18 @@
|
|||
use std::{
|
||||
any::TypeId,
|
||||
sync::{
|
||||
atomic::{AtomicBool, Ordering::SeqCst},
|
||||
Arc,
|
||||
},
|
||||
};
|
||||
pub mod active_file;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use collections::HashMap;
|
||||
use assistant_tooling::{LanguageModelAttachment, ProjectContext, ToolOutput};
|
||||
use editor::Editor;
|
||||
use futures::future::join_all;
|
||||
use gpui::{AnyView, Render, Task, View, WeakView};
|
||||
use gpui::{Render, Task, View, WeakModel, WeakView};
|
||||
use language::Buffer;
|
||||
use project::ProjectPath;
|
||||
use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
|
||||
use util::{maybe, ResultExt};
|
||||
use util::maybe;
|
||||
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>,
|
||||
buffer: WeakModel<Buffer>,
|
||||
path: Option<ProjectPath>,
|
||||
}
|
||||
|
||||
pub struct FileAttachmentView {
|
||||
|
|
@ -142,7 +23,13 @@ impl Render for FileAttachmentView {
|
|||
fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
|
||||
match &self.output {
|
||||
Ok(attachment) => {
|
||||
let filename = attachment.filename.clone();
|
||||
let filename: SharedString = attachment
|
||||
.path
|
||||
.as_ref()
|
||||
.and_then(|p| p.path.file_name()?.to_str())
|
||||
.unwrap_or("Untitled")
|
||||
.to_string()
|
||||
.into();
|
||||
|
||||
// todo!(): make the button link to the actual file to open
|
||||
ButtonLike::new("file-attachment")
|
||||
|
|
@ -152,7 +39,7 @@ impl Render for FileAttachmentView {
|
|||
.bg(cx.theme().colors().editor_background)
|
||||
.rounded_md()
|
||||
.child(ui::Icon::new(IconName::File))
|
||||
.child(filename.to_string()),
|
||||
.child(filename.clone()),
|
||||
)
|
||||
.tooltip({
|
||||
move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
|
||||
|
|
@ -164,6 +51,20 @@ impl Render for FileAttachmentView {
|
|||
}
|
||||
}
|
||||
|
||||
impl ToolOutput for FileAttachmentView {
|
||||
fn generate(&self, project: &mut ProjectContext, cx: &mut WindowContext) -> String {
|
||||
if let Ok(result) = &self.output {
|
||||
if let Some(path) = &result.path {
|
||||
project.add_file(path.clone());
|
||||
return format!("current file: {}", path.path.display());
|
||||
} else if let Some(buffer) = result.buffer.upgrade() {
|
||||
return format!("current untitled buffer text:\n{}", buffer.read(cx).text());
|
||||
}
|
||||
}
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ActiveEditorAttachmentTool {
|
||||
workspace: WeakView<Workspace>,
|
||||
}
|
||||
|
|
@ -174,7 +75,7 @@ impl ActiveEditorAttachmentTool {
|
|||
}
|
||||
}
|
||||
|
||||
impl AttachmentTool for ActiveEditorAttachmentTool {
|
||||
impl LanguageModelAttachment for ActiveEditorAttachmentTool {
|
||||
type Output = ActiveEditorAttachment;
|
||||
type View = FileAttachmentView;
|
||||
|
||||
|
|
@ -191,47 +92,22 @@ impl AttachmentTool for ActiveEditorAttachmentTool {
|
|||
|
||||
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();
|
||||
|
||||
if let Some(buffer) = buffer.as_singleton() {
|
||||
let path =
|
||||
project::File::from_dyn(buffer.read(cx).file()).map(|file| ProjectPath {
|
||||
worktree_id: file.worktree_id(cx),
|
||||
path: file.path.clone(),
|
||||
});
|
||||
return Ok(ActiveEditorAttachment {
|
||||
filename: filename.into(),
|
||||
language: language.into(),
|
||||
text: text.into(),
|
||||
buffer: buffer.downgrade(),
|
||||
path,
|
||||
});
|
||||
} else {
|
||||
Err(anyhow!("no active buffer"))
|
||||
}
|
||||
|
||||
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 })
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue