chat_prompts/chat/
mediatek.rs

1use crate::{
2    error::{PromptError, Result},
3    BuildChatPrompt,
4};
5use endpoints::chat::{
6    ChatCompletionAssistantMessage, ChatCompletionRequestMessage, ChatCompletionSystemMessage,
7    ChatCompletionUserMessage, ChatCompletionUserMessageContent, ContentPart,
8};
9
10/// Generate prompts for the `Breeze-7B-Instruct-v1_0` model
11#[derive(Debug, Default, Clone)]
12pub struct BreezeInstructPrompt;
13impl BreezeInstructPrompt {
14    /// Create a system prompt from a chat completion request message.
15    fn create_system_prompt(&self, message: &ChatCompletionSystemMessage) -> String {
16        let content = message.content();
17        match content.is_empty() {
18            true => String::from("<s>You are a helpful AI assistant built by MediaTek Research. The user you are helping speaks Traditional Chinese and comes from Taiwan."),
19            false => format!("<s>{content}"),
20        }
21    }
22
23    /// Create a user prompt from a chat completion request message.
24    fn append_user_message(
25        &self,
26        chat_history: impl AsRef<str>,
27        system_prompt: impl AsRef<str>,
28        message: &ChatCompletionUserMessage,
29    ) -> String {
30        let content = match message.content() {
31            ChatCompletionUserMessageContent::Text(text) => text.to_string(),
32            ChatCompletionUserMessageContent::Parts(parts) => {
33                let mut content = String::new();
34                for part in parts {
35                    if let ContentPart::Text(text_content) = part {
36                        content.push_str(text_content.text());
37                        content.push('\n');
38                    }
39                }
40                content
41            }
42        };
43
44        match chat_history.as_ref().is_empty() {
45            true => format!(
46                "{system_prompt}  [INST] {user_message} [/INST]",
47                system_prompt = system_prompt.as_ref().trim(),
48                user_message = content.trim(),
49            ),
50            false => format!(
51                "{chat_history} [INST] {user_message} [/INST]",
52                chat_history = chat_history.as_ref().trim(),
53                user_message = content.trim(),
54            ),
55        }
56    }
57
58    /// create an assistant prompt from a chat completion request message.
59    fn append_assistant_message(
60        &self,
61        chat_history: impl AsRef<str>,
62        message: &ChatCompletionAssistantMessage,
63    ) -> Result<String> {
64        let content = match message.content() {
65            Some(content) => content.to_string(),
66            // Note that the content is optional if `tool_calls` is specified.
67            None => match message.tool_calls().is_some() {
68                true => String::new(),
69                false => return Err(PromptError::NoAssistantMessage),
70            },
71        };
72
73        Ok(format!(
74            "{chat_history} {assistant_message}",
75            chat_history = chat_history.as_ref().trim(),
76            assistant_message = content.trim(),
77        ))
78    }
79}
80impl BuildChatPrompt for BreezeInstructPrompt {
81    fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
82        if messages.is_empty() {
83            return Err(crate::error::PromptError::NoMessages);
84        }
85
86        // system prompt
87        let system_prompt = match messages[0] {
88            ChatCompletionRequestMessage::System(ref message) => self.create_system_prompt(message),
89            _ => String::from("<s>You are a helpful AI assistant built by MediaTek Research. The user you are helping speaks Traditional Chinese and comes from Taiwan."),
90        };
91
92        // append user/assistant messages
93        let mut prompt = String::new();
94        for message in messages {
95            match message {
96                ChatCompletionRequestMessage::User(message) => {
97                    prompt = self.append_user_message(&prompt, &system_prompt, message);
98                }
99                ChatCompletionRequestMessage::Assistant(message) => {
100                    prompt = self.append_assistant_message(&prompt, message)?;
101                }
102                _ => continue,
103            }
104        }
105
106        Ok(prompt)
107    }
108}
chat_prompts/chat/mediatek.rs

chat_prompts/chat/
mediatek.rs
