1use crate::codegen::CodegenKind;
2use language::{BufferSnapshot, OffsetRangeExt, ToOffset};
3use std::cmp::{self, Reverse};
4use std::fmt::Write;
5use std::ops::Range;
6
7fn summarize(buffer: &BufferSnapshot, selected_range: Range<impl ToOffset>) -> String {
8 #[derive(Debug)]
9 struct Match {
10 collapse: Range<usize>,
11 keep: Vec<Range<usize>>,
12 }
13
14 let selected_range = selected_range.to_offset(buffer);
15 let mut ts_matches = buffer.matches(0..buffer.len(), |grammar| {
16 Some(&grammar.embedding_config.as_ref()?.query)
17 });
18 let configs = ts_matches
19 .grammars()
20 .iter()
21 .map(|g| g.embedding_config.as_ref().unwrap())
22 .collect::<Vec<_>>();
23 let mut matches = Vec::new();
24 while let Some(mat) = ts_matches.peek() {
25 let config = &configs[mat.grammar_index];
26 if let Some(collapse) = mat.captures.iter().find_map(|cap| {
27 if Some(cap.index) == config.collapse_capture_ix {
28 Some(cap.node.byte_range())
29 } else {
30 None
31 }
32 }) {
33 let mut keep = Vec::new();
34 for capture in mat.captures.iter() {
35 if Some(capture.index) == config.keep_capture_ix {
36 keep.push(capture.node.byte_range());
37 } else {
38 continue;
39 }
40 }
41 ts_matches.advance();
42 matches.push(Match { collapse, keep });
43 } else {
44 ts_matches.advance();
45 }
46 }
47 matches.sort_unstable_by_key(|mat| (mat.collapse.start, Reverse(mat.collapse.end)));
48 let mut matches = matches.into_iter().peekable();
49
50 let mut summary = String::new();
51 let mut offset = 0;
52 let mut flushed_selection = false;
53 while let Some(mat) = matches.next() {
54 // Keep extending the collapsed range if the next match surrounds
55 // the current one.
56 while let Some(next_mat) = matches.peek() {
57 if mat.collapse.start <= next_mat.collapse.start
58 && mat.collapse.end >= next_mat.collapse.end
59 {
60 matches.next().unwrap();
61 } else {
62 break;
63 }
64 }
65
66 if offset > mat.collapse.start {
67 // Skip collapsed nodes that have already been summarized.
68 offset = cmp::max(offset, mat.collapse.end);
69 continue;
70 }
71
72 if offset <= selected_range.start && selected_range.start <= mat.collapse.end {
73 if !flushed_selection {
74 // The collapsed node ends after the selection starts, so we'll flush the selection first.
75 summary.extend(buffer.text_for_range(offset..selected_range.start));
76 summary.push_str("<|START|");
77 if selected_range.end == selected_range.start {
78 summary.push_str(">");
79 } else {
80 summary.extend(buffer.text_for_range(selected_range.clone()));
81 summary.push_str("|END|>");
82 }
83 offset = selected_range.end;
84 flushed_selection = true;
85 }
86
87 // If the selection intersects the collapsed node, we won't collapse it.
88 if selected_range.end >= mat.collapse.start {
89 continue;
90 }
91 }
92
93 summary.extend(buffer.text_for_range(offset..mat.collapse.start));
94 for keep in mat.keep {
95 summary.extend(buffer.text_for_range(keep));
96 }
97 offset = mat.collapse.end;
98 }
99
100 // Flush selection if we haven't already done so.
101 if !flushed_selection && offset <= selected_range.start {
102 summary.extend(buffer.text_for_range(offset..selected_range.start));
103 summary.push_str("<|START|");
104 if selected_range.end == selected_range.start {
105 summary.push_str(">");
106 } else {
107 summary.extend(buffer.text_for_range(selected_range.clone()));
108 summary.push_str("|END|>");
109 }
110 offset = selected_range.end;
111 }
112
113 summary.extend(buffer.text_for_range(offset..buffer.len()));
114 summary
115}
116
117pub fn generate_content_prompt(
118 user_prompt: String,
119 language_name: Option<&str>,
120 buffer: &BufferSnapshot,
121 range: Range<impl ToOffset>,
122 kind: CodegenKind,
123) -> String {
124 let mut prompt = String::new();
125
126 // General Preamble
127 if let Some(language_name) = language_name {
128 writeln!(prompt, "You're an expert {language_name} engineer.\n").unwrap();
129 } else {
130 writeln!(prompt, "You're an expert engineer.\n").unwrap();
131 }
132
133 let outline = summarize(buffer, range);
134 writeln!(
135 prompt,
136 "The file you are currently working on has the following outline:"
137 )
138 .unwrap();
139 if let Some(language_name) = language_name {
140 let language_name = language_name.to_lowercase();
141 writeln!(prompt, "```{language_name}\n{outline}\n```").unwrap();
142 } else {
143 writeln!(prompt, "```\n{outline}\n```").unwrap();
144 }
145
146 match kind {
147 CodegenKind::Generate { position: _ } => {
148 writeln!(prompt, "In particular, the user's cursor is current on the '<|START|>' span in the above outline, with no text selected.").unwrap();
149 writeln!(
150 prompt,
151 "Assume the cursor is located where the `<|START|` marker is."
152 )
153 .unwrap();
154 writeln!(
155 prompt,
156 "Text can't be replaced, so assume your answer will be inserted at the cursor."
157 )
158 .unwrap();
159 writeln!(
160 prompt,
161 "Generate text based on the users prompt: {user_prompt}"
162 )
163 .unwrap();
164 }
165 CodegenKind::Transform { range: _ } => {
166 writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
167 writeln!(
168 prompt,
169 "Modify the users code selected text based upon the users prompt: {user_prompt}"
170 )
171 .unwrap();
172 writeln!(
173 prompt,
174 "You MUST reply with only the adjusted code (within the '<|START|' and '|END|>' spans), not the entire file."
175 )
176 .unwrap();
177 }
178 }
179
180 if let Some(language_name) = language_name {
181 writeln!(prompt, "Your answer MUST always be valid {language_name}").unwrap();
182 }
183 writeln!(prompt, "Always wrap your response in a Markdown codeblock").unwrap();
184 writeln!(prompt, "Never make remarks about the output.").unwrap();
185
186 prompt
187}
188
189#[cfg(test)]
190pub(crate) mod tests {
191
192 use super::*;
193 use std::sync::Arc;
194
195 use gpui::AppContext;
196 use indoc::indoc;
197 use language::{language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, Point};
198 use settings::SettingsStore;
199
200 pub(crate) fn rust_lang() -> Language {
201 Language::new(
202 LanguageConfig {
203 name: "Rust".into(),
204 path_suffixes: vec!["rs".to_string()],
205 ..Default::default()
206 },
207 Some(tree_sitter_rust::language()),
208 )
209 .with_embedding_query(
210 r#"
211 (
212 [(line_comment) (attribute_item)]* @context
213 .
214 [
215 (struct_item
216 name: (_) @name)
217
218 (enum_item
219 name: (_) @name)
220
221 (impl_item
222 trait: (_)? @name
223 "for"? @name
224 type: (_) @name)
225
226 (trait_item
227 name: (_) @name)
228
229 (function_item
230 name: (_) @name
231 body: (block
232 "{" @keep
233 "}" @keep) @collapse)
234
235 (macro_definition
236 name: (_) @name)
237 ] @item
238 )
239 "#,
240 )
241 .unwrap()
242 }
243
244 #[gpui::test]
245 fn test_outline_for_prompt(cx: &mut AppContext) {
246 cx.set_global(SettingsStore::test(cx));
247 language_settings::init(cx);
248 let text = indoc! {"
249 struct X {
250 a: usize,
251 b: usize,
252 }
253
254 impl X {
255
256 fn new() -> Self {
257 let a = 1;
258 let b = 2;
259 Self { a, b }
260 }
261
262 pub fn a(&self, param: bool) -> usize {
263 self.a
264 }
265
266 pub fn b(&self) -> usize {
267 self.b
268 }
269 }
270 "};
271 let buffer =
272 cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx));
273 let snapshot = buffer.read(cx).snapshot();
274
275 assert_eq!(
276 summarize(&snapshot, Point::new(1, 4)..Point::new(1, 4)),
277 indoc! {"
278 struct X {
279 <|START|>a: usize,
280 b: usize,
281 }
282
283 impl X {
284
285 fn new() -> Self {}
286
287 pub fn a(&self, param: bool) -> usize {}
288
289 pub fn b(&self) -> usize {}
290 }
291 "}
292 );
293
294 assert_eq!(
295 summarize(&snapshot, Point::new(8, 12)..Point::new(8, 14)),
296 indoc! {"
297 struct X {
298 a: usize,
299 b: usize,
300 }
301
302 impl X {
303
304 fn new() -> Self {
305 let <|START|a |END|>= 1;
306 let b = 2;
307 Self { a, b }
308 }
309
310 pub fn a(&self, param: bool) -> usize {}
311
312 pub fn b(&self) -> usize {}
313 }
314 "}
315 );
316
317 assert_eq!(
318 summarize(&snapshot, Point::new(6, 0)..Point::new(6, 0)),
319 indoc! {"
320 struct X {
321 a: usize,
322 b: usize,
323 }
324
325 impl X {
326 <|START|>
327 fn new() -> Self {}
328
329 pub fn a(&self, param: bool) -> usize {}
330
331 pub fn b(&self) -> usize {}
332 }
333 "}
334 );
335
336 assert_eq!(
337 summarize(&snapshot, Point::new(21, 0)..Point::new(21, 0)),
338 indoc! {"
339 struct X {
340 a: usize,
341 b: usize,
342 }
343
344 impl X {
345
346 fn new() -> Self {}
347
348 pub fn a(&self, param: bool) -> usize {}
349
350 pub fn b(&self) -> usize {}
351 }
352 <|START|>"}
353 );
354
355 // Ensure nested functions get collapsed properly.
356 let text = indoc! {"
357 struct X {
358 a: usize,
359 b: usize,
360 }
361
362 impl X {
363
364 fn new() -> Self {
365 let a = 1;
366 let b = 2;
367 Self { a, b }
368 }
369
370 pub fn a(&self, param: bool) -> usize {
371 let a = 30;
372 fn nested() -> usize {
373 3
374 }
375 self.a + nested()
376 }
377
378 pub fn b(&self) -> usize {
379 self.b
380 }
381 }
382 "};
383 buffer.update(cx, |buffer, cx| buffer.set_text(text, cx));
384 let snapshot = buffer.read(cx).snapshot();
385 assert_eq!(
386 summarize(&snapshot, Point::new(0, 0)..Point::new(0, 0)),
387 indoc! {"
388 <|START|>struct X {
389 a: usize,
390 b: usize,
391 }
392
393 impl X {
394
395 fn new() -> Self {}
396
397 pub fn a(&self, param: bool) -> usize {}
398
399 pub fn b(&self) -> usize {}
400 }
401 "}
402 );
403 }
404}