From 32b59bfa0e09ab249fcd2b75d0521d9e7ad76e37 Mon Sep 17 00:00:00 2001 From: Kyle Kelley Date: Mon, 6 May 2024 10:37:31 -0700 Subject: [PATCH] Trim index output (#11445) Trims down the project index output view in assistant2 to just list the filenames and hideaway the query. image Introduces a way for tools to render running. Release Notes: - N/A --------- Co-authored-by: Max --- crates/assistant2/src/assistant2.rs | 16 +-- crates/assistant2/src/tools/project_index.rs | 135 +++++++++++-------- crates/assistant_tooling/src/registry.rs | 24 +++- crates/assistant_tooling/src/tool.rs | 6 +- 4 files changed, 106 insertions(+), 75 deletions(-) diff --git a/crates/assistant2/src/assistant2.rs b/crates/assistant2/src/assistant2.rs index 9c208dc037..7ca8e3b57f 100644 --- a/crates/assistant2/src/assistant2.rs +++ b/crates/assistant2/src/assistant2.rs @@ -724,21 +724,7 @@ impl AssistantChat { let tools = tool_calls .iter() - .map(|tool_call| { - let result = &tool_call.result; - let name = tool_call.name.clone(); - match result { - Some(result) => div() - .p_2() - .child(result.into_any_element(&name)) - .into_any_element(), - None => div() - .p_2() - .child(Label::new(name).color(Color::Modified)) - .child("Running...") - .into_any_element(), - } - }) + .map(|tool_call| self.tool_registry.render_tool_call(tool_call, cx)) .collect::>(); let tools_body = if tools.is_empty() { diff --git a/crates/assistant2/src/tools/project_index.rs b/crates/assistant2/src/tools/project_index.rs index b6b8258964..7ccbae79d3 100644 --- a/crates/assistant2/src/tools/project_index.rs +++ b/crates/assistant2/src/tools/project_index.rs @@ -1,14 +1,12 @@ use anyhow::Result; -use assistant_tooling::{ - // assistant_tool_button::{AssistantToolButton, ToolStatus}, - LanguageModelTool, -}; +use assistant_tooling::LanguageModelTool; use gpui::{prelude::*, Model, Task}; use project::Fs; use schemars::JsonSchema; use semantic_index::{ProjectIndex, Status}; use serde::Deserialize; -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; + use ui::{ div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString, WindowContext, @@ -22,8 +20,6 @@ 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. @@ -40,21 +36,26 @@ pub struct CodebaseQuery { pub struct ProjectIndexView { input: CodebaseQuery, output: Result, + element_id: ElementId, + expanded_header: bool, } 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(); - } + fn new(input: CodebaseQuery, output: Result) -> Self { + let element_id = ElementId::Name(nanoid::nanoid!().into()); + + Self { + input, + output, + element_id, + expanded_header: false, } } + + fn toggle_header(&mut self, cx: &mut ViewContext) { + self.expanded_header = !self.expanded_header; + cx.notify(); + } } impl Render for ProjectIndexView { @@ -70,42 +71,47 @@ impl Render for ProjectIndexView { 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; + let num_files_searched = output.files_searched.len(); - 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()), - ) - })) + let header = h_flex() + .gap_2() + .child(Icon::new(IconName::File)) + .child(format!( + "Read {} {}", + num_files_searched, + if num_files_searched == 1 { + "file" + } else { + "files" + } + )); + + v_flex().gap_3().child( + CollapsibleContainer::new(self.element_id.clone(), self.expanded_header) + .start_slot(header) + .on_click(cx.listener(move |this, _, cx| { + this.toggle_header(cx); + })) + .child( + v_flex() + .gap_3() + .p_3() + .child( + h_flex() + .gap_2() + .child(Icon::new(IconName::MagnifyingGlass)) + .child(Label::new(format!("`{}`", query)).color(Color::Muted)), + ) + .child(v_flex().gap_2().children(output.files_searched.iter().map( + |path| { + h_flex() + .gap_2() + .child(Icon::new(IconName::File)) + .child(Label::new(path.clone()).color(Color::Muted)) + }, + ))), + ), + ) } } @@ -117,6 +123,7 @@ pub struct ProjectIndexTool { pub struct ProjectIndexOutput { excerpts: Vec, status: Status, + files_searched: HashSet, } impl ProjectIndexTool { @@ -138,7 +145,7 @@ impl LanguageModelTool for ProjectIndexTool { } 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() + "Semantic search against the user's current codebase, returning excerpts related to the query by computing a dot product against embeddings of code chunks in the code base and an embedding of the query.".to_string() } fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task> { @@ -175,8 +182,6 @@ impl LanguageModelTool for ProjectIndexTool { } 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, @@ -184,12 +189,21 @@ impl LanguageModelTool for ProjectIndexTool { } }); + let mut files_searched = HashSet::new(); let excerpts = futures::future::join_all(excerpts) .await .into_iter() .filter_map(|result| result.log_err()) - .collect(); - anyhow::Ok(ProjectIndexOutput { excerpts, status }) + .inspect(|excerpt| { + files_searched.insert(excerpt.path.clone()); + }) + .collect::>(); + + anyhow::Ok(ProjectIndexOutput { + excerpts, + status, + files_searched, + }) }) } @@ -199,7 +213,12 @@ impl LanguageModelTool for ProjectIndexTool { output: Result, cx: &mut WindowContext, ) -> gpui::View { - cx.new_view(|_cx| ProjectIndexView { input, output }) + cx.new_view(|_cx| ProjectIndexView::new(input, output)) + } + + fn render_running(_: &mut WindowContext) -> impl IntoElement { + CollapsibleContainer::new(ElementId::Name(nanoid::nanoid!().into()), false) + .start_slot("Searching code base") } fn format(_input: &Self::Input, output: &Result) -> String { diff --git a/crates/assistant_tooling/src/registry.rs b/crates/assistant_tooling/src/registry.rs index b0e7dc4f2e..4c3c1a082d 100644 --- a/crates/assistant_tooling/src/registry.rs +++ b/crates/assistant_tooling/src/registry.rs @@ -1,5 +1,5 @@ use anyhow::{anyhow, Result}; -use gpui::{Task, WindowContext}; +use gpui::{div, AnyElement, IntoElement as _, ParentElement, Styled, Task, WindowContext}; use std::{ any::TypeId, collections::HashMap, @@ -15,6 +15,7 @@ pub struct Tool { enabled: AtomicBool, type_id: TypeId, call: Box Task>>, + render_running: Box gpui::AnyElement>, definition: ToolFunctionDefinition, } @@ -22,12 +23,14 @@ impl Tool { fn new( type_id: TypeId, call: Box Task>>, + render_running: Box gpui::AnyElement>, definition: ToolFunctionDefinition, ) -> Self { Self { enabled: AtomicBool::new(true), type_id, call, + render_running, definition, } } @@ -70,6 +73,24 @@ impl ToolRegistry { .collect() } + pub fn render_tool_call( + &self, + tool_call: &ToolFunctionCall, + cx: &mut WindowContext, + ) -> AnyElement { + match &tool_call.result { + Some(result) => div() + .p_2() + .child(result.into_any_element(&tool_call.name)) + .into_any_element(), + None => self + .tools + .get(&tool_call.name) + .map(|tool| (tool.render_running)(cx)) + .unwrap_or_else(|| div().into_any_element()), + } + } + pub fn register( &mut self, tool: T, @@ -115,6 +136,7 @@ impl ToolRegistry { }) }, ), + Box::new(|cx| T::render_running(cx).into_any_element()), definition, ); diff --git a/crates/assistant_tooling/src/tool.rs b/crates/assistant_tooling/src/tool.rs index 256d2eef8a..8bc55bac80 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, IntoElement as _, Render, Task, View, WindowContext}; +use gpui::{div, AnyElement, AnyView, IntoElement, Render, Task, View, WindowContext}; use schemars::{schema::RootSchema, schema_for, JsonSchema}; use serde::Deserialize; use std::fmt::Display; @@ -104,4 +104,8 @@ pub trait LanguageModelTool { output: Result, cx: &mut WindowContext, ) -> View; + + fn render_running(_cx: &mut WindowContext) -> impl IntoElement { + div() + } }