1use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10pub struct Message {
11 pub role: Role,
12 #[serde(skip_serializing_if = "Option::is_none")]
13 pub content: Option<Content>,
14 #[serde(skip_serializing_if = "Option::is_none")]
15 pub tool_calls: Option<Vec<crate::ToolCall>>,
16 #[serde(skip_serializing_if = "Option::is_none")]
17 pub tool_call_id: Option<String>,
18 #[serde(skip_serializing_if = "Option::is_none")]
20 pub name: Option<String>,
21 #[serde(skip_serializing_if = "Option::is_none")]
23 pub reasoning: Option<String>,
24 #[serde(skip_serializing_if = "Option::is_none")]
25 pub finish_reason: Option<String>,
26}
27
28impl Default for Message {
29 fn default() -> Self {
30 Self {
31 role: Role::User,
32 content: None,
33 tool_calls: None,
34 tool_call_id: None,
35 name: None,
36 reasoning: None,
37 finish_reason: None,
38 }
39 }
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
43#[serde(rename_all = "lowercase")]
44pub enum Role {
45 System,
46 User,
47 Assistant,
48 Tool,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
53#[serde(untagged)]
54pub enum Content {
55 Text(String),
56 Parts(Vec<ContentPart>),
57}
58
59#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
60#[serde(tag = "type")]
61pub enum ContentPart {
62 #[serde(rename = "text")]
63 Text { text: String },
64 #[serde(rename = "image_url")]
65 ImageUrl { image_url: ImageUrl },
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
69pub struct ImageUrl {
70 pub url: String,
71 #[serde(skip_serializing_if = "Option::is_none")]
72 pub detail: Option<String>,
73}
74
75impl Message {
80 pub fn user(text: &str) -> Self {
81 Self {
82 role: Role::User,
83 content: Some(Content::Text(text.to_string())),
84 ..Default::default()
85 }
86 }
87
88 pub fn assistant(text: &str) -> Self {
89 Self {
90 role: Role::Assistant,
91 content: Some(Content::Text(text.to_string())),
92 ..Default::default()
93 }
94 }
95
96 pub fn system(text: &str) -> Self {
97 Self {
98 role: Role::System,
99 content: Some(Content::Text(text.to_string())),
100 ..Default::default()
101 }
102 }
103
104 pub fn tool_result(tool_call_id: &str, name: &str, content: &str) -> Self {
105 Self {
106 role: Role::Tool,
107 content: Some(Content::Text(content.to_string())),
108 tool_call_id: Some(tool_call_id.to_string()),
109 name: Some(name.to_string()),
110 ..Default::default()
111 }
112 }
113
114 pub fn assistant_with_tool_calls(text: &str, tool_calls: Vec<crate::ToolCall>) -> Self {
122 Self {
123 role: Role::Assistant,
124 content: if text.is_empty() {
125 None
126 } else {
127 Some(Content::Text(text.to_string()))
128 },
129 tool_calls: Some(tool_calls),
130 ..Default::default()
131 }
132 }
133
134 pub fn system_summary(text: String) -> Self {
136 Self {
137 role: Role::System,
138 content: Some(Content::Text(text)),
139 name: Some("context_summary".to_string()),
140 ..Default::default()
141 }
142 }
143
144 pub fn text_content(&self) -> String {
146 match &self.content {
147 Some(Content::Text(t)) => t.clone(),
148 Some(Content::Parts(parts)) => parts
149 .iter()
150 .filter_map(|p| match p {
151 ContentPart::Text { text } => Some(text.as_str()),
152 _ => None,
153 })
154 .collect::<Vec<_>>()
155 .join("\n"),
156 None => String::new(),
157 }
158 }
159
160 pub fn has_tool_calls(&self) -> bool {
162 self.tool_calls
163 .as_ref()
164 .is_some_and(|calls| !calls.is_empty())
165 }
166}
167
168impl std::fmt::Display for Role {
169 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170 f.write_str(self.as_str())
171 }
172}
173
174impl Role {
175 pub fn as_str(&self) -> &'static str {
176 match self {
177 Role::System => "system",
178 Role::User => "user",
179 Role::Assistant => "assistant",
180 Role::Tool => "tool",
181 }
182 }
183}
184
185#[cfg(test)]
190mod tests {
191 use super::*;
192
193 #[test]
194 fn user_message_roundtrip() {
195 let msg = Message::user("hello world");
196 let json = serde_json::to_string(&msg).expect("serialize");
197 let deser: Message = serde_json::from_str(&json).expect("deserialize");
198 assert_eq!(msg, deser);
199 assert_eq!(deser.text_content(), "hello world");
200 }
201
202 #[test]
203 fn assistant_message_with_tool_calls() {
204 let msg = Message {
205 role: Role::Assistant,
206 content: Some(Content::Text("I'll read that file.".into())),
207 tool_calls: Some(vec![crate::ToolCall {
208 id: "call_1".into(),
209 r#type: "function".into(),
210 function: crate::FunctionCall {
211 name: "read_file".into(),
212 arguments: r#"{"path":"src/main.rs"}"#.into(),
213 },
214 thought_signature: None,
215 }]),
216 tool_call_id: None,
217 name: None,
218 reasoning: None,
219 finish_reason: None,
220 };
221 assert!(msg.has_tool_calls());
222 let json = serde_json::to_string(&msg).expect("serialize");
223 let deser: Message = serde_json::from_str(&json).expect("deserialize");
224 assert_eq!(msg, deser);
225 }
226
227 #[test]
228 fn tool_result_message() {
229 let msg = Message::tool_result("call_1", "read_file", "fn main() {}");
230 assert_eq!(msg.role, Role::Tool);
231 assert_eq!(msg.tool_call_id.as_deref(), Some("call_1"));
232 assert_eq!(msg.text_content(), "fn main() {}");
233 }
234
235 #[test]
236 fn multimodal_content_text_extraction() {
237 let msg = Message {
238 role: Role::User,
239 content: Some(Content::Parts(vec![
240 ContentPart::Text {
241 text: "Look at this:".into(),
242 },
243 ContentPart::ImageUrl {
244 image_url: ImageUrl {
245 url: "data:image/png;base64,abc".into(),
246 detail: Some("high".into()),
247 },
248 },
249 ContentPart::Text {
250 text: "What do you see?".into(),
251 },
252 ])),
253 tool_calls: None,
254 tool_call_id: None,
255 name: None,
256 reasoning: None,
257 finish_reason: None,
258 };
259 assert_eq!(msg.text_content(), "Look at this:\nWhat do you see?");
260 }
261
262 #[test]
263 fn empty_content_returns_empty_string() {
264 let msg = Message {
265 role: Role::Assistant,
266 content: None,
267 tool_calls: None,
268 tool_call_id: None,
269 name: None,
270 reasoning: None,
271 finish_reason: None,
272 };
273 assert_eq!(msg.text_content(), "");
274 }
275
276 #[test]
277 fn role_display() {
278 assert_eq!(format!("{}", Role::System), "system");
279 assert_eq!(format!("{}", Role::User), "user");
280 assert_eq!(format!("{}", Role::Assistant), "assistant");
281 assert_eq!(format!("{}", Role::Tool), "tool");
282 }
283
284 #[test]
285 fn role_serde_roundtrip() {
286 for role in [Role::System, Role::User, Role::Assistant, Role::Tool] {
287 let json = serde_json::to_string(&role).expect("serialize");
288 let deser: Role = serde_json::from_str(&json).expect("deserialize");
289 assert_eq!(role, deser);
290 }
291 }
292}
293
294#[cfg(test)]
296mod proptests {
297 use super::*;
298 use proptest::prelude::*;
299
300 fn arb_role() -> impl Strategy<Value = Role> {
301 prop_oneof![
302 Just(Role::System),
303 Just(Role::User),
304 Just(Role::Assistant),
305 Just(Role::Tool),
306 ]
307 }
308
309 fn arb_content() -> impl Strategy<Value = Content> {
310 prop_oneof![
311 ".*".prop_map(Content::Text),
312 prop::collection::vec(".*".prop_map(|t| ContentPart::Text { text: t }), 0..5)
313 .prop_map(Content::Parts),
314 ]
315 }
316
317 fn arb_message() -> impl Strategy<Value = Message> {
318 (arb_role(), proptest::option::of(arb_content())).prop_map(|(role, content)| Message {
319 role,
320 content,
321 tool_calls: None,
322 tool_call_id: None,
323 name: None,
324 reasoning: None,
325 finish_reason: None,
326 })
327 }
328
329 proptest! {
330 #[test]
331 fn message_serde_roundtrip(msg in arb_message()) {
332 let json = serde_json::to_string(&msg).expect("serialize");
333 let deser: Message = serde_json::from_str(&json).expect("deserialize");
334 assert_eq!(msg, deser);
335 }
336
337 #[test]
338 fn text_content_never_panics(msg in arb_message()) {
339 let _ = msg.text_content(); }
341 }
342}