From d01428e69c411482da1ad95e8f20716b626d1ec8 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 30 Apr 2024 13:43:25 -0400 Subject: [PATCH] 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 --- .../examples/chat_with_functions.rs | 12 +- crates/assistant2/src/assistant2.rs | 24 +++- crates/assistant2/src/tools.rs | 2 + crates/assistant2/src/tools/create_buffer.rs | 111 ++++++++++++++++++ crates/assistant2/src/tools/project_index.rs | 4 +- crates/assistant2/src/ui/chat_message.rs | 24 ++-- crates/assistant_tooling/src/registry.rs | 5 +- crates/assistant_tooling/src/tool.rs | 4 +- crates/gpui/src/app/test_context.rs | 4 +- crates/gpui/src/view.rs | 10 ++ 10 files changed, 173 insertions(+), 27 deletions(-) create mode 100644 crates/assistant2/src/tools/create_buffer.rs diff --git a/crates/assistant2/examples/chat_with_functions.rs b/crates/assistant2/examples/chat_with_functions.rs index 58207741bf..7c7011caa3 100644 --- a/crates/assistant2/examples/chat_with_functions.rs +++ b/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> { + fn execute( + &self, + input: &Self::Input, + _cx: &mut WindowContext, + ) -> Task> { 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> { + fn execute( + &self, + input: &Self::Input, + cx: &mut WindowContext, + ) -> Task> { cx.spawn({ let fs = self.fs.clone(); let root_dir = self.root_dir.clone(); diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs index a4f95b6f3a..54ffc686b1 100644 --- a/crates/assistant2/src/assistant2.rs +++ b/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; diff --git a/crates/assistant2/src/tools.rs b/crates/assistant2/src/tools.rs index 46578ff34e..d1a676e74e 100644 --- a/crates/assistant2/src/tools.rs +++ b/crates/assistant2/src/tools.rs @@ -1,3 +1,5 @@ +mod create_buffer; mod project_index; +pub use create_buffer::*; pub use project_index::*; diff --git a/crates/assistant2/src/tools/create_buffer.rs b/crates/assistant2/src/tools/create_buffer.rs new file mode 100644 index 0000000000..ddc66ba15a --- /dev/null +++ b/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, + project: Model, +} + +impl CreateBufferTool { + pub fn new(workspace: WeakView, project: Model) -> 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> { + 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) -> 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, + cx: &mut WindowContext, + ) -> View { + cx.new_view(|_cx| CreateBufferView {}) + } +} + +pub struct CreateBufferView {} + +impl Render for CreateBufferView { + fn render(&mut self, _cx: &mut ViewContext) -> impl IntoElement { + div().child("Opening a buffer") + } +} diff --git a/crates/assistant2/src/tools/project_index.rs b/crates/assistant2/src/tools/project_index.rs index 6ffddfe51d..20f6c51add 100644 --- a/crates/assistant2/src/tools/project_index.rs +++ b/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> { + fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task> { let project_index = self.project_index.read(cx); let status = project_index.status(); diff --git a/crates/assistant2/src/ui/chat_message.rs b/crates/assistant2/src/ui/chat_message.rs index d51f62aba5..1a2ded6582 100644 --- a/crates/assistant2/src/ui/chat_message.rs +++ b/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, collapsed: bool, on_collapse_handle_click: Box, } @@ -24,7 +24,7 @@ impl ChatMessage { pub fn new( id: MessageId, player: UserOrAssistant, - message: AnyElement, + message: Option, collapsed: bool, on_collapse_handle_click: Box, ) -> 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)) } } diff --git a/crates/assistant_tooling/src/registry.rs b/crates/assistant_tooling/src/registry.rs index 136f012d33..da9fe94b9e 100644 --- a/crates/assistant_tooling/src/registry.rs +++ b/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> { 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 { diff --git a/crates/assistant_tooling/src/tool.rs b/crates/assistant_tooling/src/tool.rs index 3b728e4ba7..82536b9e8a 100644 --- a/crates/assistant_tooling/src/tool.rs +++ b/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>; + fn execute(&self, input: &Self::Input, cx: &mut WindowContext) -> Task>; fn format(input: &Self::Input, output: &Result) -> String; diff --git a/crates/gpui/src/app/test_context.rs b/crates/gpui/src/app/test_context.rs index a17f8defe9..df622b5382 100644 --- a/crates/gpui/src/app/test_context.rs +++ b/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(&mut self, build_window: F) -> (View, &mut VisualTestContext) + pub fn add_window_view(&mut self, build_root_view: F) -> (View, &mut VisualTestContext) where F: FnOnce(&mut ViewContext) -> 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(); diff --git a/crates/gpui/src/view.rs b/crates/gpui/src/view.rs index d67f335fbb..f6e859ce44 100644 --- a/crates/gpui/src/view.rs +++ b/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) -> impl IntoElement { + Empty + } +}