1use anyhow::anyhow;
 2use language::ToOffset;
 3
 4use crate::templates::base::PromptArguments;
 5use crate::templates::base::PromptTemplate;
 6use std::fmt::Write;
 7
 8pub struct FileContext {}
 9
10impl PromptTemplate for FileContext {
11    fn generate(
12        &self,
13        args: &PromptArguments,
14        max_token_length: Option<usize>,
15    ) -> anyhow::Result<(String, usize)> {
16        if let Some(buffer) = &args.buffer {
17            let mut prompt = String::new();
18            // Add Initial Preamble
19            // TODO: Do we want to add the path in here?
20            writeln!(
21                prompt,
22                "The file you are currently working on has the following content:"
23            )
24            .unwrap();
25
26            let language_name = args
27                .language_name
28                .clone()
29                .unwrap_or("".to_string())
30                .to_lowercase();
31            writeln!(prompt, "```{language_name}").unwrap();
32
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                writeln!(prompt, "```").unwrap();
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        } else {
88            Err(anyhow!("no buffer provided to retrieve file context from"))
89        }
90    }
91}