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