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, Serialize};
9use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
10use ui::{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#[derive(Default)]
19enum ProjectIndexToolState {
20 #[default]
21 CollectingQuery,
22 Searching,
23 Error(anyhow::Error),
24 Finished {
25 excerpts: BTreeMap<ProjectPath, Vec<Range<usize>>>,
26 index_status: Status,
27 },
28}
29
30pub struct ProjectIndexView {
31 project_index: Model<ProjectIndex>,
32 input: CodebaseQuery,
33 expanded_header: bool,
34 state: ProjectIndexToolState,
35}
36
37#[derive(Default, Deserialize, JsonSchema)]
38pub struct CodebaseQuery {
39 /// Semantic search query
40 query: String,
41 /// Maximum number of results to return, defaults to 20
42 limit: Option<usize>,
43}
44
45#[derive(Serialize, Deserialize)]
46pub struct SerializedState {
47 index_status: Status,
48 error_message: Option<String>,
49 worktrees: BTreeMap<Arc<Path>, WorktreeIndexOutput>,
50}
51
52#[derive(Default, Serialize, Deserialize)]
53struct WorktreeIndexOutput {
54 excerpts: BTreeMap<Arc<Path>, Vec<Range<usize>>>,
55}
56
57impl ProjectIndexView {
58 fn toggle_header(&mut self, cx: &mut ViewContext<Self>) {
59 self.expanded_header = !self.expanded_header;
60 cx.notify();
61 }
62}
63
64impl Render for ProjectIndexView {
65 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
66 let query = self.input.query.clone();
67
68 let (header_text, content) = match &self.state {
69 ProjectIndexToolState::Error(error) => {
70 return format!("failed to search: {error:?}").into_any_element()
71 }
72 ProjectIndexToolState::CollectingQuery | ProjectIndexToolState::Searching => {
73 ("Searching...".to_string(), div())
74 }
75 ProjectIndexToolState::Finished { excerpts, .. } => {
76 let file_count = excerpts.len();
77
78 let header_text = format!(
79 "Read {} {}",
80 file_count,
81 if file_count == 1 { "file" } else { "files" }
82 );
83
84 let el = v_flex().gap_2().children(excerpts.keys().map(|path| {
85 h_flex().gap_2().child(Icon::new(IconName::File)).child(
86 Label::new(path.path.to_string_lossy().to_string()).color(Color::Muted),
87 )
88 }));
89
90 (header_text, el)
91 }
92 };
93
94 let header = h_flex()
95 .gap_2()
96 .child(Icon::new(IconName::File))
97 .child(header_text);
98
99 v_flex()
100 .gap_3()
101 .child(
102 CollapsibleContainer::new("collapsible-container", self.expanded_header)
103 .start_slot(header)
104 .on_click(cx.listener(move |this, _, cx| {
105 this.toggle_header(cx);
106 }))
107 .child(
108 v_flex()
109 .gap_3()
110 .p_3()
111 .child(
112 h_flex()
113 .gap_2()
114 .child(Icon::new(IconName::MagnifyingGlass))
115 .child(Label::new(format!("`{}`", query)).color(Color::Muted)),
116 )
117 .child(content),
118 ),
119 )
120 .into_any_element()
121 }
122}
123
124impl ToolOutput for ProjectIndexView {
125 type Input = CodebaseQuery;
126 type SerializedState = SerializedState;
127
128 fn generate(
129 &self,
130 context: &mut assistant_tooling::ProjectContext,
131 _: &mut ViewContext<Self>,
132 ) -> String {
133 match &self.state {
134 ProjectIndexToolState::CollectingQuery => String::new(),
135 ProjectIndexToolState::Searching => String::new(),
136 ProjectIndexToolState::Error(error) => format!("failed to search: {error:?}"),
137 ProjectIndexToolState::Finished {
138 excerpts,
139 index_status,
140 } => {
141 let mut body = "found results in the following paths:\n".to_string();
142
143 for (project_path, ranges) in excerpts {
144 context.add_excerpts(project_path.clone(), ranges);
145 writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
146 }
147
148 if *index_status != Status::Idle {
149 body.push_str("Still indexing. Results may be incomplete.\n");
150 }
151
152 body
153 }
154 }
155 }
156
157 fn set_input(&mut self, input: Self::Input, cx: &mut ViewContext<Self>) {
158 self.input = input;
159 cx.notify();
160 }
161
162 fn execute(&mut self, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
163 self.state = ProjectIndexToolState::Searching;
164 cx.notify();
165
166 let project_index = self.project_index.read(cx);
167 let index_status = project_index.status();
168 let search = project_index.search(
169 self.input.query.clone(),
170 self.input.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
171 cx,
172 );
173
174 cx.spawn(|this, mut cx| async move {
175 let search_result = search.await;
176 this.update(&mut cx, |this, cx| {
177 match search_result {
178 Ok(search_results) => {
179 let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
180 for search_result in search_results {
181 let project_path = ProjectPath {
182 worktree_id: search_result.worktree.read(cx).id(),
183 path: search_result.path,
184 };
185 excerpts
186 .entry(project_path)
187 .or_default()
188 .push(search_result.range);
189 }
190 this.state = ProjectIndexToolState::Finished {
191 excerpts,
192 index_status,
193 };
194 }
195 Err(error) => {
196 this.state = ProjectIndexToolState::Error(error);
197 }
198 }
199 cx.notify();
200 })
201 })
202 }
203
204 fn serialize(&self, cx: &mut ViewContext<Self>) -> Self::SerializedState {
205 let mut serialized = SerializedState {
206 error_message: None,
207 index_status: Status::Idle,
208 worktrees: Default::default(),
209 };
210 match &self.state {
211 ProjectIndexToolState::Error(err) => serialized.error_message = Some(err.to_string()),
212 ProjectIndexToolState::Finished {
213 excerpts,
214 index_status,
215 } => {
216 serialized.index_status = *index_status;
217 if let Some(project) = self.project_index.read(cx).project().upgrade() {
218 let project = project.read(cx);
219 for (project_path, excerpts) in excerpts {
220 if let Some(worktree) =
221 project.worktree_for_id(project_path.worktree_id, cx)
222 {
223 let worktree_path = worktree.read(cx).abs_path();
224 serialized
225 .worktrees
226 .entry(worktree_path)
227 .or_default()
228 .excerpts
229 .insert(project_path.path.clone(), excerpts.clone());
230 }
231 }
232 }
233 }
234 _ => {}
235 }
236 serialized
237 }
238
239 fn deserialize(
240 &mut self,
241 serialized: Self::SerializedState,
242 cx: &mut ViewContext<Self>,
243 ) -> Result<()> {
244 if !serialized.worktrees.is_empty() {
245 let mut excerpts = BTreeMap::<ProjectPath, Vec<Range<usize>>>::new();
246 if let Some(project) = self.project_index.read(cx).project().upgrade() {
247 let project = project.read(cx);
248 for (worktree_path, worktree_state) in serialized.worktrees {
249 if let Some(worktree) = project
250 .worktrees()
251 .find(|worktree| worktree.read(cx).abs_path() == worktree_path)
252 {
253 let worktree_id = worktree.read(cx).id();
254 for (path, serialized_excerpts) in worktree_state.excerpts {
255 excerpts.insert(ProjectPath { worktree_id, path }, serialized_excerpts);
256 }
257 }
258 }
259 }
260 self.state = ProjectIndexToolState::Finished {
261 excerpts,
262 index_status: serialized.index_status,
263 };
264 }
265 cx.notify();
266 Ok(())
267 }
268}
269
270impl ProjectIndexTool {
271 pub fn new(project_index: Model<ProjectIndex>) -> Self {
272 Self { project_index }
273 }
274}
275
276impl LanguageModelTool for ProjectIndexTool {
277 type View = ProjectIndexView;
278
279 fn name(&self) -> String {
280 "query_codebase".to_string()
281 }
282
283 fn description(&self) -> String {
284 "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()
285 }
286
287 fn view(&self, cx: &mut WindowContext) -> gpui::View<Self::View> {
288 cx.new_view(|_| ProjectIndexView {
289 state: ProjectIndexToolState::CollectingQuery,
290 input: Default::default(),
291 expanded_header: false,
292 project_index: self.project_index.clone(),
293 })
294 }
295}