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