1use language::BufferSnapshot;
2use std::{fmt::Write, ops::Range};
3
4pub fn generate_content_prompt(
5 user_prompt: String,
6 language_name: Option<&str>,
7 buffer: BufferSnapshot,
8 range: Range<usize>,
9) -> String {
10 let mut prompt = String::new();
11
12 let content_type = match language_name {
13 None | Some("Markdown" | "Plain Text") => {
14 writeln!(
15 prompt,
16 "Here's a file of text that I'm going to ask you to make an edit to."
17 )
18 .unwrap();
19 "text"
20 }
21 Some(language_name) => {
22 writeln!(
23 prompt,
24 "Here's a file of {language_name} that I'm going to ask you to make an edit to."
25 )
26 .unwrap();
27 "code"
28 }
29 };
30
31 const MAX_CTX: usize = 50000;
32 let mut is_truncated = false;
33 if range.is_empty() {
34 prompt.push_str("The point you'll need to insert at is marked with <insert_here></insert_here>.\n\n<document>");
35 } else {
36 prompt.push_str("The section you'll need to rewrite is marked with <rewrite_this></rewrite_this> tags.\n\n<document>");
37 }
38 // Include file content.
39 let before_range = 0..range.start;
40 let truncated_before = if before_range.len() > MAX_CTX {
41 is_truncated = true;
42 range.start - MAX_CTX..range.start
43 } else {
44 before_range
45 };
46 let mut non_rewrite_len = truncated_before.len();
47 for chunk in buffer.text_for_range(truncated_before) {
48 prompt.push_str(chunk);
49 }
50 if !range.is_empty() {
51 prompt.push_str("<rewrite_this>\n");
52 for chunk in buffer.text_for_range(range.clone()) {
53 prompt.push_str(chunk);
54 }
55 prompt.push_str("\n<rewrite_this>");
56 } else {
57 prompt.push_str("<insert_here></insert_here>");
58 }
59 let after_range = range.end..buffer.len();
60 let truncated_after = if after_range.len() > MAX_CTX {
61 is_truncated = true;
62 range.end..range.end + MAX_CTX
63 } else {
64 after_range
65 };
66 non_rewrite_len += truncated_after.len();
67 for chunk in buffer.text_for_range(truncated_after) {
68 prompt.push_str(chunk);
69 }
70
71 write!(prompt, "</document>\n\n").unwrap();
72
73 if is_truncated {
74 writeln!(prompt, "The context around the relevant section has been truncated (possibly in the middle of a line) for brevity.\n").unwrap();
75 }
76
77 if range.is_empty() {
78 writeln!(
79 prompt,
80 "You can't replace {content_type}, your answer will be inserted in place of the `<insert_here></insert_here>` tags. Don't include the insert_here tags in your output.",
81 )
82 .unwrap();
83 writeln!(
84 prompt,
85 "Generate {content_type} based on the following prompt:\n\n<prompt>\n{user_prompt}\n</prompt>",
86 )
87 .unwrap();
88 writeln!(prompt, "Match the indentation in the original file in the inserted {content_type}, don't include any indentation on blank lines.\n").unwrap();
89 prompt.push_str("Immediately start with the following format with no remarks:\n\n```\n{{INSERTED_CODE}}\n```");
90 } else {
91 writeln!(prompt, "Edit the section of {content_type} in <rewrite_this></rewrite_this> tags based on the following prompt:'").unwrap();
92 writeln!(prompt, "\n<prompt>\n{user_prompt}\n</prompt>\n").unwrap();
93 let rewrite_len = range.end - range.start;
94 if rewrite_len < 20000 && rewrite_len * 2 < non_rewrite_len {
95 writeln!(prompt, "And here's the section to rewrite based on that prompt again for reference:\n\n<rewrite_this>\n").unwrap();
96 for chunk in buffer.text_for_range(range.clone()) {
97 prompt.push_str(chunk);
98 }
99 writeln!(prompt, "\n</rewrite_this>\n").unwrap();
100 }
101 writeln!(prompt, "Only make changes that are necessary to fulfill the prompt, leave everything else as-is. All surrounding {content_type} will be preserved.\n").unwrap();
102 write!(
103 prompt,
104 "Start at the indentation level in the original file in the rewritten {content_type}. "
105 )
106 .unwrap();
107 prompt.push_str("Don't stop until you've rewritten the entire section, even if you have no more changes to make, always write out the whole section with no unnecessary elisions.");
108 prompt.push_str("\n\nImmediately start with the following format with no remarks:\n\n```\n{{REWRITTEN_CODE}}\n```");
109 }
110
111 prompt
112}
113
114pub fn generate_terminal_assistant_prompt(
115 user_prompt: &str,
116 shell: Option<&str>,
117 working_directory: Option<&str>,
118 latest_output: &[String],
119) -> String {
120 let mut prompt = String::new();
121 writeln!(&mut prompt, "You are an expert terminal user.").unwrap();
122 writeln!(&mut prompt, "You will be given a description of a command and you need to respond with a command that matches the description.").unwrap();
123 writeln!(&mut prompt, "Do not include markdown blocks or any other text formatting in your response, always respond with a single command that can be executed in the given shell.").unwrap();
124 writeln!(
125 &mut prompt,
126 "Current OS name is '{}', architecture is '{}'.",
127 std::env::consts::OS,
128 std::env::consts::ARCH,
129 )
130 .unwrap();
131 if let Some(shell) = shell {
132 writeln!(&mut prompt, "Current shell is '{shell}'.").unwrap();
133 }
134 if let Some(working_directory) = working_directory {
135 writeln!(
136 &mut prompt,
137 "Current working directory is '{working_directory}'."
138 )
139 .unwrap();
140 }
141 if !latest_output.is_empty() {
142 writeln!(
143 &mut prompt,
144 "Latest non-empty {} lines of the terminal output: {:?}",
145 latest_output.len(),
146 latest_output
147 )
148 .unwrap();
149 }
150 writeln!(&mut prompt, "Here is the description of the command:").unwrap();
151 prompt.push_str(user_prompt);
152 prompt
153}