1use std::cmp::Reverse;
2use std::ops::Range;
3use std::sync::Arc;
4
5use language::BufferSnapshot;
6use util::ResultExt;
7
8use crate::models::LanguageModel;
9use crate::templates::repository_context::PromptCodeSnippet;
10
11pub(crate) enum PromptFileType {
12 Text,
13 Code,
14}
15
16// TODO: Set this up to manage for defaults well
17pub struct PromptArguments {
18 pub model: Arc<dyn LanguageModel>,
19 pub user_prompt: Option<String>,
20 pub language_name: Option<String>,
21 pub project_name: Option<String>,
22 pub snippets: Vec<PromptCodeSnippet>,
23 pub reserved_tokens: usize,
24 pub buffer: Option<BufferSnapshot>,
25 pub selected_range: Option<Range<usize>>,
26}
27
28impl PromptArguments {
29 pub(crate) fn get_file_type(&self) -> PromptFileType {
30 if self
31 .language_name
32 .as_ref()
33 .and_then(|name| Some(!["Markdown", "Plain Text"].contains(&name.as_str())))
34 .unwrap_or(true)
35 {
36 PromptFileType::Code
37 } else {
38 PromptFileType::Text
39 }
40 }
41}
42
43pub trait PromptTemplate {
44 fn generate(
45 &self,
46 args: &PromptArguments,
47 max_token_length: Option<usize>,
48 ) -> anyhow::Result<(String, usize)>;
49}
50
51#[repr(i8)]
52#[derive(PartialEq, Eq, Ord)]
53pub enum PromptPriority {
54 Mandatory, // Ignores truncation
55 Ordered { order: usize }, // Truncates based on priority
56}
57
58impl PartialOrd for PromptPriority {
59 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
60 match (self, other) {
61 (Self::Mandatory, Self::Mandatory) => Some(std::cmp::Ordering::Equal),
62 (Self::Mandatory, Self::Ordered { .. }) => Some(std::cmp::Ordering::Greater),
63 (Self::Ordered { .. }, Self::Mandatory) => Some(std::cmp::Ordering::Less),
64 (Self::Ordered { order: a }, Self::Ordered { order: b }) => b.partial_cmp(a),
65 }
66 }
67}
68
69pub struct PromptChain {
70 args: PromptArguments,
71 templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
72}
73
74impl PromptChain {
75 pub fn new(
76 args: PromptArguments,
77 templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)>,
78 ) -> Self {
79 PromptChain { args, templates }
80 }
81
82 pub fn generate(&self, truncate: bool) -> anyhow::Result<(String, usize)> {
83 // Argsort based on Prompt Priority
84 let seperator = "\n";
85 let seperator_tokens = self.args.model.count_tokens(seperator)?;
86 let mut sorted_indices = (0..self.templates.len()).collect::<Vec<_>>();
87 sorted_indices.sort_by_key(|&i| Reverse(&self.templates[i].0));
88
89 // If Truncate
90 let mut tokens_outstanding = if truncate {
91 Some(self.args.model.capacity()? - self.args.reserved_tokens)
92 } else {
93 None
94 };
95
96 let mut prompts = vec!["".to_string(); sorted_indices.len()];
97 for idx in sorted_indices {
98 let (_, template) = &self.templates[idx];
99
100 if let Some((template_prompt, prompt_token_count)) =
101 template.generate(&self.args, tokens_outstanding).log_err()
102 {
103 if template_prompt != "" {
104 prompts[idx] = template_prompt;
105
106 if let Some(remaining_tokens) = tokens_outstanding {
107 let new_tokens = prompt_token_count + seperator_tokens;
108 tokens_outstanding = if remaining_tokens > new_tokens {
109 Some(remaining_tokens - new_tokens)
110 } else {
111 Some(0)
112 };
113 }
114 }
115 }
116 }
117
118 prompts.retain(|x| x != "");
119
120 let full_prompt = prompts.join(seperator);
121 let total_token_count = self.args.model.count_tokens(&full_prompt)?;
122 anyhow::Ok((prompts.join(seperator), total_token_count))
123 }
124}
125
126#[cfg(test)]
127pub(crate) mod tests {
128 use super::*;
129
130 #[test]
131 pub fn test_prompt_chain() {
132 struct TestPromptTemplate {}
133 impl PromptTemplate for TestPromptTemplate {
134 fn generate(
135 &self,
136 args: &PromptArguments,
137 max_token_length: Option<usize>,
138 ) -> anyhow::Result<(String, usize)> {
139 let mut content = "This is a test prompt template".to_string();
140
141 let mut token_count = args.model.count_tokens(&content)?;
142 if let Some(max_token_length) = max_token_length {
143 if token_count > max_token_length {
144 content = args.model.truncate(&content, max_token_length)?;
145 token_count = max_token_length;
146 }
147 }
148
149 anyhow::Ok((content, token_count))
150 }
151 }
152
153 struct TestLowPriorityTemplate {}
154 impl PromptTemplate for TestLowPriorityTemplate {
155 fn generate(
156 &self,
157 args: &PromptArguments,
158 max_token_length: Option<usize>,
159 ) -> anyhow::Result<(String, usize)> {
160 let mut content = "This is a low priority test prompt template".to_string();
161
162 let mut token_count = args.model.count_tokens(&content)?;
163 if let Some(max_token_length) = max_token_length {
164 if token_count > max_token_length {
165 content = args.model.truncate(&content, max_token_length)?;
166 token_count = max_token_length;
167 }
168 }
169
170 anyhow::Ok((content, token_count))
171 }
172 }
173
174 #[derive(Clone)]
175 struct DummyLanguageModel {
176 capacity: usize,
177 }
178
179 impl LanguageModel for DummyLanguageModel {
180 fn name(&self) -> String {
181 "dummy".to_string()
182 }
183 fn count_tokens(&self, content: &str) -> anyhow::Result<usize> {
184 anyhow::Ok(content.chars().collect::<Vec<char>>().len())
185 }
186 fn truncate(&self, content: &str, length: usize) -> anyhow::Result<String> {
187 anyhow::Ok(
188 content.chars().collect::<Vec<char>>()[..length]
189 .into_iter()
190 .collect::<String>(),
191 )
192 }
193 fn capacity(&self) -> anyhow::Result<usize> {
194 anyhow::Ok(self.capacity)
195 }
196 }
197
198 let model: Arc<dyn LanguageModel> = Arc::new(DummyLanguageModel { capacity: 100 });
199 let args = PromptArguments {
200 model: model.clone(),
201 language_name: None,
202 project_name: None,
203 snippets: Vec::new(),
204 reserved_tokens: 0,
205 buffer: None,
206 selected_range: None,
207 user_prompt: None,
208 };
209
210 let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
211 (
212 PromptPriority::Ordered { order: 0 },
213 Box::new(TestPromptTemplate {}),
214 ),
215 (
216 PromptPriority::Ordered { order: 1 },
217 Box::new(TestLowPriorityTemplate {}),
218 ),
219 ];
220 let chain = PromptChain::new(args, templates);
221
222 let (prompt, token_count) = chain.generate(false).unwrap();
223
224 assert_eq!(
225 prompt,
226 "This is a test prompt template\nThis is a low priority test prompt template"
227 .to_string()
228 );
229
230 assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
231
232 // Testing with Truncation Off
233 // Should ignore capacity and return all prompts
234 let model: Arc<dyn LanguageModel> = Arc::new(DummyLanguageModel { capacity: 20 });
235 let args = PromptArguments {
236 model: model.clone(),
237 language_name: None,
238 project_name: None,
239 snippets: Vec::new(),
240 reserved_tokens: 0,
241 buffer: None,
242 selected_range: None,
243 user_prompt: None,
244 };
245
246 let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
247 (
248 PromptPriority::Ordered { order: 0 },
249 Box::new(TestPromptTemplate {}),
250 ),
251 (
252 PromptPriority::Ordered { order: 1 },
253 Box::new(TestLowPriorityTemplate {}),
254 ),
255 ];
256 let chain = PromptChain::new(args, templates);
257
258 let (prompt, token_count) = chain.generate(false).unwrap();
259
260 assert_eq!(
261 prompt,
262 "This is a test prompt template\nThis is a low priority test prompt template"
263 .to_string()
264 );
265
266 assert_eq!(model.count_tokens(&prompt).unwrap(), token_count);
267
268 // Testing with Truncation Off
269 // Should ignore capacity and return all prompts
270 let capacity = 20;
271 let model: Arc<dyn LanguageModel> = Arc::new(DummyLanguageModel { capacity });
272 let args = PromptArguments {
273 model: model.clone(),
274 language_name: None,
275 project_name: None,
276 snippets: Vec::new(),
277 reserved_tokens: 0,
278 buffer: None,
279 selected_range: None,
280 user_prompt: None,
281 };
282
283 let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
284 (
285 PromptPriority::Ordered { order: 0 },
286 Box::new(TestPromptTemplate {}),
287 ),
288 (
289 PromptPriority::Ordered { order: 1 },
290 Box::new(TestLowPriorityTemplate {}),
291 ),
292 (
293 PromptPriority::Ordered { order: 2 },
294 Box::new(TestLowPriorityTemplate {}),
295 ),
296 ];
297 let chain = PromptChain::new(args, templates);
298
299 let (prompt, token_count) = chain.generate(true).unwrap();
300
301 assert_eq!(prompt, "This is a test promp".to_string());
302 assert_eq!(token_count, capacity);
303
304 // Change Ordering of Prompts Based on Priority
305 let capacity = 120;
306 let reserved_tokens = 10;
307 let model: Arc<dyn LanguageModel> = Arc::new(DummyLanguageModel { capacity });
308 let args = PromptArguments {
309 model: model.clone(),
310 language_name: None,
311 project_name: None,
312 snippets: Vec::new(),
313 reserved_tokens,
314 buffer: None,
315 selected_range: None,
316 user_prompt: None,
317 };
318 let templates: Vec<(PromptPriority, Box<dyn PromptTemplate>)> = vec![
319 (
320 PromptPriority::Mandatory,
321 Box::new(TestLowPriorityTemplate {}),
322 ),
323 (
324 PromptPriority::Ordered { order: 0 },
325 Box::new(TestPromptTemplate {}),
326 ),
327 (
328 PromptPriority::Ordered { order: 1 },
329 Box::new(TestLowPriorityTemplate {}),
330 ),
331 ];
332 let chain = PromptChain::new(args, templates);
333
334 let (prompt, token_count) = chain.generate(true).unwrap();
335
336 assert_eq!(
337 prompt,
338 "This is a low priority test prompt template\nThis is a test prompt template\nThis is a low priority test prompt "
339 .to_string()
340 );
341 assert_eq!(token_count, capacity - reserved_tokens);
342 }
343}