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 let status = project_index.status();
144 let results = project_index.search(
145 query.query.clone(),
146 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
147 cx,
148 );
149
150 let fs = self.fs.clone();
151
152 cx.spawn(|cx| async move {
153 let results = results.await?;
154
155 let excerpts = results.into_iter().map(|result| {
156 let abs_path = result
157 .worktree
158 .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
159 let fs = fs.clone();
160
161 async move {
162 let path = result.path.clone();
163 let text = fs.load(&abs_path?).await?;
164
165 let mut start = result.range.start;
166 let mut end = result.range.end.min(text.len());
167 while !text.is_char_boundary(start) {
168 start += 1;
169 }
170 while !text.is_char_boundary(end) {
171 end -= 1;
172 }
173
174 anyhow::Ok(CodebaseExcerpt {
175 element_id: ElementId::Name(nanoid::nanoid!().into()),
176 expanded: false,
177 path: path.to_string_lossy().to_string().into(),
178 text: SharedString::from(text[start..end].to_string()),
179 score: result.score,
180 })
181 }
182 });
183
184 let excerpts = futures::future::join_all(excerpts)
185 .await
186 .into_iter()
187 .filter_map(|result| result.log_err())
188 .collect();
189 anyhow::Ok(ProjectIndexOutput { excerpts, status })
190 })
191 }
192
193 fn output_view(
194 _tool_call_id: String,
195 input: Self::Input,
196 output: Result<Self::Output>,
197 cx: &mut WindowContext,
198 ) -> gpui::View<Self::View> {
199 cx.new_view(|_cx| ProjectIndexView { input, output })
200 }
201
202 fn status_view(&self, cx: &mut WindowContext) -> Option<AnyView> {
203 Some(
204 cx.new_view(|cx| ProjectIndexStatusView::new(self.project_index.clone(), cx))
205 .into(),
206 )
207 }
208
209 fn format(_input: &Self::Input, output: &Result<Self::Output>) -> String {
210 match &output {
211 Ok(output) => {
212 let mut body = "Semantic search results:\n".to_string();
213
214 if output.status != Status::Idle {
215 body.push_str("Still indexing. Results may be incomplete.\n");
216 }
217
218 if output.excerpts.is_empty() {
219 body.push_str("No results found");
220 return body;
221 }
222
223 for excerpt in &output.excerpts {
224 body.push_str("Excerpt from ");
225 body.push_str(excerpt.path.as_ref());
226 body.push_str(", score ");
227 body.push_str(&excerpt.score.to_string());
228 body.push_str(":\n");
229 body.push_str("~~~\n");
230 body.push_str(excerpt.text.as_ref());
231 body.push_str("~~~\n");
232 }
233 body
234 }
235 Err(err) => format!("Error: {}", err),
236 }
237 }
238}
239
240struct ProjectIndexStatusView {
241 project_index: Model<ProjectIndex>,
242}
243
244impl ProjectIndexStatusView {
245 pub fn new(project_index: Model<ProjectIndex>, cx: &mut ViewContext<Self>) -> Self {
246 cx.subscribe(&project_index, |_this, _, _status: &Status, cx| {
247 cx.notify();
248 })
249 .detach();
250 Self { project_index }
251 }
252}
253
254impl Render for ProjectIndexStatusView {
255 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
256 let status = self.project_index.read(cx).status();
257
258 h_flex().gap_2().map(|element| match status {
259 Status::Idle => element.child(Label::new("Project index ready")),
260 Status::Loading => element.child(Label::new("Project index loading...")),
261 Status::Scanning { remaining_count } => element.child(Label::new(format!(
262 "Project index scanning: {remaining_count} remaining..."
263 ))),
264 })
265 }
266}