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