rag_toolchain/clients/
types.rs

1/// # [`PromptMessage`]
2/// This enum is used to represent the different types of messages that can be sent to the LLM.
3/// we will map the PromptMessage within the client into the compatible format.
4/// * [`PromptMessage::SystemMessage`] - This is a message that typically we asign the model a role.
5/// * [`PromptMessage::HumanMessage`] - This is a message that is from a human i.e you.
6/// * [`PromptMessage::AIMessage`] - This is a message that we get back from the LLM.
7#[derive(Debug, PartialEq, Eq, Clone)]
8pub enum PromptMessage {
9    SystemMessage(String),
10    HumanMessage(String),
11    AIMessage(String),
12}
13
14impl PromptMessage {
15    /// # [`PromptMessage::content`]
16    ///
17    /// Given that the clients will return a message that we only care for the message
18    /// this function will return the message as a string to avoid pattern matching.
19    ///
20    /// # Returns
21    /// * &[`str`] - the message content
22    pub fn content(&self) -> &str {
23        match self {
24            PromptMessage::SystemMessage(message) => message,
25            PromptMessage::HumanMessage(message) => message,
26            PromptMessage::AIMessage(message) => message,
27        }
28    }
29}
30
31#[cfg(test)]
32mod tests {
33    use super::*;
34
35    #[test]
36    fn prompt_message_content() {
37        let test_string = String::from("Test String");
38        assert_eq!(
39            &test_string,
40            PromptMessage::HumanMessage(test_string.clone()).content()
41        );
42        assert_eq!(
43            &test_string,
44            PromptMessage::AIMessage(test_string.clone()).content()
45        );
46        assert_eq!(
47            &test_string,
48            PromptMessage::SystemMessage(test_string.clone()).content()
49        );
50    }
51}
rag_toolchain/clients/types.rs

rag_toolchain/clients/
types.rs
