use super::BuildChatPrompt;
use crate::error::Result;
use endpoints::chat::{ChatCompletionRequestMessage, ChatCompletionRole};
#[derive(Debug, Default, Clone)]
pub struct MistralInstructPrompt;
impl MistralInstructPrompt {
fn append_user_message(
&self,
chat_history: impl AsRef<str>,
content: impl AsRef<str>,
) -> String {
match chat_history.as_ref().is_empty() {
true => format!(
"<s>[INST] {user_message} [/INST]",
user_message = content.as_ref().trim(),
),
false => format!(
"{chat_history}[INST] {user_message} [/INST]",
chat_history = chat_history.as_ref().trim(),
user_message = content.as_ref().trim(),
),
}
}
fn append_assistant_message(
&self,
chat_history: impl AsRef<str>,
content: impl AsRef<str>,
) -> String {
format!(
"{chat_history}{assistant_message}</s>",
chat_history = chat_history.as_ref().trim(),
assistant_message = content.as_ref().trim(),
)
}
}
impl BuildChatPrompt for MistralInstructPrompt {
fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
if messages.is_empty() {
return Err(crate::error::PromptError::NoMessages);
}
let mut prompt = String::new();
for message in messages {
match message.role {
ChatCompletionRole::System => continue,
ChatCompletionRole::User => {
prompt = self.append_user_message(&prompt, message.content.as_str());
}
ChatCompletionRole::Assistant => {
prompt = self.append_assistant_message(&prompt, message.content.as_str());
}
_ => {
return Err(crate::error::PromptError::UnknownRole(message.role));
}
}
}
Ok(prompt)
}
}
#[derive(Debug, Default, Clone)]
pub struct MistralLitePrompt;
impl MistralLitePrompt {
fn append_user_message(
&self,
chat_history: impl AsRef<str>,
content: impl AsRef<str>,
) -> String {
match chat_history.as_ref().is_empty() {
true => format!(
"<|prompter|>{user_message}</s>",
user_message = content.as_ref().trim(),
),
false => format!(
"{chat_history}<|prompter|>{user_message}</s>",
chat_history = chat_history.as_ref().trim(),
user_message = content.as_ref().trim(),
),
}
}
fn append_assistant_message(
&self,
chat_history: impl AsRef<str>,
content: impl AsRef<str>,
) -> String {
format!(
"{chat_history}<|assistant|>{assistant_message}</s>",
chat_history = chat_history.as_ref().trim(),
assistant_message = content.as_ref().trim(),
)
}
}
impl BuildChatPrompt for MistralLitePrompt {
fn build(&self, messages: &mut Vec<ChatCompletionRequestMessage>) -> Result<String> {
if messages.is_empty() {
return Err(crate::error::PromptError::NoMessages);
}
let mut prompt = String::new();
for message in messages {
match message.role {
ChatCompletionRole::System => continue,
ChatCompletionRole::User => {
prompt = self.append_user_message(&prompt, message.content.as_str());
}
ChatCompletionRole::Assistant => {
prompt = self.append_assistant_message(&prompt, message.content.as_str());
}
_ => {
return Err(crate::error::PromptError::UnknownRole(message.role));
}
}
}
prompt.push_str("<|assistant|>");
Ok(prompt)
}
}