unai 0.0.1

Common AI abstractions
Documentation
use super::content::Content;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MessageRole {
    /// The system role, typically used for instructions or context
    System,

    /// The assistant role, typically used for model responses
    Assistant,

    /// The user role, typically used for user inputs
    User,

    /// The tool role, typically used for tool outputs
    Tool,
}

#[derive(Debug, Clone)]
pub struct Message {
    /// The role of the message
    pub role: MessageRole,

    /// The content of the message
    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
    }
}