1use language::ToOffset;
 2
 3use crate::templates::base::PromptArguments;
 4use crate::templates::base::PromptTemplate;
 5use std::fmt::Write;
 6
 7pub struct FileContext {}
 8
 9impl PromptTemplate for FileContext {
10    fn generate(
11        &self,
12        args: &PromptArguments,
13        max_token_length: Option<usize>,
14    ) -> anyhow::Result<(String, usize)> {
15        let mut prompt = String::new();
16
17        // Add Initial Preamble
18        // TODO: Do we want to add the path in here?
19        writeln!(
20            prompt,
21            "The file you are currently working on has the following content:"
22        )
23        .unwrap();
24
25        let language_name = args
26            .language_name
27            .clone()
28            .unwrap_or("".to_string())
29            .to_lowercase();
30        writeln!(prompt, "```{language_name}").unwrap();
31
32        if let Some(buffer) = &args.buffer {
33            if let Some(selected_range) = &args.selected_range {
34                let start = selected_range.start.to_offset(buffer);
35                let end = selected_range.end.to_offset(buffer);
36
37                writeln!(
38                    prompt,
39                    "{}",
40                    buffer.text_for_range(0..start).collect::<String>()
41                )
42                .unwrap();
43
44                if start == end {
45                    write!(prompt, "<|START|>").unwrap();
46                } else {
47                    write!(prompt, "<|START|").unwrap();
48                }
49
50                write!(
51                    prompt,
52                    "{}",
53                    buffer.text_for_range(start..end).collect::<String>()
54                )
55                .unwrap();
56                if start != end {
57                    write!(prompt, "|END|>").unwrap();
58                }
59
60                write!(
61                    prompt,
62                    "{}",
63                    buffer.text_for_range(end..buffer.len()).collect::<String>()
64                )
65                .unwrap();
66
67                writeln!(prompt, "```").unwrap();
68
69                if start == end {
70                    writeln!(prompt, "In particular, the user's cursor is currently on the '<|START|>' span in the above content, with no text selected.").unwrap();
71                } else {
72                    writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap();
73                }
74            } else {
75                // If we dont have a selected range, include entire file.
76                writeln!(prompt, "{}", &buffer.text()).unwrap();
77            }
78        }
79
80        // Really dumb truncation strategy
81        if let Some(max_tokens) = max_token_length {
82            prompt = args.model.truncate(&prompt, max_tokens)?;
83        }
84
85        let token_count = args.model.count_tokens(&prompt)?;
86        anyhow::Ok((prompt, token_count))
87    }
88}