1use anyhow::{anyhow, Result};
2use assistant_tooling::{LanguageModelTool, ToolOutput};
3use collections::BTreeMap;
4use gpui::{prelude::*, Model, Task};
5use project::ProjectPath;
6use schemars::JsonSchema;
7use semantic_index::{ProjectIndex, Status};
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::{fmt::Write as _, ops::Range, path::Path, sync::Arc};
11use ui::{div, prelude::*, CollapsibleContainer, Color, Icon, IconName, Label, WindowContext};
12
13const DEFAULT_SEARCH_LIMIT: usize = 20;
14
15pub struct ProjectIndexTool {
16 project_index: Model<ProjectIndex>,
17}
18
19// Note: Comments on a `LanguageModelTool::Input` become descriptions on the generated JSON schema as shown to the language model.
20// Any changes or deletions to the `CodebaseQuery` comments will change model behavior.
21
22#[derive(Deserialize, JsonSchema)]
23pub struct CodebaseQuery {
24 /// Semantic search query
25 query: String,
26 /// Maximum number of results to return, defaults to 20
27 limit: Option<usize>,
28}
29
30pub struct ProjectIndexView {
31 input: CodebaseQuery,
32 status: Status,
33 excerpts: Result<BTreeMap<ProjectPath, Vec<Range<usize>>>>,
34 element_id: ElementId,
35 expanded_header: bool,
36}
37
38#[derive(Serialize, Deserialize)]
39pub struct ProjectIndexOutput {
40 status: Status,
41 worktrees: BTreeMap<Arc<Path>, WorktreeIndexOutput>,
42}
43
44#[derive(Serialize, Deserialize)]
45struct WorktreeIndexOutput {
46 excerpts: BTreeMap<Arc<Path>, Vec<Range<usize>>>,
47}
48
49impl ProjectIndexView {
50 fn toggle_header(&mut self, cx: &mut ViewContext<Self>) {
51 self.expanded_header = !self.expanded_header;
52 cx.notify();
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 let excerpts = match &self.excerpts {
60 Err(err) => {
61 return div().child(Label::new(format!("Error: {}", err)).color(Color::Error));
62 }
63 Ok(excerpts) => excerpts,
64 };
65
66 let file_count = excerpts.len();
67 let header = h_flex()
68 .gap_2()
69 .child(Icon::new(IconName::File))
70 .child(format!(
71 "Read {} {}",
72 file_count,
73 if file_count == 1 { "file" } else { "files" }
74 ));
75
76 v_flex().gap_3().child(
77 CollapsibleContainer::new(self.element_id.clone(), self.expanded_header)
78 .start_slot(header)
79 .on_click(cx.listener(move |this, _, cx| {
80 this.toggle_header(cx);
81 }))
82 .child(
83 v_flex()
84 .gap_3()
85 .p_3()
86 .child(
87 h_flex()
88 .gap_2()
89 .child(Icon::new(IconName::MagnifyingGlass))
90 .child(Label::new(format!("`{}`", query)).color(Color::Muted)),
91 )
92 .child(v_flex().gap_2().children(excerpts.keys().map(|path| {
93 h_flex().gap_2().child(Icon::new(IconName::File)).child(
94 Label::new(path.path.to_string_lossy().to_string())
95 .color(Color::Muted),
96 )
97 }))),
98 ),
99 )
100 }
101}
102
103impl ToolOutput for ProjectIndexView {
104 fn generate(
105 &self,
106 context: &mut assistant_tooling::ProjectContext,
107 _: &mut WindowContext,
108 ) -> String {
109 match &self.excerpts {
110 Ok(excerpts) => {
111 let mut body = "found results in the following paths:\n".to_string();
112
113 for (project_path, ranges) in excerpts {
114 context.add_excerpts(project_path.clone(), ranges);
115 writeln!(&mut body, "* {}", &project_path.path.display()).unwrap();
116 }
117
118 if self.status != Status::Idle {
119 body.push_str("Still indexing. Results may be incomplete.\n");
120 }
121
122 body
123 }
124 Err(err) => format!("Error: {}", err),
125 }
126 }
127}
128
129impl ProjectIndexTool {
130 pub fn new(project_index: Model<ProjectIndex>) -> Self {
131 Self { project_index }
132 }
133}
134
135impl LanguageModelTool for ProjectIndexTool {
136 type Input = CodebaseQuery;
137 type Output = ProjectIndexOutput;
138 type View = ProjectIndexView;
139
140 fn name(&self) -> String {
141 "query_codebase".to_string()
142 }
143
144 fn description(&self) -> String {
145 "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()
146 }
147
148 fn execute(&self, query: &Self::Input, cx: &mut WindowContext) -> Task<Result<Self::Output>> {
149 let project_index = self.project_index.read(cx);
150 let status = project_index.status();
151 let search = project_index.search(
152 query.query.clone(),
153 query.limit.unwrap_or(DEFAULT_SEARCH_LIMIT),
154 cx,
155 );
156
157 cx.spawn(|mut cx| async move {
158 let search_results = search.await?;
159
160 cx.update(|cx| {
161 let mut output = ProjectIndexOutput {
162 status,
163 worktrees: Default::default(),
164 };
165
166 for search_result in search_results {
167 let worktree_path = search_result.worktree.read(cx).abs_path();
168 let excerpts = &mut output
169 .worktrees
170 .entry(worktree_path)
171 .or_insert(WorktreeIndexOutput {
172 excerpts: Default::default(),
173 })
174 .excerpts;
175
176 let excerpts_for_path = excerpts.entry(search_result.path).or_default();
177 let ix = match excerpts_for_path
178 .binary_search_by_key(&search_result.range.start, |r| r.start)
179 {
180 Ok(ix) | Err(ix) => ix,
181 };
182 excerpts_for_path.insert(ix, search_result.range);
183 }
184
185 output
186 })
187 })
188 }
189
190 fn view(
191 &self,
192 input: Self::Input,
193 output: Result<Self::Output>,
194 cx: &mut WindowContext,
195 ) -> gpui::View<Self::View> {
196 cx.new_view(|cx| {
197 let status;
198 let excerpts;
199 match output {
200 Ok(output) => {
201 status = output.status;
202 let project_index = self.project_index.read(cx);
203 if let Some(project) = project_index.project().upgrade() {
204 let project = project.read(cx);
205 excerpts = Ok(output
206 .worktrees
207 .into_iter()
208 .filter_map(|(abs_path, output)| {
209 for worktree in project.worktrees() {
210 let worktree = worktree.read(cx);
211 if worktree.abs_path() == abs_path {
212 return Some((worktree.id(), output.excerpts));
213 }
214 }
215 None
216 })
217 .flat_map(|(worktree_id, excerpts)| {
218 excerpts.into_iter().map(move |(path, ranges)| {
219 (ProjectPath { worktree_id, path }, ranges)
220 })
221 })
222 .collect::<BTreeMap<_, _>>());
223 } else {
224 excerpts = Err(anyhow!("project was dropped"));
225 }
226 }
227 Err(err) => {
228 status = Status::Idle;
229 excerpts = Err(err);
230 }
231 };
232
233 ProjectIndexView {
234 input,
235 status,
236 excerpts,
237 element_id: ElementId::Name(nanoid::nanoid!().into()),
238 expanded_header: false,
239 }
240 })
241 }
242
243 fn render_running(arguments: &Option<Value>, _: &mut WindowContext) -> impl IntoElement {
244 let text: String = arguments
245 .as_ref()
246 .and_then(|arguments| arguments.get("query"))
247 .and_then(|query| query.as_str())
248 .map(|query| format!("Searching for: {}", query))
249 .unwrap_or_else(|| "Preparing search...".to_string());
250
251 CollapsibleContainer::new(ElementId::Name(nanoid::nanoid!().into()), false).start_slot(text)
252 }
253}