1use anyhow::Result;
2use assistant_tooling::LanguageModelTool;
3use gpui::{prelude::*, AnyView, Model, Task};
4use project::Fs;
5use schemars::JsonSchema;
6use semantic_index::{ProjectIndex, Status};
7use serde::Deserialize;
8use std::sync::Arc;
9use ui::{
10 div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, SharedString,
11 WindowContext,
12};
13use util::ResultExt as _;
14
15const DEFAULT_SEARCH_LIMIT: usize = 20;
16
17#[derive(Clone)]
18pub struct CodebaseExcerpt {
19 path: SharedString,
20 text: SharedString,
21 score: f32,
22 element_id: ElementId,
23 expanded: bool,
24}
25
26// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model.
27// Any changes or deletions to the `CodebaseQuery` comments will change model behavior.
28
29#[derive(Deserialize, JsonSchema)]
30pub struct CodebaseQuery {
31 /// Semantic search query
32 query: String,
33 /// Maximum number of results to return, defaults to 20
34 limit: Option<usize>,
35}
36
37pub struct ProjectIndexView {
38 input: CodebaseQuery,
39 output: Result<ProjectIndexOutput>,
40}
41
42impl ProjectIndexView {
43 fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext<Self>) {
44 if let Ok(output) = &mut self.output {
45 if let Some(excerpt) = output
46 .excerpts
47 .iter_mut()
48 .find(|excerpt| excerpt.element_id == element_id)
49 {
50 excerpt.expanded = !excerpt.expanded;
51 cx.notify();
52 }
53 }
54 }
55}
56
57impl Render for ProjectIndexView {
58 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
59 let query = self.input.query.clone();
60
61 let result = &self.output;
62
63 let output = match result {
64 Err(err) => {
65 return div().child(Label::new(format!("Error: {}", err)).color(Color::Error));
66 }
67 Ok(output) => output,
68 };
69
70 div()
71 .v_flex()
72 .gap_2()
73 .child(
74 div()
75 .p_2()
76 .rounded_md()
77 .bg(cx.theme().colors().editor_background)
78 .child(
79 h_flex()
80 .child(Label::new("Query: ").color(Color::Modified))
81 .child(Label::new(query).color(Color::Muted)),
82 ),
83 )
84 .children(output.excerpts.iter().map(|excerpt| {
85 let element_id = excerpt.element_id.clone();
86 let expanded = excerpt.expanded;
87
88 CollapsibleContainer::new(element_id.clone(), expanded)
89 .start_slot(
90 h_flex()
91 .gap_1()
92 .child(Icon::new(IconName::File).color(Color::Muted))
93 .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
94 )
95 .on_click(cx.listener(move |this, _, cx| {
96 this.toggle_expanded(element_id.clone(), cx);
97 }))
98 .child(
99 div()
100 .p_2()
101 .rounded_md()
102 .bg(cx.theme().colors().editor_background)
103 .child(excerpt.text.clone()),
104 )
105 }))
106 }
107}
108
109pub struct ProjectIndexTool {
110 project_index: Model<ProjectIndex>,
111 fs: Arc<dyn Fs>,
112}
113
114pub struct ProjectIndexOutput {
115 excerpts: Vec<CodebaseExcerpt>,
116 status: Status,
117}
118
119impl ProjectIndexTool {
120 pub fn new(project_index: Model<ProjectIndex>, fs: Arc<dyn Fs>) -> Self {
121 // Listen for project index status and update the ProjectIndexTool directly
122
123 // TODO: setup a better description based on the user's current codebase.
124 Self { project_index, fs }
125 }
126}
127
128impl LanguageModelTool for ProjectIndexTool {
129 type Input = CodebaseQuery;
130 type Output = ProjectIndexOutput;
131 type View = ProjectIndexView;
132
133 fn name(&self) -> String {
134 "query_codebase".to_string()
135 }
136
137 fn description(&self) -> String {
138 "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()
139 }
140
141 fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
142 let project_index = self.project_index.read(cx);
143
144 let status = project_index.status();
145 let results = project_index.search(
146 query.query.as_str(),
147 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
148 cx,
149 );
150
151 let fs = self.fs.clone();
152
153 cx.spawn(|cx| async move {
154 let results = results.await;
155
156 let excerpts = results.into_iter().map(|result| {
157 let abs_path = result
158 .worktree
159 .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
160 let fs = fs.clone();
161
162 async move {
163 let path = result.path.clone();
164 let text = fs.load(&abs_path?).await?;
165
166 let mut start = result.range.start;
167 let mut end = result.range.end.min(text.len());
168 while !text.is_char_boundary(start) {
169 start += 1;
170 }
171 while !text.is_char_boundary(end) {
172 end -= 1;
173 }
174
175 anyhow::Ok(CodebaseExcerpt {
176 element_id: ElementId::Name(nanoid::nanoid!().into()),
177 expanded: false,
178 path: path.to_string_lossy().to_string().into(),
179 text: SharedString::from(text[start..end].to_string()),
180 score: result.score,
181 })
182 }
183 });
184
185 let excerpts = futures::future::join_all(excerpts)
186 .await
187 .into_iter()
188 .filter_map(|result| result.log_err())
189 .collect();
190 anyhow::Ok(ProjectIndexOutput { excerpts, status })
191 })
192 }
193
194 fn output_view(
195 _tool_call_id: String,
196 input: Self::Input,
197 output: Result<Self::Output>,
198 cx: &mut WindowContext,
199 ) -> gpui::View<Self::View> {
200 cx.new_view(|_cx| ProjectIndexView { input, output })
201 }
202
203 fn status_view(&self, cx: &mut WindowContext) -> Option<AnyView> {
204 Some(
205 cx.new_view(|cx| ProjectIndexStatusView::new(self.project_index.clone(), cx))
206 .into(),
207 )
208 }
209
210 fn format(_input: &Self::Input, output: &Result<Self::Output>) -> String {
211 match &output {
212 Ok(output) => {
213 let mut body = "Semantic search results:\n".to_string();
214
215 if output.status != Status::Idle {
216 body.push_str("Still indexing. Results may be incomplete.\n");
217 }
218
219 if output.excerpts.is_empty() {
220 body.push_str("No results found");
221 return body;
222 }
223
224 for excerpt in &output.excerpts {
225 body.push_str("Excerpt from ");
226 body.push_str(excerpt.path.as_ref());
227 body.push_str(", score ");
228 body.push_str(&excerpt.score.to_string());
229 body.push_str(":\n");
230 body.push_str("~~~\n");
231 body.push_str(excerpt.text.as_ref());
232 body.push_str("~~~\n");
233 }
234 body
235 }
236 Err(err) => format!("Error: {}", err),
237 }
238 }
239}
240
241struct ProjectIndexStatusView {
242 project_index: Model<ProjectIndex>,
243}
244
245impl ProjectIndexStatusView {
246 pub fn new(project_index: Model<ProjectIndex>, cx: &mut ViewContext<Self>) -> Self {
247 cx.subscribe(&project_index, |_this, _, _status: &Status, cx| {
248 cx.notify();
249 })
250 .detach();
251 Self { project_index }
252 }
253}
254
255impl Render for ProjectIndexStatusView {
256 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
257 let status = self.project_index.read(cx).status();
258
259 h_flex().gap_2().map(|element| match status {
260 Status::Idle => element.child(Label::new("Project index ready")),
261 Status::Loading => element.child(Label::new("Project index loading...")),
262 Status::Scanning { remaining_count } => element.child(Label::new(format!(
263 "Project index scanning: {remaining_count} remaining..."
264 ))),
265 })
266 }
267}