ragit_pdl/
message.rs

1use std::fmt;
2use crate::image::ImageType;
3use crate::role::Role;
4use crate::util::encode_base64;
5
6#[derive(Clone, Debug, Eq, Hash, PartialEq)]
7pub struct Message {
8    pub role: Role,
9    pub content: Vec<MessageContent>,
10}
11
12impl Message {
13    pub fn simple_message(role: Role, message: String) -> Self {
14        Message {
15            role,
16            content: vec![MessageContent::String(message)],
17        }
18    }
19
20    pub fn is_valid_system_prompt(&self) -> bool {
21        self.role == Role::System
22        && self.content.len() == 1
23        && matches!(&self.content[0], MessageContent::String(_))
24    }
25
26    pub fn is_user_prompt(&self) -> bool {
27        self.role == Role::User
28    }
29
30    pub fn is_assistant_prompt(&self) -> bool {
31        self.role == Role::Assistant
32    }
33
34    pub fn has_image(&self) -> bool {
35        self.content.iter().any(|content| matches!(content, MessageContent::Image { .. }))
36    }
37}
38
39#[derive(Clone, Debug, Eq, Hash, PartialEq)]
40pub enum MessageContent {
41    String(String),
42    Image {
43        image_type: ImageType,
44        bytes: Vec<u8>,
45    },
46}
47
48impl MessageContent {
49    pub fn unwrap_str(&self) -> &str {
50        match self {
51            MessageContent::String(s) => s.as_str(),
52            _ => panic!("{self:?} is not a string"),
53        }
54    }
55
56    pub fn simple_message(s: String) -> Vec<Self> {
57        vec![MessageContent::String(s)]
58    }
59
60    pub fn is_string(&self) -> bool {
61        matches!(self, MessageContent::String(_))
62    }
63}
64
65impl fmt::Display for MessageContent {
66    fn fmt(&self, fmt: &mut fmt::Formatter) -> Result<(), fmt::Error> {
67        match self {
68            MessageContent::String(s) => write!(fmt, "{s}"),
69            MessageContent::Image { image_type, bytes } => write!(
70                fmt,
71                "<|raw_media({}:{})|>",
72                image_type.to_extension(),
73                encode_base64(bytes),
74            ),
75        }
76    }
77}
ragit_pdl/message.rs

ragit_pdl/
message.rs
