assistant2: Add create buffer tool (#11219)

This PR adds a new tool to the `assistant2` crate that allows the
assistant to create a new buffer with some content.

Release Notes:

- N/A

---------

Co-authored-by: Nathan <nathan@zed.dev>

crates/assistant2/examples/chat_with_functions.rs

@@ -122,7 +122,11 @@ impl LanguageModelTool for RollDiceTool {
         "Rolls N many dice and returns the results.".to_string()
     }
 
-    fn execute(&self, input: &Self::Input, _cx: &AppContext) -> Task<gpui::Result<Self::Output>> {
+    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();
@@ -223,7 +227,11 @@ impl LanguageModelTool for FileBrowserTool {
         "A tool for browsing the filesystem.".to_string()
     }
 
-    fn execute(&self, input: &Self::Input, cx: &AppContext) -> Task<gpui::Result<Self::Output>> {
+    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();

crates/assistant2/src/assistant2.rs

@@ -32,7 +32,7 @@ use workspace::{
 
 pub use assistant_settings::AssistantSettings;
 
-use crate::tools::ProjectIndexTool;
+use crate::tools::{CreateBufferTool, ProjectIndexTool};
 use crate::ui::UserOrAssistant;
 
 const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5;
@@ -121,6 +121,13 @@ impl AssistantPanel {
                     )
                     .context("failed to register ProjectIndexTool")
                     .log_err();
+                tool_registry
+                    .register(
+                        CreateBufferTool::new(workspace.clone(), project.clone()),
+                        cx,
+                    )
+                    .context("failed to register CreateBufferTool")
+                    .log_err();
 
                 let tool_registry = Arc::new(tool_registry);
 
@@ -542,7 +549,7 @@ impl AssistantChat {
                 .child(crate::ui::ChatMessage::new(
                     *id,
                     UserOrAssistant::User(self.user_store.read(cx).current_user()),
-                    body.clone().into_any_element(),
+                    Some(body.clone().into_any_element()),
                     self.is_message_collapsed(id),
                     Box::new(cx.listener({
                         let id = *id;
@@ -559,10 +566,15 @@ impl AssistantChat {
                 tool_calls,
                 ..
             }) => {
-                let assistant_body = if body.text.is_empty() && !tool_calls.is_empty() {
-                    div()
+                let assistant_body = if body.text.is_empty() {
+                    None
                 } else {
-                    div().p_2().child(body.element(ElementId::from(id.0), cx))
+                    Some(
+                        div()
+                            .p_2()
+                            .child(body.element(ElementId::from(id.0), cx))
+                            .into_any_element(),
+                    )
                 };
 
                 div()
@@ -570,7 +582,7 @@ impl AssistantChat {
                     .child(crate::ui::ChatMessage::new(
                         *id,
                         UserOrAssistant::Assistant,
-                        assistant_body.into_any_element(),
+                        assistant_body,
                         self.is_message_collapsed(id),
                         Box::new(cx.listener({
                             let id = *id;

crates/assistant2/src/tools.rs

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

crates/assistant2/src/tools/create_buffer.rs

@@ -0,0 +1,111 @@
+use anyhow::Result;
+use assistant_tooling::LanguageModelTool;
+use editor::Editor;
+use gpui::{prelude::*, Model, Task, View, WeakView};
+use project::Project;
+use schemars::JsonSchema;
+use serde::Deserialize;
+use ui::prelude::*;
+use util::ResultExt;
+use workspace::Workspace;
+
+pub struct CreateBufferTool {
+    workspace: WeakView<Workspace>,
+    project: Model<Project>,
+}
+
+impl CreateBufferTool {
+    pub fn new(workspace: WeakView<Workspace>, project: Model<Project>) -> Self {
+        Self { workspace, project }
+    }
+}
+
+#[derive(Debug, Deserialize, JsonSchema)]
+pub struct CreateBufferInput {
+    /// The contents of the buffer.
+    text: String,
+
+    /// The name of the language to use for the buffer.
+    ///
+    /// This should be a human-readable name, like "Rust", "JavaScript", or "Python".
+    language: String,
+}
+
+pub struct CreateBufferOutput {}
+
+impl LanguageModelTool for CreateBufferTool {
+    type Input = CreateBufferInput;
+    type Output = CreateBufferOutput;
+    type View = CreateBufferView;
+
+    fn name(&self) -> String {
+        "create_buffer".to_string()
+    }
+
+    fn description(&self) -> String {
+        "Create a new buffer in the current codebase".to_string()
+    }
+
+    fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
+        cx.spawn({
+            let workspace = self.workspace.clone();
+            let project = self.project.clone();
+            let text = input.text.clone();
+            let language_name = input.language.clone();
+            |mut cx| async move {
+                let language = cx
+                    .update(|cx| {
+                        project
+                            .read(cx)
+                            .languages()
+                            .language_for_name(&language_name)
+                    })?
+                    .await?;
+
+                let buffer = cx.update(|cx| {
+                    project.update(cx, |project, cx| {
+                        project.create_buffer(&text, Some(language), cx)
+                    })
+                })??;
+
+                workspace
+                    .update(&mut cx, |workspace, cx| {
+                        workspace.add_item_to_active_pane(
+                            Box::new(
+                                cx.new_view(|cx| Editor::for_buffer(buffer, Some(project), cx)),
+                            ),
+                            None,
+                            cx,
+                        );
+                    })
+                    .log_err();
+
+                Ok(CreateBufferOutput {})
+            }
+        })
+    }
+
+    fn format(input: &Self::Input, output: &Result<Self::Output>) -> String {
+        match output {
+            Ok(_) => format!("Created a new {} buffer", input.language),
+            Err(err) => format!("Failed to create buffer: {err:?}"),
+        }
+    }
+
+    fn output_view(
+        _tool_call_id: String,
+        _input: Self::Input,
+        _output: Result<Self::Output>,
+        cx: &mut WindowContext,
+    ) -> View<Self::View> {
+        cx.new_view(|_cx| CreateBufferView {})
+    }
+}
+
+pub struct CreateBufferView {}
+
+impl Render for CreateBufferView {
+    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+        div().child("Opening a buffer")
+    }
+}

crates/assistant2/src/tools/project_index.rs

@@ -1,6 +1,6 @@
 use anyhow::Result;
 use assistant_tooling::LanguageModelTool;
-use gpui::{prelude::*, AnyView, AppContext, Model, Task};
+use gpui::{prelude::*, AnyView, Model, Task};
 use project::Fs;
 use schemars::JsonSchema;
 use semantic_index::{ProjectIndex, Status};
@@ -138,7 +138,7 @@ impl LanguageModelTool for ProjectIndexTool {
         "Semantic search against the user's current codebase, returning excerpts related to the query by computing a dot product against embeddings of chunks and an embedding of the query".to_string()
     }
 
-    fn execute(&self, query: &Self::Input, cx: &AppContext) -> Task<Result<Self::Output>> {
+    fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
         let project_index = self.project_index.read(cx);
 
         let status = project_index.status();

crates/assistant2/src/ui/chat_message.rs

@@ -15,7 +15,7 @@ pub enum UserOrAssistant {
 pub struct ChatMessage {
     id: MessageId,
     player: UserOrAssistant,
-    message: AnyElement,
+    message: Option<AnyElement>,
     collapsed: bool,
     on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
 }
@@ -24,7 +24,7 @@ impl ChatMessage {
     pub fn new(
         id: MessageId,
         player: UserOrAssistant,
-        message: AnyElement,
+        message: Option<AnyElement>,
         collapsed: bool,
         on_collapse_handle_click: Box<dyn Fn(&ClickEvent, &mut WindowContext) + 'static>,
     ) -> Self {
@@ -65,19 +65,21 @@ impl RenderOnce for ChatMessage {
                         this.bg(cx.theme().colors().element_hover)
                     }),
             );
-        let content = div()
-            .overflow_hidden()
-            .w_full()
-            .p_4()
-            .rounded_lg()
-            .when(self.collapsed, |this| this.h(collapsed_height))
-            .bg(cx.theme().colors().surface_background)
-            .child(self.message);
+        let content = self.message.map(|message| {
+            div()
+                .overflow_hidden()
+                .w_full()
+                .p_4()
+                .rounded_lg()
+                .when(self.collapsed, |this| this.h(collapsed_height))
+                .bg(cx.theme().colors().surface_background)
+                .child(message)
+        });
 
         v_flex()
             .gap_1()
             .child(ChatMessageHeader::new(self.player))
-            .child(h_flex().gap_3().child(collapse_handle).child(content))
+            .child(h_flex().gap_3().child(collapse_handle).children(content))
     }
 }
 

crates/assistant_tooling/src/registry.rs

@@ -120,8 +120,8 @@ impl ToolRegistry {
 #[cfg(test)]
 mod test {
     use super::*;
-    use gpui::View;
     use gpui::{div, prelude::*, Render, TestAppContext};
+    use gpui::{EmptyView, View};
     use schemars::schema_for;
     use schemars::JsonSchema;
     use serde::{Deserialize, Serialize};
@@ -170,7 +170,7 @@ mod test {
         fn execute(
             &self,
             input: &Self::Input,
-            _cx: &gpui::AppContext,
+            _cx: &mut WindowContext,
         ) -> Task<Result<Self::Output>> {
             let _location = input.location.clone();
             let _unit = input.unit.clone();
@@ -200,6 +200,7 @@ mod test {
     #[gpui::test]
     async fn test_openai_weather_example(cx: &mut TestAppContext) {
         cx.background_executor.run_until_parked();
+        let (_, cx) = cx.add_window_view(|_cx| EmptyView);
 
         let tool = WeatherTool {
             current_weather: WeatherResult {

crates/assistant_tooling/src/tool.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use gpui::{AnyElement, AnyView, AppContext, IntoElement as _, Render, Task, View, WindowContext};
+use gpui::{AnyElement, AnyView, IntoElement as _, Render, Task, View, WindowContext};
 use schemars::{schema::RootSchema, schema_for, JsonSchema};
 use serde::Deserialize;
 use std::fmt::Display;
@@ -94,7 +94,7 @@ pub trait LanguageModelTool {
     }
 
     /// Executes the tool with the given input.
-    fn execute(&self, input: &Self::Input, cx: &AppContext) -> Task<Result<Self::Output>>;
+    fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>>;
 
     fn format(input: &Self::Input, output: &Result<Self::Output>) -> String;
 

crates/gpui/src/app/test_context.rs

@@ -218,7 +218,7 @@ impl TestAppContext {
     /// Adds a new window, and returns its root view and a `VisualTestContext` which can be used
     /// as a `WindowContext` for the rest of the test. Typically you would shadow this context with
     /// the returned one. `let (view, cx) = cx.add_window_view(...);`
-    pub fn add_window_view<F, V>(&mut self, build_window: F) -> (View<V>, &mut VisualTestContext)
+    pub fn add_window_view<F, V>(&mut self, build_root_view: F) -> (View<V>, &mut VisualTestContext)
     where
         F: FnOnce(&mut ViewContext<V>) -> V,
         V: 'static + Render,
@@ -230,7 +230,7 @@ impl TestAppContext {
                 bounds: Some(bounds),
                 ..Default::default()
             },
-            |cx| cx.new_view(build_window),
+            |cx| cx.new_view(build_root_view),
         );
         drop(cx);
         let view = window.root_view(self).unwrap();

crates/gpui/src/view.rs

@@ -1,3 +1,4 @@
+use crate::Empty;
 use crate::{
     seal::Sealed, AnyElement, AnyModel, AnyWeakModel, AppContext, Bounds, ContentMask, Element,
     ElementId, Entity, EntityId, Flatten, FocusHandle, FocusableView, GlobalElementId, IntoElement,
@@ -457,3 +458,12 @@ mod any_view {
         view.update(cx, |view, cx| view.render(cx).into_any_element())
     }
 }
+
+/// A view that renders nothing
+pub struct EmptyView;
+
+impl Render for EmptyView {
+    fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
+        Empty
+    }
+}