1use crate::prompts::base::{PromptArguments, PromptFileType, PromptTemplate};
 2use std::fmt::Write;
 3
 4pub struct EngineerPreamble {}
 5
 6impl PromptTemplate for EngineerPreamble {
 7    fn generate(
 8        &self,
 9        args: &PromptArguments,
10        max_token_length: Option<usize>,
11    ) -> anyhow::Result<(String, usize)> {
12        let mut prompts = Vec::new();
13
14        match args.get_file_type() {
15            PromptFileType::Code => {
16                prompts.push(format!(
17                    "You are an expert {}engineer.",
18                    args.language_name.clone().unwrap_or("".to_string()) + " "
19                ));
20            }
21            PromptFileType::Text => {
22                prompts.push("You are an expert engineer.".to_string());
23            }
24        }
25
26        if let Some(project_name) = args.project_name.clone() {
27            prompts.push(format!(
28                "You are currently working inside the '{project_name}' project in code editor Zed."
29            ));
30        }
31
32        if let Some(mut remaining_tokens) = max_token_length {
33            let mut prompt = String::new();
34            let mut total_count = 0;
35            for prompt_piece in prompts {
36                let prompt_token_count =
37                    args.model.count_tokens(&prompt_piece)? + args.model.count_tokens("\n")?;
38                if remaining_tokens > prompt_token_count {
39                    writeln!(prompt, "{prompt_piece}").unwrap();
40                    remaining_tokens -= prompt_token_count;
41                    total_count += prompt_token_count;
42                }
43            }
44
45            anyhow::Ok((prompt, total_count))
46        } else {
47            let prompt = prompts.join("\n");
48            let token_count = args.model.count_tokens(&prompt)?;
49            anyhow::Ok((prompt, token_count))
50        }
51    }
52}