chat_prompts/chat/
megrez.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 models using ChatML template.
11#[derive(Debug, Default, Clone)]
12pub struct MegrezPrompt;
13impl MegrezPrompt {
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("<|role_start|>system<|role_end|>You are a helpful assistant. Answer questions as concisely and accurately as possible.<|turn_end|>"),
19            false => format!(
20                "<|role_start|>system<|role_end|>{content}<|turn_end|>"
21            ),
22        }
23    }
24
25    /// Create a user prompt from a chat completion request message.
26    fn append_user_message(
27        &self,
28        chat_history: impl AsRef<str>,
29        system_prompt: impl AsRef<str>,
30        message: &ChatCompletionUserMessage,
31    ) -> String {
32        let content = match message.content() {
33            ChatCompletionUserMessageContent::Text(text) => text.to_string(),
34            ChatCompletionUserMessageContent::Parts(parts) => {
35                let mut content = String::new();
36                for part in parts {
37                    if let ContentPart::Text(text_content) = part {
38                        content.push_str(text_content.text());
39                        content.push('\n');
40                    }
41                }
42                content
43            }
44        };
45
46        match chat_history.as_ref().is_empty() {
47            true => match system_prompt.as_ref().is_empty() {
48                true => {
49                    format!(
50                        "<|role_start|>user<|role_end|>{user_message}<|turn_end|>",
51                        user_message = content.trim(),
52                    )
53                }
54                false => {
55                    format!(
56                        "{system_prompt}<|role_start|>user<|role_end|>{user_message}<|turn_end|>",
57                        system_prompt = system_prompt.as_ref().trim(),
58                        user_message = content.trim(),
59                    )
60                }
61            },
62            false => format!(
63                "{chat_history}<|role_start|>user<|role_end|>{user_message}<|turn_end|>",
64                chat_history = chat_history.as_ref().trim(),
65                user_message = content.trim(),
66            ),
67        }
68    }
69
70    /// create an assistant prompt from a chat completion request message.
71    fn append_assistant_message(
72        &self,
73        chat_history: impl AsRef<str>,
74        message: &ChatCompletionAssistantMessage,
75    ) -> Result<String> {
76        let content = match message.content() {
77            Some(content) => content.to_string(),
78            // Note that the content is optional if `tool_calls` is specified.
79            None => match message.tool_calls().is_some() {
80                true => String::new(),
81                false => return Err(PromptError::NoAssistantMessage),
82            },
83        };
84
85        Ok(format!(
86            "{chat_history}<|role_start|>assistant<|role_end|>{assistant_message}<|turn_end|>",
87            chat_history = chat_history.as_ref().trim(),
88            assistant_message = content.trim(),
89        ))
90    }
91}
92impl BuildChatPrompt for MegrezPrompt {
93    fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
94        if messages.is_empty() {
95            return Err(crate::error::PromptError::NoMessages);
96        }
97
98        // system prompt
99        let system_prompt = match messages[0] {
100            ChatCompletionRequestMessage::System(ref message) => self.create_system_prompt(message),
101            _ => String::from("<|role_start|>system<|role_end|>You are a helpful assistant. Answer questions as concisely and accurately as possible.<|turn_end|>"),
102        };
103
104        // append user/assistant messages
105        let mut prompt = String::new();
106        for message in messages {
107            match message {
108                ChatCompletionRequestMessage::User(message) => {
109                    prompt = self.append_user_message(&prompt, &system_prompt, message);
110                }
111                ChatCompletionRequestMessage::Assistant(message) => {
112                    prompt = self.append_assistant_message(&prompt, message)?;
113                }
114                _ => continue,
115            }
116        }
117
118        prompt.push_str("<|role_start|>assistant<|role_end|>");
119
120        Ok(prompt)
121    }
122}
chat_prompts/chat/megrez.rs

chat_prompts/chat/
megrez.rs
