1use anyhow::Result;
2use assistant_tooling::{LanguageModelTool, ToolView};
3use collections::BTreeMap;
4use file_icons::FileIcons;
5use gpui::{prelude::*, AnyElement, Model, Task};
6use project::ProjectPath;
7use schemars::JsonSchema;
8use semantic_index::{ProjectIndex, Status};
9use serde::{Deserialize, Serialize};
10use std::{
11 fmt::Write as _,
12 ops::Range,
13 path::{Path, PathBuf},
14 str::FromStr as _,
15 sync::Arc,
16};
17use ui::{prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
18
19const DEFAULT_SEARCH_LIMIT: usize = 20;
20
21pub struct ProjectIndexTool {
22 project_index: Model<ProjectIndex>,
23}
24
25#[derive(Default)]
26enum ProjectIndexToolState {
27 #[default]
28 CollectingQuery,
29 Searching,
30 Error(anyhow::Error),
31 Finished {
32 excerpts: BTreeMap<ProjectPath, Vec<Range<usize>>>,
33 index_status: Status,
34 },
35}
36
37pub struct ProjectIndexView {
38 project_index: Model<ProjectIndex>,
39 input: CodebaseQuery,
40 expanded_header: bool,
41 state: ProjectIndexToolState,
42}
43
44#[derive(Default, Deserialize, JsonSchema)]
45pub struct CodebaseQuery {
46 /// Semantic search query
47 query: String,
48 /// Criteria to include results
49 includes: Option<SearchFilter>,
50 /// Criteria to exclude results
51 excludes: Option<SearchFilter>,
52}
53
54#[derive(Deserialize, JsonSchema, Clone, Default)]
55pub struct SearchFilter {
56 /// Filter by file path prefix
57 prefix_path: Option<String>,
58 /// Filter by file extension
59 extension: Option<String>,
60 // Note: we possibly can't do content filtering very easily given the project context handling
61 // the final results, so we're leaving out direct string matches for now
62}
63
64fn project_starts_with(prefix_path: Option<String>, project_path: ProjectPath) -> bool {
65 if let Some(path) = &prefix_path {
66 if let Some(path) = PathBuf::from_str(path).ok() {
67 return project_path.path.starts_with(path);
68 }
69 }
70
71 return false;
72}
73
74impl SearchFilter {
75 fn matches(&self, project_path: &ProjectPath) -> bool {
76 let path_match = project_starts_with(self.prefix_path.clone(), project_path.clone());
77
78 path_match
79 && (if let Some(extension) = &self.extension {
80 project_path
81 .path
82 .extension()
83 .and_then(|ext| ext.to_str())
84 .map(|ext| ext == extension)
85 .unwrap_or(false)
86 } else {
87 true
88 })
89 }
90}
91
92#[derive(Serialize, Deserialize)]
93pub struct SerializedState {
94 index_status: Status,
95 error_message: Option<String>,
96 worktrees: BTreeMap<Arc<Path>, WorktreeIndexOutput>,
97}
98
99#[derive(Default, Serialize, Deserialize)]
100struct WorktreeIndexOutput {
101 excerpts: BTreeMap<Arc<Path>, Vec<Range<usize>>>,
102}
103
104impl ProjectIndexView {
105 fn toggle_header(&mut self, cx: &mut ViewContext<Self>) {
106 self.expanded_header = !self.expanded_header;
107 cx.notify();
108 }
109
110 fn render_filter_section(
111 &mut self,
112 heading: &str,
113 filter: Option<SearchFilter>,
114 cx: &mut ViewContext<Self>,
115 ) -> Option<AnyElement> {
116 let filter = match filter {
117 Some(filter) => filter,
118 None => return None,
119 };
120
121 // Any of the filter fields can be empty. We'll show nothing if they're all empty.
122 let path = filter.prefix_path.as_ref().map(|path| {
123 let icon_path = FileIcons::get_icon(Path::new(path), cx)
124 .map(SharedString::from)
125 .unwrap_or_else(|| SharedString::from("icons/file_icons/file.svg"));
126
127 h_flex()
128 .gap_1()
129 .child("Paths: ")
130 .child(Icon::from_path(icon_path))
131 .child(ui::Label::new(path.clone()).color(Color::Muted))
132 });
133
134 let extension = filter.extension.as_ref().map(|extension| {
135 let icon_path = FileIcons::get_icon(Path::new(extension), cx)
136 .map(SharedString::from)
137 .unwrap_or_else(|| SharedString::from("icons/file_icons/file.svg"));
138
139 h_flex()
140 .gap_1()
141 .child("Extensions: ")
142 .child(Icon::from_path(icon_path))
143 .child(ui::Label::new(extension.clone()).color(Color::Muted))
144 });
145
146 if path.is_none() && extension.is_none() {
147 return None;
148 }
149
150 Some(
151 v_flex()
152 .child(ui::Label::new(heading.to_string()))
153 .gap_1()
154 .children(path)
155 .children(extension)
156 .into_any_element(),
157 )
158 }
159}
160
161impl Render for ProjectIndexView {
162 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
163 let query = self.input.query.clone();
164
165 let (header_text, content) = match &self.state {
166 ProjectIndexToolState::Error(error) => {
167 return format!("failed to search: {error:?}").into_any_element()
168 }
169 ProjectIndexToolState::CollectingQuery | ProjectIndexToolState::Searching => {
170 ("Searching...".to_string(), div())
171 }
172 ProjectIndexToolState::Finished { excerpts, .. } => {
173 let file_count = excerpts.len();
174
175 if excerpts.is_empty() {
176 ("No results found".to_string(), div())
177 } else {
178 let header_text = format!(
179 "Read {} {}",
180 file_count,
181 if file_count == 1 { "file" } else { "files" }
182 );
183
184 let el = v_flex().gap_2().children(excerpts.keys().map(|path| {
185 h_flex().gap_2().child(Icon::new(IconName::File)).child(
186 Label::new(path.path.to_string_lossy().to_string()).color(Color::Muted),
187 )
188 }));
189
190 (header_text, el)
191 }
192 }
193 };
194
195 let header = h_flex()
196 .gap_2()
197 .child(Icon::new(IconName::File))
198 .child(header_text);
199
200 v_flex()
201 .gap_3()
202 .child(
203 CollapsibleContainer::new("collapsible-container", self.expanded_header)
204 .start_slot(header)
205 .on_click(cx.listener(move |this, _, cx| {
206 this.toggle_header(cx);
207 }))
208 .child(
209 v_flex()
210 .gap_3()
211 .p_3()
212 .child(
213 h_flex()
214 .gap_2()
215 .child(Icon::new(IconName::MagnifyingGlass))
216 .child(Label::new(format!("`{}`", query)).color(Color::Muted)),
217 )
218 .children(self.render_filter_section(
219 "Includes",
220 self.input.includes.clone(),
221 cx,
222 ))
223 .children(self.render_filter_section(
224 "Excludes",
225 self.input.excludes.clone(),
226 cx,
227 ))
228 .child(content),
229 ),
230 )
231 .into_any_element()
232 }
233}
234
235impl ToolView for ProjectIndexView {
236 type Input = CodebaseQuery;
237 type SerializedState = SerializedState;
238
239 fn generate(
240 &self,
241 context: &mut assistant_tooling::ProjectContext,
242 _: &mut ViewContext<Self>,
243 ) -> String {
244 match &self.state {
245 ProjectIndexToolState::CollectingQuery => String::new(),
246 ProjectIndexToolState::Searching => String::new(),
247 ProjectIndexToolState::Error(error) => format!("failed to search: {error:?}"),
248 ProjectIndexToolState::Finished {
249 excerpts,
250 index_status,
251 } => {
252 let mut body = "found results in the following paths:\n".to_string();
253
254 for (project_path, ranges) in excerpts {
255 context.add_excerpts(project_path.clone(), ranges);
256 writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
257 }
258
259 if *index_status != Status::Idle {
260 body.push_str("Still indexing. Results may be incomplete.\n");
261 }
262
263 body
264 }
265 }
266 }
267
268 fn set_input(&mut self, input: Self::Input, cx: &mut ViewContext<Self>) {
269 self.input = input;
270 cx.notify();
271 }
272
273 fn execute(&mut self, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
274 self.state = ProjectIndexToolState::Searching;
275 cx.notify();
276
277 let project_index = self.project_index.read(cx);
278 let index_status = project_index.status();
279
280 // TODO: wire the filters into the search here instead of processing after.
281 // Otherwise we'll get zero results sometimes.
282 let search = project_index.search(self.input.query.clone(), DEFAULT_SEARCH_LIMIT, cx);
283
284 let includes = self.input.includes.clone();
285 let excludes = self.input.excludes.clone();
286
287 cx.spawn(|this, mut cx| async move {
288 let search_result = search.await;
289 this.update(&mut cx, |this, cx| {
290 match search_result {
291 Ok(search_results) => {
292 let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
293 for search_result in search_results {
294 let project_path = ProjectPath {
295 worktree_id: search_result.worktree.read(cx).id(),
296 path: search_result.path,
297 };
298
299 if let Some(includes) = &includes {
300 if !includes.matches(&project_path) {
301 continue;
302 }
303 } else if let Some(excludes) = &excludes {
304 if excludes.matches(&project_path) {
305 continue;
306 }
307 }
308
309 excerpts
310 .entry(project_path)
311 .or_default()
312 .push(search_result.range);
313 }
314 this.state = ProjectIndexToolState::Finished {
315 excerpts,
316 index_status,
317 };
318 }
319 Err(error) => {
320 this.state = ProjectIndexToolState::Error(error);
321 }
322 }
323 cx.notify();
324 })
325 })
326 }
327
328 fn serialize(&self, cx: &mut ViewContext<Self>) -> Self::SerializedState {
329 let mut serialized = SerializedState {
330 error_message: None,
331 index_status: Status::Idle,
332 worktrees: Default::default(),
333 };
334 match &self.state {
335 ProjectIndexToolState::Error(err) => serialized.error_message = Some(err.to_string()),
336 ProjectIndexToolState::Finished {
337 excerpts,
338 index_status,
339 } => {
340 serialized.index_status = *index_status;
341 if let Some(project) = self.project_index.read(cx).project().upgrade() {
342 let project = project.read(cx);
343 for (project_path, excerpts) in excerpts {
344 if let Some(worktree) =
345 project.worktree_for_id(project_path.worktree_id, cx)
346 {
347 let worktree_path = worktree.read(cx).abs_path();
348 serialized
349 .worktrees
350 .entry(worktree_path)
351 .or_default()
352 .excerpts
353 .insert(project_path.path.clone(), excerpts.clone());
354 }
355 }
356 }
357 }
358 _ => {}
359 }
360 serialized
361 }
362
363 fn deserialize(
364 &mut self,
365 serialized: Self::SerializedState,
366 cx: &mut ViewContext<Self>,
367 ) -> Result<()> {
368 if !serialized.worktrees.is_empty() {
369 let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
370 if let Some(project) = self.project_index.read(cx).project().upgrade() {
371 let project = project.read(cx);
372 for (worktree_path, worktree_state) in serialized.worktrees {
373 if let Some(worktree) = project
374 .worktrees()
375 .find(|worktree| worktree.read(cx).abs_path() == worktree_path)
376 {
377 let worktree_id = worktree.read(cx).id();
378 for (path, serialized_excerpts) in worktree_state.excerpts {
379 excerpts.insert(ProjectPath { worktree_id, path }, serialized_excerpts);
380 }
381 }
382 }
383 }
384 self.state = ProjectIndexToolState::Finished {
385 excerpts,
386 index_status: serialized.index_status,
387 };
388 }
389 cx.notify();
390 Ok(())
391 }
392}
393
394impl ProjectIndexTool {
395 pub fn new(project_index: Model<ProjectIndex>) -> Self {
396 Self { project_index }
397 }
398}
399
400impl LanguageModelTool for ProjectIndexTool {
401 type View = ProjectIndexView;
402
403 fn name(&self) -> String {
404 "semantic_search_codebase".to_string()
405 }
406
407 fn description(&self) -> String {
408 unindent::unindent(
409 r#"This search tool uses a semantic index to perform search queries across your codebase, identifying and returning excerpts of text and code possibly related to the query.
410
411 Ideal for:
412 - Discovering implementations of similar logic within the project
413 - Finding usage examples of functions, classes/structures, libraries, and other code elements
414 - Developing understanding of the codebase's architecture and design
415
416 Note: The search's effectiveness is directly related to the current state of the codebase and the specificity of your query. It is recommended that you use snippets of code that are similar to the code you wish to find."#,
417 )
418 }
419
420 fn view(&self, cx: &mut WindowContext) -> gpui::View<Self::View> {
421 cx.new_view(|_| ProjectIndexView {
422 state: ProjectIndexToolState::CollectingQuery,
423 input: Default::default(),
424 expanded_header: false,
425 project_index: self.project_index.clone(),
426 })
427 }
428}