1use anyhow::Result;
2use assistant_tooling::{LanguageModelTool, ToolOutput};
3use collections::BTreeMap;
4use gpui::{prelude::*, Model, Task};
5use project::ProjectPath;
6use schemars::JsonSchema;
7use semantic_index::{ProjectIndex, Status};
8use serde::Deserialize;
9use std::{fmt::Write as _, ops::Range};
10use ui::{div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
11
12const DEFAULT_SEARCH_LIMIT: usize = 20;
13
14pub struct ProjectIndexTool {
15 project_index: Model<ProjectIndex>,
16}
17
18// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model.
19// Any changes or deletions to the `CodebaseQuery` comments will change model behavior.
20
21#[derive(Deserialize, JsonSchema)]
22pub struct CodebaseQuery {
23 /// Semantic search query
24 query: String,
25 /// Maximum number of results to return, defaults to 20
26 limit: Option<usize>,
27}
28
29pub struct ProjectIndexView {
30 input: CodebaseQuery,
31 output: Result<ProjectIndexOutput>,
32 element_id: ElementId,
33 expanded_header: bool,
34}
35
36pub struct ProjectIndexOutput {
37 status: Status,
38 excerpts: BTreeMap<ProjectPath, Vec<Range<usize>>>,
39}
40
41impl ProjectIndexView {
42 fn new(input: CodebaseQuery, output: Result<ProjectIndexOutput>) -> Self {
43 let element_id = ElementId::Name(nanoid::nanoid!().into());
44
45 Self {
46 input,
47 output,
48 element_id,
49 expanded_header: false,
50 }
51 }
52
53 fn toggle_header(&mut self, cx: &mut ViewContext<Self>) {
54 self.expanded_header = !self.expanded_header;
55 cx.notify();
56 }
57}
58
59impl Render for ProjectIndexView {
60 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
61 let query = self.input.query.clone();
62
63 let result = &self.output;
64
65 let output = match result {
66 Err(err) => {
67 return div().child(Label::new(format!("Error: {}", err)).color(Color::Error));
68 }
69 Ok(output) => output,
70 };
71
72 let file_count = output.excerpts.len();
73
74 let header = h_flex()
75 .gap_2()
76 .child(Icon::new(IconName::File))
77 .child(format!(
78 "Read {} {}",
79 file_count,
80 if file_count == 1 { "file" } else { "files" }
81 ));
82
83 v_flex().gap_3().child(
84 CollapsibleContainer::new(self.element_id.clone(), self.expanded_header)
85 .start_slot(header)
86 .on_click(cx.listener(move |this, _, cx| {
87 this.toggle_header(cx);
88 }))
89 .child(
90 v_flex()
91 .gap_3()
92 .p_3()
93 .child(
94 h_flex()
95 .gap_2()
96 .child(Icon::new(IconName::MagnifyingGlass))
97 .child(Label::new(format!("`{}`", query)).color(Color::Muted)),
98 )
99 .child(
100 v_flex()
101 .gap_2()
102 .children(output.excerpts.keys().map(|path| {
103 h_flex().gap_2().child(Icon::new(IconName::File)).child(
104 Label::new(path.path.to_string_lossy().to_string())
105 .color(Color::Muted),
106 )
107 })),
108 ),
109 ),
110 )
111 }
112}
113
114impl ToolOutput for ProjectIndexView {
115 fn generate(
116 &self,
117 context: &mut assistant_tooling::ProjectContext,
118 _: &mut WindowContext,
119 ) -> String {
120 match &self.output {
121 Ok(output) => {
122 let mut body = "found results in the following paths:\n".to_string();
123
124 for (project_path, ranges) in &output.excerpts {
125 context.add_excerpts(project_path.clone(), ranges);
126 writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
127 }
128
129 if output.status != Status::Idle {
130 body.push_str("Still indexing. Results may be incomplete.\n");
131 }
132
133 body
134 }
135 Err(err) => format!("Error: {}", err),
136 }
137 }
138}
139
140impl ProjectIndexTool {
141 pub fn new(project_index: Model<ProjectIndex>) -> Self {
142 Self { project_index }
143 }
144}
145
146impl LanguageModelTool for ProjectIndexTool {
147 type Input = CodebaseQuery;
148 type Output = ProjectIndexOutput;
149 type View = ProjectIndexView;
150
151 fn name(&self) -> String {
152 "query_codebase".to_string()
153 }
154
155 fn description(&self) -> String {
156 "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()
157 }
158
159 fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
160 let project_index = self.project_index.read(cx);
161 let status = project_index.status();
162 let search = project_index.search(
163 query.query.clone(),
164 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
165 cx,
166 );
167
168 cx.spawn(|mut cx| async move {
169 let search_results = search.await?;
170
171 cx.update(|cx| {
172 let mut output = ProjectIndexOutput {
173 status,
174 excerpts: Default::default(),
175 };
176
177 for search_result in search_results {
178 let path = ProjectPath {
179 worktree_id: search_result.worktree.read(cx).id(),
180 path: search_result.path.clone(),
181 };
182
183 let excerpts_for_path = output.excerpts.entry(path).or_default();
184 let ix = match excerpts_for_path
185 .binary_search_by_key(&search_result.range.start, |r| r.start)
186 {
187 Ok(ix) | Err(ix) => ix,
188 };
189 excerpts_for_path.insert(ix, search_result.range);
190 }
191
192 output
193 })
194 })
195 }
196
197 fn output_view(
198 input: Self::Input,
199 output: Result<Self::Output>,
200 cx: &mut WindowContext,
201 ) -> gpui::View<Self::View> {
202 cx.new_view(|_cx| ProjectIndexView::new(input, output))
203 }
204
205 fn render_running(_: &mut WindowContext) -> impl IntoElement {
206 CollapsibleContainer::new(ElementId::Name(nanoid::nanoid!().into()), false)
207 .start_slot("Searching code base")
208 }
209}