Trim index output (#11445)

Trims down the project index output view in assistant2 to just list the
filenames and hideaway the query.

<img width="374" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/8603e3cf-9fdc-4661-bc45-1d87621a006f">

Introduces a way for tools to render running.

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>

crates/assistant2/src/assistant2.rs

@@ -724,21 +724,7 @@ impl AssistantChat {
 
                 let tools = tool_calls
                     .iter()
-                    .map(|tool_call| {
+                    .map(|tool_call| self.tool_registry.render_tool_call(tool_call, cx))
-                        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(),
-                        }
-                    })
                     .collect::<Vec<AnyElement>>();
 
                 let tools_body = if tools.is_empty() {

crates/assistant2/src/tools/project_index.rs

@@ -1,14 +1,12 @@
 use anyhow::Result;
-use assistant_tooling::{
+use assistant_tooling::LanguageModelTool;
-    // assistant_tool_button::{AssistantToolButton, ToolStatus},
-    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<ProjectIndexOutput>,
+    element_id: ElementId,
+    expanded_header: bool,
 }
 
 impl ProjectIndexView {
-    fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext<Self>) {
+    fn new(input: CodebaseQuery, output: Result<ProjectIndexOutput>) -> Self {
-        if let Ok(output) = &mut self.output {
+        let element_id = ElementId::Name(nanoid::nanoid!().into());
-            if let Some(excerpt) = output
+
-                .excerpts
+        Self {
-                .iter_mut()
+            input,
-                .find(|excerpt| excerpt.element_id == element_id)
+            output,
-            {
+            element_id,
-                excerpt.expanded = !excerpt.expanded;
+            expanded_header: false,
-                cx.notify();
-            }
         }
     }
+
+    fn toggle_header(&mut self, cx: &mut ViewContext<Self>) {
+        self.expanded_header = !self.expanded_header;
+        cx.notify();
+    }
 }
 
 impl Render for ProjectIndexView {
@@ -70,42 +71,47 @@ impl Render for ProjectIndexView {
             Ok(output) => output,
         };
 
-        div()
+        let num_files_searched = output.files_searched.len();
-            .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)
+        let header = h_flex()
-                    .start_slot(
+            .gap_2()
-                        h_flex()
+            .child(Icon::new(IconName::File))
-                            .gap_1()
+            .child(format!(
-                            .child(Icon::new(IconName::File).color(Color::Muted))
+                "Read {} {}",
-                            .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
+                num_files_searched,
-                    )
+                if num_files_searched == 1 {
-                    .on_click(cx.listener(move |this, _, cx| {
+                    "file"
-                        this.toggle_expanded(element_id.clone(), cx);
+                } else {
-                    }))
+                    "files"
-                    .child(
+                }
-                        div()
+            ));
-                            .p_2()
+
-                            .rounded_md()
+        v_flex().gap_3().child(
-                            .bg(cx.theme().colors().editor_background)
+            CollapsibleContainer::new(self.element_id.clone(), self.expanded_header)
-                            .child(excerpt.text.clone()),
+                .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<CodebaseExcerpt>,
     status: Status,
+    files_searched: HashSet<SharedString>,
 }
 
 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<Result<Self::Output>> {
@@ -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();
+                .inspect(|excerpt| {
-            anyhow::Ok(ProjectIndexOutput { excerpts, status })
+                    files_searched.insert(excerpt.path.clone());
+                })
+                .collect::<Vec<_>>();
+
+            anyhow::Ok(ProjectIndexOutput {
+                excerpts,
+                status,
+                files_searched,
+            })
         })
     }
 
@@ -199,7 +213,12 @@ impl LanguageModelTool for ProjectIndexTool {
         output: Result<Self::Output>,
         cx: &mut WindowContext,
     ) -> gpui::View<Self::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<Self::Output>) -> String {

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<dyn Fn(&ToolFunctionCall, &mut WindowContext) -> Task<Result<ToolFunctionCall>>>,
+    render_running: Box<dyn Fn(&mut WindowContext) -> gpui::AnyElement>,
     definition: ToolFunctionDefinition,
 }
 
@@ -22,12 +23,14 @@ impl Tool {
     fn new(
         type_id: TypeId,
         call: Box<dyn Fn(&ToolFunctionCall, &mut WindowContext) -> Task<Result<ToolFunctionCall>>>,
+        render_running: Box<dyn Fn(&mut WindowContext) -> 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<T: 'static + LanguageModelTool>(
         &mut self,
         tool: T,
@@ -115,6 +136,7 @@ impl ToolRegistry {
                     })
                 },
             ),
+            Box::new(|cx| T::render_running(cx).into_any_element()),
             definition,
         );
 

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<Self::Output>,
         cx: &mut WindowContext,
     ) -> View<Self::View>;
+
+    fn render_running(_cx: &mut WindowContext) -> impl IntoElement {
+        div()
+    }
 }