Skip to main content

chat_prompts/chat/
wizard.rs

1use crate::{error::Result, BuildChatPrompt};
2use endpoints::chat::{
3    ChatCompletionRequestMessage, ChatCompletionSystemMessage, ChatCompletionUserMessage,
4    ChatCompletionUserMessageContent, ContentPart,
5};
6
7/// Generate prompts for the `wizard-vicuna` model.
8#[derive(Debug, Default, Clone)]
9pub struct WizardCoderPrompt;
10impl WizardCoderPrompt {
11    /// Create a system prompt from a chat completion request message.
12    fn create_system_prompt(&self, message: &ChatCompletionSystemMessage) -> String {
13        let content = message.content();
14        match content.is_empty() {
15            true => String::from("Below is an instruction that describes a task. Write a response that appropriately completes the request."),
16            false => content.to_string(),
17        }
18    }
19
20    /// Create a user prompt from a chat completion request message.
21    fn append_user_message(
22        &self,
23        system_prompt: impl AsRef<str>,
24        message: &ChatCompletionUserMessage,
25    ) -> String {
26        let content = match message.content() {
27            ChatCompletionUserMessageContent::Text(text) => text.to_string(),
28            ChatCompletionUserMessageContent::Parts(parts) => {
29                let mut content = String::new();
30                for part in parts {
31                    if let ContentPart::Text(text_content) = part {
32                        content.push_str(text_content.text());
33                        content.push('\n');
34                    }
35                }
36                content
37            }
38        };
39
40        format!(
41            "{system_prompt}\n\n### Instruction:\n{user_message}",
42            system_prompt = system_prompt.as_ref().trim(),
43            user_message = content.trim(),
44        )
45    }
46}
47impl BuildChatPrompt for WizardCoderPrompt {
48    fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
49        if messages.is_empty() {
50            return Err(crate::error::PromptError::NoMessages);
51        }
52
53        // system prompt
54        let system_prompt = match messages[0] {
55            ChatCompletionRequestMessage::System(ref message) => {
56                self.create_system_prompt(message)
57            }
58            _ => String::from("Below is an instruction that describes a task. Write a response that appropriately completes the request."),
59        };
60
61        let message = messages.last().unwrap();
62        let mut prompt = match message {
63            ChatCompletionRequestMessage::User(message) => {
64                self.append_user_message(system_prompt, message)
65            }
66            _ => return Err(crate::error::PromptError::NoUserMessage),
67        };
68
69        prompt.push_str("\n\n### Response:");
70
71        Ok(prompt)
72    }
73}