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>
This commit is contained in:
Marshall Bowers 2024-04-30 13:43:25 -04:00 committed by GitHub
parent ada2791fa3
commit d01428e69c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 173 additions and 27 deletions

View File

@ -122,7 +122,11 @@ impl LanguageModelTool for RollDiceTool {
"Rolls N many dice and returns the results.".to_string() "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) let rolls = (0..input.num_dice)
.map(|_| { .map(|_| {
let die_type = input.die_type.as_ref().unwrap_or(&Die::D6).clone(); 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() "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({ cx.spawn({
let fs = self.fs.clone(); let fs = self.fs.clone();
let root_dir = self.root_dir.clone(); let root_dir = self.root_dir.clone();

View File

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

View File

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

View File

@ -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")
}
}

View File

@ -1,6 +1,6 @@
use anyhow::Result; use anyhow::Result;
use assistant_tooling::LanguageModelTool; use assistant_tooling::LanguageModelTool;
use gpui::{prelude::*, AnyView, AppContext, Model, Task}; use gpui::{prelude::*, AnyView, Model, Task};
use project::Fs; use project::Fs;
use schemars::JsonSchema; use schemars::JsonSchema;
use semantic_index::{ProjectIndex, Status}; 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() "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 project_index = self.project_index.read(cx);
let status = project_index.status(); let status = project_index.status();

View File

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

View File

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

View File

@ -1,5 +1,5 @@
use anyhow::Result; 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 schemars::{schema::RootSchema, schema_for, JsonSchema};
use serde::Deserialize; use serde::Deserialize;
use std::fmt::Display; use std::fmt::Display;
@ -94,7 +94,7 @@ pub trait LanguageModelTool {
} }
/// Executes the tool with the given input. /// 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; fn format(input: &Self::Input, output: &Result<Self::Output>) -> String;

View File

@ -218,7 +218,7 @@ impl TestAppContext {
/// Adds a new window, and returns its root view and a `VisualTestContext` which can be used /// 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 /// 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(...);` /// 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 where
F: FnOnce(&mut ViewContext<V>) -> V, F: FnOnce(&mut ViewContext<V>) -> V,
V: 'static + Render, V: 'static + Render,
@ -230,7 +230,7 @@ impl TestAppContext {
bounds: Some(bounds), bounds: Some(bounds),
..Default::default() ..Default::default()
}, },
|cx| cx.new_view(build_window), |cx| cx.new_view(build_root_view),
); );
drop(cx); drop(cx);
let view = window.root_view(self).unwrap(); let view = window.root_view(self).unwrap();

View File

@ -1,3 +1,4 @@
use crate::Empty;
use crate::{ use crate::{
seal::Sealed, AnyElement, AnyModel, AnyWeakModel, AppContext, Bounds, ContentMask, Element, seal::Sealed, AnyElement, AnyModel, AnyWeakModel, AppContext, Bounds, ContentMask, Element,
ElementId, Entity, EntityId, Flatten, FocusHandle, FocusableView, GlobalElementId, IntoElement, 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()) 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
}
}