use super::content::Content;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MessageRole {
System,
Assistant,
User,
Tool,
}
#[derive(Debug, Clone)]
pub struct Message {
pub role: MessageRole,
pub content: Vec<Content>,
}
impl From<String> for Message {
fn from(content: String) -> Self {
Self {
role: MessageRole::System,
content: vec![content.into()],
}
}
}
impl From<&str> for Message {
fn from(content: &str) -> Self {
Self::from(content.to_owned())
}
}
impl From<Content> for Message {
fn from(content: Content) -> Self {
Self {
role: MessageRole::System,
content: vec![content],
}
}
}
impl From<(MessageRole, Content)> for Message {
fn from((role, content): (MessageRole, Content)) -> Self {
Self {
role,
content: vec![content],
}
}
}
impl From<(MessageRole, String)> for Message {
fn from((role, content): (MessageRole, String)) -> Self {
Self {
role,
content: vec![content.into()],
}
}
}
impl From<(MessageRole, &str)> for Message {
fn from((role, content): (MessageRole, &str)) -> Self {
Self::from((role, content.to_owned()))
}
}
#[derive(Debug, Clone)]
pub struct Prompt {
messages: Vec<Message>,
}
impl Prompt {
pub fn new() -> Self {
Self {
messages: Vec::new(),
}
}
pub fn add_message(&mut self, message: Message) -> &mut Self {
self.messages.push(message);
self
}
pub fn get_messages(&self) -> &[Message] {
&self.messages
}
}
impl IntoIterator for Prompt {
type Item = Message;
type IntoIter = std::vec::IntoIter<Message>;
fn into_iter(self) -> Self::IntoIter {
self.messages.into_iter()
}
}
impl Default for Prompt {
fn default() -> Self {
Self::new()
}
}
impl From<String> for Prompt {
fn from(content: String) -> Self {
let mut prompt = Self::new();
prompt.add_message(Message::from(content));
prompt
}
}
impl From<&str> for Prompt {
fn from(content: &str) -> Self {
Self::from(content.to_owned())
}
}
impl From<Message> for Prompt {
fn from(message: Message) -> Self {
let mut prompt = Self::new();
prompt.add_message(message);
prompt
}
}
impl From<Vec<Message>> for Prompt {
fn from(messages: Vec<Message>) -> Self {
let mut prompt = Self::new();
for message in messages {
prompt.add_message(message);
}
prompt
}
}