From e3de440715c57c5be4e9ed0b5de284609acdd309 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Tue, 30 Apr 2024 11:12:44 -0400 Subject: [PATCH] assistant2: Restructure tools in preparation for adding more (#11213) This PR does a slight restructuring of how tools are defined in the `assistant2` crate to make it more amenable to adding more tools in the near future. Release Notes: - N/A --- crates/assistant2/src/assistant2.rs | 4 +- crates/assistant2/src/tools.rs | 268 +------------------ crates/assistant2/src/tools/project_index.rs | 267 ++++++++++++++++++ 3 files changed, 271 insertions(+), 268 deletions(-) create mode 100644 crates/assistant2/src/tools/project_index.rs diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs index 3ee03f6d2917d251a588aa693c6bce8794f1b2d3..a4f95b6f3aecd6fc6e5adb139ecc61d91b1240c2 100644 --- a/crates/assistant2/src/assistant2.rs +++ b/crates/assistant2/src/assistant2.rs @@ -1,6 +1,6 @@ mod assistant_settings; mod completion_provider; -pub mod tools; +mod tools; mod ui; use ::ui::{div, prelude::*, Color, ViewContext}; @@ -23,7 +23,6 @@ use semantic_index::{CloudEmbeddingProvider, SemanticIndex}; use serde::Deserialize; use settings::Settings; use std::sync::Arc; -use tools::ProjectIndexTool; use ui::Composer; use util::{paths::EMBEDDINGS_DIR, ResultExt}; use workspace::{ @@ -33,6 +32,7 @@ use workspace::{ pub use assistant_settings::AssistantSettings; +use crate::tools::ProjectIndexTool; use crate::ui::UserOrAssistant; const MAX_COMPLETION_CALLS_PER_SUBMISSION: usize = 5; diff --git a/crates/assistant2/src/tools.rs b/crates/assistant2/src/tools.rs index 6ffddfe51d460da016e0cfa518727a0a219bd48f..46578ff34edead9a9c40f945918acb16332dc56b 100644 --- a/crates/assistant2/src/tools.rs +++ b/crates/assistant2/src/tools.rs @@ -1,267 +1,3 @@ -use anyhow::Result; -use assistant_tooling::LanguageModelTool; -use gpui::{prelude::*, AnyView, AppContext, Model, Task}; -use project::Fs; -use schemars::JsonSchema; -use semantic_index::{ProjectIndex, Status}; -use serde::Deserialize; -use std::sync::Arc; -use ui::{ - div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString, - WindowContext, -}; -use util::ResultExt as _; +mod project_index; -const DEFAULT_SEARCH_LIMIT: usize = 20; - -#[derive(Clone)] -pub struct CodebaseExcerpt { - path: SharedString, - text: SharedString, - score: f32, - element_id: ElementId, - expanded: bool, -} - -// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model. -// Any changes or deletions to the `CodebaseQuery` comments will change model behavior. - -#[derive(Deserialize, JsonSchema)] -pub struct CodebaseQuery { - /// Semantic search query - query: String, - /// Maximum number of results to return, defaults to 20 - limit: Option, -} - -pub struct ProjectIndexView { - input: CodebaseQuery, - output: Result, -} - -impl ProjectIndexView { - fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext) { - if let Ok(output) = &mut self.output { - if let Some(excerpt) = output - .excerpts - .iter_mut() - .find(|excerpt| excerpt.element_id == element_id) - { - excerpt.expanded = !excerpt.expanded; - cx.notify(); - } - } - } -} - -impl Render for ProjectIndexView { - fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { - let query = self.input.query.clone(); - - let result = &self.output; - - let output = match result { - Err(err) => { - return div().child(Label::new(format!("Error: {}", err)).color(Color::Error)); - } - Ok(output) => output, - }; - - div() - .v_flex() - .gap_2() - .child( - div() - .p_2() - .rounded_md() - .bg(cx.theme().colors().editor_background) - .child( - h_flex() - .child(Label::new("Query: ").color(Color::Modified)) - .child(Label::new(query).color(Color::Muted)), - ), - ) - .children(output.excerpts.iter().map(|excerpt| { - let element_id = excerpt.element_id.clone(); - let expanded = excerpt.expanded; - - CollapsibleContainer::new(element_id.clone(), expanded) - .start_slot( - h_flex() - .gap_1() - .child(Icon::new(IconName::File).color(Color::Muted)) - .child(Label::new(excerpt.path.clone()).color(Color::Muted)), - ) - .on_click(cx.listener(move |this, _, cx| { - this.toggle_expanded(element_id.clone(), cx); - })) - .child( - div() - .p_2() - .rounded_md() - .bg(cx.theme().colors().editor_background) - .child(excerpt.text.clone()), - ) - })) - } -} - -pub struct ProjectIndexTool { - project_index: Model, - fs: Arc, -} - -pub struct ProjectIndexOutput { - excerpts: Vec, - status: Status, -} - -impl ProjectIndexTool { - pub fn new(project_index: Model, fs: Arc) -> Self { - // Listen for project index status and update the ProjectIndexTool directly - - // TODO: setup a better description based on the user's current codebase. - Self { project_index, fs } - } -} - -impl LanguageModelTool for ProjectIndexTool { - type Input = CodebaseQuery; - type Output = ProjectIndexOutput; - type View = ProjectIndexView; - - fn name(&self) -> String { - "query_codebase".to_string() - } - - fn description(&self) -> 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> { - let project_index = self.project_index.read(cx); - - let status = project_index.status(); - let results = project_index.search( - query.query.as_str(), - query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT), - cx, - ); - - let fs = self.fs.clone(); - - cx.spawn(|cx| async move { - let results = results.await; - - let excerpts = results.into_iter().map(|result| { - let abs_path = result - .worktree - .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path)); - let fs = fs.clone(); - - async move { - let path = result.path.clone(); - let text = fs.load(&abs_path?).await?; - - let mut start = result.range.start; - let mut end = result.range.end.min(text.len()); - while !text.is_char_boundary(start) { - start += 1; - } - while !text.is_char_boundary(end) { - end -= 1; - } - - anyhow::Ok(CodebaseExcerpt { - element_id: ElementId::Name(nanoid::nanoid!().into()), - expanded: false, - path: path.to_string_lossy().to_string().into(), - text: SharedString::from(text[start..end].to_string()), - score: result.score, - }) - } - }); - - let excerpts = futures::future::join_all(excerpts) - .await - .into_iter() - .filter_map(|result| result.log_err()) - .collect(); - anyhow::Ok(ProjectIndexOutput { excerpts, status }) - }) - } - - fn output_view( - _tool_call_id: String, - input: Self::Input, - output: Result, - cx: &mut WindowContext, - ) -> gpui::View { - cx.new_view(|_cx| ProjectIndexView { input, output }) - } - - fn status_view(&self, cx: &mut WindowContext) -> Option { - Some( - cx.new_view(|cx| ProjectIndexStatusView::new(self.project_index.clone(), cx)) - .into(), - ) - } - - fn format(_input: &Self::Input, output: &Result) -> String { - match &output { - Ok(output) => { - let mut body = "Semantic search results:\n".to_string(); - - if output.status != Status::Idle { - body.push_str("Still indexing. Results may be incomplete.\n"); - } - - if output.excerpts.is_empty() { - body.push_str("No results found"); - return body; - } - - for excerpt in &output.excerpts { - body.push_str("Excerpt from "); - body.push_str(excerpt.path.as_ref()); - body.push_str(", score "); - body.push_str(&excerpt.score.to_string()); - body.push_str(":\n"); - body.push_str("~~~\n"); - body.push_str(excerpt.text.as_ref()); - body.push_str("~~~\n"); - } - body - } - Err(err) => format!("Error: {}", err), - } - } -} - -struct ProjectIndexStatusView { - project_index: Model, -} - -impl ProjectIndexStatusView { - pub fn new(project_index: Model, cx: &mut ViewContext) -> Self { - cx.subscribe(&project_index, |_this, _, _status: &Status, cx| { - cx.notify(); - }) - .detach(); - Self { project_index } - } -} - -impl Render for ProjectIndexStatusView { - fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { - let status = self.project_index.read(cx).status(); - - h_flex().gap_2().map(|element| match status { - Status::Idle => element.child(Label::new("Project index ready")), - Status::Loading => element.child(Label::new("Project index loading...")), - Status::Scanning { remaining_count } => element.child(Label::new(format!( - "Project index scanning: {remaining_count} remaining..." - ))), - }) - } -} +pub use project_index::*; diff --git a/crates/assistant2/src/tools/project_index.rs b/crates/assistant2/src/tools/project_index.rs new file mode 100644 index 0000000000000000000000000000000000000000..6ffddfe51d460da016e0cfa518727a0a219bd48f --- /dev/null +++ b/crates/assistant2/src/tools/project_index.rs @@ -0,0 +1,267 @@ +use anyhow::Result; +use assistant_tooling::LanguageModelTool; +use gpui::{prelude::*, AnyView, AppContext, Model, Task}; +use project::Fs; +use schemars::JsonSchema; +use semantic_index::{ProjectIndex, Status}; +use serde::Deserialize; +use std::sync::Arc; +use ui::{ + div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString, + WindowContext, +}; +use util::ResultExt as _; + +const DEFAULT_SEARCH_LIMIT: usize = 20; + +#[derive(Clone)] +pub struct CodebaseExcerpt { + path: SharedString, + text: SharedString, + score: f32, + element_id: ElementId, + expanded: bool, +} + +// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model. +// Any changes or deletions to the `CodebaseQuery` comments will change model behavior. + +#[derive(Deserialize, JsonSchema)] +pub struct CodebaseQuery { + /// Semantic search query + query: String, + /// Maximum number of results to return, defaults to 20 + limit: Option, +} + +pub struct ProjectIndexView { + input: CodebaseQuery, + output: Result, +} + +impl ProjectIndexView { + fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext) { + if let Ok(output) = &mut self.output { + if let Some(excerpt) = output + .excerpts + .iter_mut() + .find(|excerpt| excerpt.element_id == element_id) + { + excerpt.expanded = !excerpt.expanded; + cx.notify(); + } + } + } +} + +impl Render for ProjectIndexView { + fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { + let query = self.input.query.clone(); + + let result = &self.output; + + let output = match result { + Err(err) => { + return div().child(Label::new(format!("Error: {}", err)).color(Color::Error)); + } + Ok(output) => output, + }; + + div() + .v_flex() + .gap_2() + .child( + div() + .p_2() + .rounded_md() + .bg(cx.theme().colors().editor_background) + .child( + h_flex() + .child(Label::new("Query: ").color(Color::Modified)) + .child(Label::new(query).color(Color::Muted)), + ), + ) + .children(output.excerpts.iter().map(|excerpt| { + let element_id = excerpt.element_id.clone(); + let expanded = excerpt.expanded; + + CollapsibleContainer::new(element_id.clone(), expanded) + .start_slot( + h_flex() + .gap_1() + .child(Icon::new(IconName::File).color(Color::Muted)) + .child(Label::new(excerpt.path.clone()).color(Color::Muted)), + ) + .on_click(cx.listener(move |this, _, cx| { + this.toggle_expanded(element_id.clone(), cx); + })) + .child( + div() + .p_2() + .rounded_md() + .bg(cx.theme().colors().editor_background) + .child(excerpt.text.clone()), + ) + })) + } +} + +pub struct ProjectIndexTool { + project_index: Model, + fs: Arc, +} + +pub struct ProjectIndexOutput { + excerpts: Vec, + status: Status, +} + +impl ProjectIndexTool { + pub fn new(project_index: Model, fs: Arc) -> Self { + // Listen for project index status and update the ProjectIndexTool directly + + // TODO: setup a better description based on the user's current codebase. + Self { project_index, fs } + } +} + +impl LanguageModelTool for ProjectIndexTool { + type Input = CodebaseQuery; + type Output = ProjectIndexOutput; + type View = ProjectIndexView; + + fn name(&self) -> String { + "query_codebase".to_string() + } + + fn description(&self) -> 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> { + let project_index = self.project_index.read(cx); + + let status = project_index.status(); + let results = project_index.search( + query.query.as_str(), + query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT), + cx, + ); + + let fs = self.fs.clone(); + + cx.spawn(|cx| async move { + let results = results.await; + + let excerpts = results.into_iter().map(|result| { + let abs_path = result + .worktree + .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path)); + let fs = fs.clone(); + + async move { + let path = result.path.clone(); + let text = fs.load(&abs_path?).await?; + + let mut start = result.range.start; + let mut end = result.range.end.min(text.len()); + while !text.is_char_boundary(start) { + start += 1; + } + while !text.is_char_boundary(end) { + end -= 1; + } + + anyhow::Ok(CodebaseExcerpt { + element_id: ElementId::Name(nanoid::nanoid!().into()), + expanded: false, + path: path.to_string_lossy().to_string().into(), + text: SharedString::from(text[start..end].to_string()), + score: result.score, + }) + } + }); + + let excerpts = futures::future::join_all(excerpts) + .await + .into_iter() + .filter_map(|result| result.log_err()) + .collect(); + anyhow::Ok(ProjectIndexOutput { excerpts, status }) + }) + } + + fn output_view( + _tool_call_id: String, + input: Self::Input, + output: Result, + cx: &mut WindowContext, + ) -> gpui::View { + cx.new_view(|_cx| ProjectIndexView { input, output }) + } + + fn status_view(&self, cx: &mut WindowContext) -> Option { + Some( + cx.new_view(|cx| ProjectIndexStatusView::new(self.project_index.clone(), cx)) + .into(), + ) + } + + fn format(_input: &Self::Input, output: &Result) -> String { + match &output { + Ok(output) => { + let mut body = "Semantic search results:\n".to_string(); + + if output.status != Status::Idle { + body.push_str("Still indexing. Results may be incomplete.\n"); + } + + if output.excerpts.is_empty() { + body.push_str("No results found"); + return body; + } + + for excerpt in &output.excerpts { + body.push_str("Excerpt from "); + body.push_str(excerpt.path.as_ref()); + body.push_str(", score "); + body.push_str(&excerpt.score.to_string()); + body.push_str(":\n"); + body.push_str("~~~\n"); + body.push_str(excerpt.text.as_ref()); + body.push_str("~~~\n"); + } + body + } + Err(err) => format!("Error: {}", err), + } + } +} + +struct ProjectIndexStatusView { + project_index: Model, +} + +impl ProjectIndexStatusView { + pub fn new(project_index: Model, cx: &mut ViewContext) -> Self { + cx.subscribe(&project_index, |_this, _, _status: &Status, cx| { + cx.notify(); + }) + .detach(); + Self { project_index } + } +} + +impl Render for ProjectIndexStatusView { + fn render(&mut self, cx: &mut ViewContext) -> impl IntoElement { + let status = self.project_index.read(cx).status(); + + h_flex().gap_2().map(|element| match status { + Status::Idle => element.child(Label::new("Project index ready")), + Status::Loading => element.child(Label::new("Project index loading...")), + Status::Scanning { remaining_count } => element.child(Label::new(format!( + "Project index scanning: {remaining_count} remaining..." + ))), + }) + } +}