1use anyhow::Result;
2use assistant_tooling::LanguageModelTool;
3use gpui::{prelude::*, AppContext, Model, Task};
4use project::Fs;
5use schemars::JsonSchema;
6use semantic_index::ProjectIndex;
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<Vec<CodebaseExcerpt>>,
40}
41
42impl ProjectIndexView {
43 fn toggle_expanded(&mut self, element_id: ElementId, cx: &mut ViewContext<Self>) {
44 if let Ok(excerpts) = &mut self.output {
45 if let Some(excerpt) = excerpts
46 .iter_mut()
47 .find(|excerpt| excerpt.element_id == element_id)
48 {
49 excerpt.expanded = !excerpt.expanded;
50 cx.notify();
51 }
52 }
53 }
54}
55
56impl Render for ProjectIndexView {
57 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
58 let query = self.input.query.clone();
59
60 let result = &self.output;
61
62 let excerpts = match result {
63 Err(err) => {
64 return div().child(Label::new(format!("Error: {}", err)).color(Color::Error));
65 }
66 Ok(excerpts) => excerpts,
67 };
68
69 div()
70 .v_flex()
71 .gap_2()
72 .child(
73 div()
74 .p_2()
75 .rounded_md()
76 .bg(cx.theme().colors().editor_background)
77 .child(
78 h_flex()
79 .child(Label::new("Query: ").color(Color::Modified))
80 .child(Label::new(query).color(Color::Muted)),
81 ),
82 )
83 .children(excerpts.iter().map(|excerpt| {
84 let element_id = excerpt.element_id.clone();
85 let expanded = excerpt.expanded;
86
87 CollapsibleContainer::new(element_id.clone(), expanded)
88 .start_slot(
89 h_flex()
90 .gap_1()
91 .child(Icon::new(IconName::File).color(Color::Muted))
92 .child(Label::new(excerpt.path.clone()).color(Color::Muted)),
93 )
94 .on_click(cx.listener(move |this, _, cx| {
95 this.toggle_expanded(element_id.clone(), cx);
96 }))
97 .child(
98 div()
99 .p_2()
100 .rounded_md()
101 .bg(cx.theme().colors().editor_background)
102 .child(
103 excerpt.text.clone(), // todo!(): Show as an editor block
104 ),
105 )
106 }))
107 }
108}
109
110pub struct ProjectIndexTool {
111 project_index: Model<ProjectIndex>,
112 fs: Arc<dyn Fs>,
113}
114
115impl ProjectIndexTool {
116 pub fn new(project_index: Model<ProjectIndex>, fs: Arc<dyn Fs>) -> Self {
117 // TODO: setup a better description based on the user's current codebase.
118 Self { project_index, fs }
119 }
120}
121
122impl LanguageModelTool for ProjectIndexTool {
123 type Input = CodebaseQuery;
124 type Output = Vec<CodebaseExcerpt>;
125 type View = ProjectIndexView;
126
127 fn name(&self) -> String {
128 "query_codebase".to_string()
129 }
130
131 fn description(&self) -> String {
132 "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()
133 }
134
135 fn execute(&self, query: &Self::Input, cx: &AppContext) -> Task<Result<Self::Output>> {
136 let project_index = self.project_index.read(cx);
137
138 let results = project_index.search(
139 query.query.as_str(),
140 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
141 cx,
142 );
143
144 let fs = self.fs.clone();
145
146 cx.spawn(|cx| async move {
147 let results = results.await;
148
149 let excerpts = results.into_iter().map(|result| {
150 let abs_path = result
151 .worktree
152 .read_with(&cx, |worktree, _| worktree.abs_path().join(&result.path));
153 let fs = fs.clone();
154
155 async move {
156 let path = result.path.clone();
157 let text = fs.load(&abs_path?).await?;
158
159 let mut start = result.range.start;
160 let mut end = result.range.end.min(text.len());
161 while !text.is_char_boundary(start) {
162 start += 1;
163 }
164 while !text.is_char_boundary(end) {
165 end -= 1;
166 }
167
168 anyhow::Ok(CodebaseExcerpt {
169 element_id: ElementId::Name(nanoid::nanoid!().into()),
170 expanded: false,
171 path: path.to_string_lossy().to_string().into(),
172 text: SharedString::from(text[start..end].to_string()),
173 score: result.score,
174 })
175 }
176 });
177
178 let excerpts = futures::future::join_all(excerpts)
179 .await
180 .into_iter()
181 .filter_map(|result| result.log_err())
182 .collect();
183 anyhow::Ok(excerpts)
184 })
185 }
186
187 fn new_view(
188 _tool_call_id: String,
189 input: Self::Input,
190 output: Result<Self::Output>,
191 cx: &mut WindowContext,
192 ) -> gpui::View<Self::View> {
193 cx.new_view(|_cx| ProjectIndexView { input, output })
194 }
195
196 fn format(_input: &Self::Input, output: &Result<Self::Output>) -> String {
197 match &output {
198 Ok(excerpts) => {
199 if excerpts.len() == 0 {
200 return "No results found".to_string();
201 }
202
203 let mut body = "Semantic search results:\n".to_string();
204
205 for excerpt in excerpts {
206 body.push_str("Excerpt from ");
207 body.push_str(excerpt.path.as_ref());
208 body.push_str(", score ");
209 body.push_str(&excerpt.score.to_string());
210 body.push_str(":\n");
211 body.push_str("~~~\n");
212 body.push_str(excerpt.text.as_ref());
213 body.push_str("~~~\n");
214 }
215 body
216 }
217 Err(err) => format!("Error: {}", err),
218 }
219 }
220}