Language Model Tool for commenting in a multibuffer (#11509)

Language Model can now open multibuffers and insert comments as block
decorations in the editor.


![image](https://github.com/zed-industries/zed/assets/836375/f4456ad0-66e7-4ad6-a2b3-63810a3223a5)


Release Notes:

- N/A

---------

Co-authored-by: max <max@zed.dev>
Co-authored-by: marshall <marshall@zed.dev>
Co-authored-by: nate <nate@zed.dev>

crates/assistant2/evals/list-of-into-element.md

@@ -1 +1 @@
-> Give me a comprehensive list of all the elements define in my project (impl Element for {}, impl<T: 'static> Element for {}, impl IntoElement for {})
+> Give me a comprehensive list of all the elements defined in my project using the following query: `impl Element for {}, impl<T: 'static> Element for {}, impl IntoElement for {})`

crates/assistant2/evals/settings-file.md

@@ -0,0 +1,3 @@
+Use tools frequently, especially when referring to files and code. I prefer to see the file directly rather than you just chatting with me.
+
+Teach me everything you can about settings files and how they're loaded.

crates/assistant2/src/assistant2.rs

@@ -31,6 +31,7 @@ use semantic_index::{CloudEmbeddingProvider, ProjectIndex, ProjectIndexDebugView
 use serde::Deserialize;
 use settings::Settings;
 use std::sync::Arc;
+use tools::OpenBufferTool;
 use ui::{ActiveFileButton, Composer, ProjectIndexButton};
 use util::{maybe, paths::EMBEDDINGS_DIR, ResultExt};
 use workspace::{
@@ -125,15 +126,16 @@ impl AssistantPanel {
                 let mut tool_registry = ToolRegistry::new();
                 tool_registry
                     .register(ProjectIndexTool::new(project_index.clone()), cx)
-                    .context("failed to register ProjectIndexTool")
+                    .unwrap();
-                    .log_err();
                 tool_registry
                     .register(
                         CreateBufferTool::new(workspace.clone(), project.clone()),
                         cx,
                     )
-                    .context("failed to register CreateBufferTool")
+                    .unwrap();
-                    .log_err();
+                tool_registry
+                    .register(OpenBufferTool::new(workspace.clone(), project.clone()), cx)
+                    .unwrap();
 
                 let mut attachment_registry = AttachmentRegistry::new();
                 attachment_registry

crates/assistant2/src/attachments.rs

@@ -1,114 +1,3 @@
-pub mod active_file;
+mod active_file;
 
-use anyhow::{anyhow, Result};
+pub use active_file::*;
-use assistant_tooling::{LanguageModelAttachment, ProjectContext, ToolOutput};
-use editor::Editor;
-use gpui::{Render, Task, View, WeakModel, WeakView};
-use language::Buffer;
-use project::ProjectPath;
-use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
-use util::maybe;
-use workspace::Workspace;
-
-pub struct ActiveEditorAttachment {
-    buffer: WeakModel<Buffer>,
-    path: Option<ProjectPath>,
-}
-
-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: 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")
-                    .child(
-                        h_flex()
-                            .gap_1()
-                            .bg(cx.theme().colors().editor_background)
-                            .rounded_md()
-                            .child(ui::Icon::new(IconName::File))
-                            .child(filename.clone()),
-                    )
-                    .tooltip({
-                        move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
-                    })
-                    .into_any_element()
-            }
-            Err(err) => div().child(err.to_string()).into_any_element(),
-        }
-    }
-}
-
-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>,
-}
-
-impl ActiveEditorAttachmentTool {
-    pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
-        Self { workspace }
-    }
-}
-
-impl LanguageModelAttachment 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(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 {
-                    buffer: buffer.downgrade(),
-                    path,
-                });
-            } else {
-                Err(anyhow!("no active buffer"))
-            }
-        }))
-    }
-
-    fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
-        cx.new_view(|_cx| FileAttachmentView { output })
-    }
-}

crates/assistant2/src/attachments/active_file.rs

@@ -1 +1,112 @@
+use anyhow::{anyhow, Result};
+use assistant_tooling::{LanguageModelAttachment, ProjectContext, ToolOutput};
+use editor::Editor;
+use gpui::{Render, Task, View, WeakModel, WeakView};
+use language::Buffer;
+use project::ProjectPath;
+use ui::{prelude::*, ButtonLike, Tooltip, WindowContext};
+use util::maybe;
+use workspace::Workspace;
 
+pub struct ActiveEditorAttachment {
+    buffer: WeakModel<Buffer>,
+    path: Option<ProjectPath>,
+}
+
+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: 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")
+                    .child(
+                        h_flex()
+                            .gap_1()
+                            .bg(cx.theme().colors().editor_background)
+                            .rounded_md()
+                            .child(ui::Icon::new(IconName::File))
+                            .child(filename.clone()),
+                    )
+                    .tooltip({
+                        move |cx| Tooltip::with_meta("File Attached", None, filename.clone(), cx)
+                    })
+                    .into_any_element()
+            }
+            Err(err) => div().child(err.to_string()).into_any_element(),
+        }
+    }
+}
+
+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>,
+}
+
+impl ActiveEditorAttachmentTool {
+    pub fn new(workspace: WeakView<Workspace>, _cx: &mut WindowContext) -> Self {
+        Self { workspace }
+    }
+}
+
+impl LanguageModelAttachment 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(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 {
+                    buffer: buffer.downgrade(),
+                    path,
+                });
+            } else {
+                Err(anyhow!("no active buffer"))
+            }
+        }))
+    }
+
+    fn view(output: Result<Self::Output>, cx: &mut WindowContext) -> View<Self::View> {
+        cx.new_view(|_cx| FileAttachmentView { output })
+    }
+}

crates/assistant2/src/tools.rs

@@ -1,5 +1,7 @@
 mod create_buffer;
+mod open_buffer;
 mod project_index;
 
 pub use create_buffer::*;
+pub use open_buffer::*;
 pub use project_index::*;

crates/assistant2/src/tools/open_buffer.rs

@@ -0,0 +1,182 @@
+use anyhow::Result;
+use assistant_tooling::{LanguageModelTool, ProjectContext, ToolOutput};
+use editor::{
+    display_map::{BlockContext, BlockDisposition, BlockProperties, BlockStyle},
+    Editor, MultiBuffer,
+};
+use gpui::{prelude::*, AnyElement, Model, Task, View, WeakView};
+use language::ToPoint;
+use project::{Project, ProjectPath};
+use schemars::JsonSchema;
+use serde::Deserialize;
+use std::path::Path;
+use ui::prelude::*;
+use util::ResultExt;
+use workspace::Workspace;
+
+pub struct OpenBufferTool {
+    workspace: WeakView<Workspace>,
+    project: Model<Project>,
+}
+
+impl OpenBufferTool {
+    pub fn new(workspace: WeakView<Workspace>, project: Model<Project>) -> Self {
+        Self { workspace, project }
+    }
+}
+
+#[derive(Debug, Deserialize, JsonSchema, Clone)]
+pub struct ExplainInput {
+    /// Name for this set of excerpts
+    title: String,
+    excerpts: Vec<ExplainedExcerpt>,
+}
+
+#[derive(Debug, Deserialize, JsonSchema, Clone)]
+struct ExplainedExcerpt {
+    /// Path to the file
+    path: String,
+    /// Name of a symbol in the buffer to show
+    symbol_name: String,
+    /// Text to display near the symbol definition
+    comment: String,
+}
+
+impl LanguageModelTool for OpenBufferTool {
+    type Input = ExplainInput;
+    type Output = String;
+    type View = OpenBufferView;
+
+    fn name(&self) -> String {
+        "explain_code".to_string()
+    }
+
+    fn description(&self) -> String {
+        "Show and explain one or more code snippets from files in the current project. Code snippets are identified using a file path and the name of a symbol defined in that file.".to_string()
+    }
+
+    fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
+        let workspace = self.workspace.clone();
+        let project = self.project.clone();
+        let excerpts = input.excerpts.clone();
+        let title = input.title.clone();
+
+        let worktree_id = project.update(cx, |project, cx| {
+            let worktree = project.worktrees().next()?;
+            let worktree_id = worktree.read(cx).id();
+            Some(worktree_id)
+        });
+
+        let worktree_id = if let Some(worktree_id) = worktree_id {
+            worktree_id
+        } else {
+            return Task::ready(Err(anyhow::anyhow!("No worktree found")));
+        };
+
+        let buffer_tasks = project.update(cx, |project, cx| {
+            let excerpts = excerpts.clone();
+            excerpts
+                .iter()
+                .map(|excerpt| {
+                    let project_path = ProjectPath {
+                        worktree_id,
+                        path: Path::new(&excerpt.path).into(),
+                    };
+                    project.open_buffer(project_path.clone(), cx)
+                })
+                .collect::<Vec<_>>()
+        });
+
+        cx.spawn(move |mut cx| async move {
+            let buffers = futures::future::try_join_all(buffer_tasks).await?;
+
+            let multibuffer = cx.new_model(|_cx| {
+                MultiBuffer::new(0, language::Capability::ReadWrite).with_title(title)
+            })?;
+            let editor =
+                cx.new_view(|cx| Editor::for_multibuffer(multibuffer, Some(project), cx))?;
+
+            for (excerpt, buffer) in excerpts.iter().zip(buffers.iter()) {
+                let snapshot = buffer.update(&mut cx, |buffer, _cx| buffer.snapshot())?;
+
+                if let Some(outline) = snapshot.outline(None) {
+                    let matches = outline
+                        .search(&excerpt.symbol_name, cx.background_executor().clone())
+                        .await;
+                    if let Some(mat) = matches.first() {
+                        let item = &outline.items[mat.candidate_id];
+                        let start = item.range.start.to_point(&snapshot);
+                        editor.update(&mut cx, |editor, cx| {
+                            let ranges = editor.buffer().update(cx, |multibuffer, cx| {
+                                multibuffer.push_excerpts_with_context_lines(
+                                    buffer.clone(),
+                                    vec![start..start],
+                                    5,
+                                    cx,
+                                )
+                            });
+                            let explanation = SharedString::from(excerpt.comment.clone());
+                            editor.insert_blocks(
+                                [BlockProperties {
+                                    position: ranges[0].start,
+                                    height: 1,
+                                    style: BlockStyle::Fixed,
+                                    render: Box::new(move |cx| {
+                                        Self::render_note_block(&explanation, cx)
+                                    }),
+                                    disposition: BlockDisposition::Above,
+                                }],
+                                None,
+                                cx,
+                            );
+                        })?;
+                    }
+                }
+            }
+
+            workspace
+                .update(&mut cx, |workspace, cx| {
+                    workspace.add_item_to_active_pane(Box::new(editor.clone()), None, cx);
+                })
+                .log_err();
+
+            anyhow::Ok("showed comments to users in a new view".into())
+        })
+    }
+
+    fn output_view(
+        _: Self::Input,
+        output: Result<Self::Output>,
+        cx: &mut WindowContext,
+    ) -> View<Self::View> {
+        cx.new_view(|_cx| OpenBufferView { output })
+    }
+}
+
+impl OpenBufferTool {
+    fn render_note_block(explanation: &SharedString, _cx: &mut BlockContext) -> AnyElement {
+        div().child(explanation.clone()).into_any_element()
+    }
+}
+
+pub struct OpenBufferView {
+    output: Result<String>,
+}
+
+impl Render for OpenBufferView {
+    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+        match &self.output {
+            Ok(output) => div().child(output.clone().into_any_element()),
+            Err(error) => div().child(format!("failed to open path: {:?}", error)),
+        }
+    }
+}
+
+impl ToolOutput for OpenBufferView {
+    fn generate(&self, _: &mut ProjectContext, _: &mut WindowContext) -> String {
+        match &self.output {
+            Ok(output) => output.clone(),
+            Err(err) => format!("Failed to create buffer: {err:?}"),
+        }
+    }
+}