1use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9use uuid::Uuid;
10
11#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
13#[serde(rename_all = "camelCase")]
14pub struct Message {
15 pub id: String,
17
18 pub role: MessageRole,
20
21 pub parts: Vec<MessagePart>,
23
24 #[serde(default, skip_serializing_if = "Option::is_none")]
26 pub metadata: Option<serde_json::Value>,
27}
28
29impl Message {
30 pub fn user(parts: Vec<MessagePart>) -> Self {
32 Self {
33 id: Uuid::new_v4().to_string(),
34 role: MessageRole::User,
35 parts,
36 metadata: None,
37 }
38 }
39
40 pub fn agent(parts: Vec<MessagePart>) -> Self {
42 Self {
43 id: Uuid::new_v4().to_string(),
44 role: MessageRole::Agent,
45 parts,
46 metadata: None,
47 }
48 }
49
50 pub fn user_text(text: impl Into<String>) -> Self {
52 Self::user(vec![MessagePart::text(text)])
53 }
54
55 pub fn agent_text(text: impl Into<String>) -> Self {
57 Self::agent(vec![MessagePart::text(text)])
58 }
59
60 pub fn text_content(&self) -> String {
62 self.parts
63 .iter()
64 .filter_map(|p| match p {
65 MessagePart::Text { text, .. } => Some(text.as_str()),
66 _ => None,
67 })
68 .collect::<Vec<_>>()
69 .join("\n")
70 }
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
75#[serde(rename_all = "lowercase")]
76pub enum MessageRole {
77 User,
79 Agent,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
87#[serde(rename_all = "camelCase", tag = "type")]
88pub enum MessagePart {
89 #[serde(rename = "text")]
91 Text {
92 text: String,
93 #[serde(skip_serializing_if = "Option::is_none")]
94 media_type: Option<String>,
95 },
96
97 #[serde(rename = "file")]
99 File { file: FilePart },
100
101 #[serde(rename = "data")]
103 Data { data: DataPart },
104}
105
106impl MessagePart {
107 pub fn text(text: impl Into<String>) -> Self {
109 Self::Text {
110 text: text.into(),
111 media_type: None,
112 }
113 }
114
115 pub fn text_with_type(text: impl Into<String>, media_type: impl Into<String>) -> Self {
117 Self::Text {
118 text: text.into(),
119 media_type: Some(media_type.into()),
120 }
121 }
122
123 pub fn file_inline(
125 name: impl Into<String>,
126 media_type: impl Into<String>,
127 data: Vec<u8>,
128 ) -> Self {
129 use base64::Engine;
130 Self::File {
131 file: FilePart {
132 name: Some(name.into()),
133 media_type: Some(media_type.into()),
134 data: Some(base64::engine::general_purpose::STANDARD.encode(data)),
135 url: None,
136 },
137 }
138 }
139
140 pub fn file_url(url: impl Into<String>, name: Option<String>) -> Self {
142 Self::File {
143 file: FilePart {
144 name,
145 media_type: None,
146 data: None,
147 url: Some(url.into()),
148 },
149 }
150 }
151
152 pub fn data(value: serde_json::Value, media_type: Option<String>) -> Self {
154 Self::Data {
155 data: DataPart { value, media_type },
156 }
157 }
158}
159
160#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
162#[serde(rename_all = "camelCase")]
163pub struct FilePart {
164 #[serde(skip_serializing_if = "Option::is_none")]
166 pub name: Option<String>,
167
168 #[serde(skip_serializing_if = "Option::is_none")]
170 pub media_type: Option<String>,
171
172 #[serde(skip_serializing_if = "Option::is_none")]
174 pub data: Option<String>,
175
176 #[serde(skip_serializing_if = "Option::is_none")]
178 pub url: Option<String>,
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
183#[serde(rename_all = "camelCase")]
184pub struct DataPart {
185 pub value: serde_json::Value,
187
188 #[serde(skip_serializing_if = "Option::is_none")]
190 pub media_type: Option<String>,
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196
197 #[test]
198 fn test_message_creation() {
199 let msg = Message::user_text("Hello, summarize this document");
200 assert_eq!(msg.role, MessageRole::User);
201 assert_eq!(msg.parts.len(), 1);
202 assert_eq!(msg.text_content(), "Hello, summarize this document");
203 }
204
205 #[test]
206 fn test_message_serialization() {
207 let msg = Message::user(vec![
208 MessagePart::text("Check this file"),
209 MessagePart::file_url("https://example.com/doc.pdf", Some("doc.pdf".into())),
210 MessagePart::data(
211 serde_json::json!({"priority": "high"}),
212 Some("application/json".into()),
213 ),
214 ]);
215
216 let json = serde_json::to_string_pretty(&msg).unwrap();
217 let parsed: Message = serde_json::from_str(&json).unwrap();
218 assert_eq!(parsed.parts.len(), 3);
219 }
220}