1use std::collections::HashMap;
7
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10use uuid::Uuid;
11use validator::Validate;
12
13use crate::error::{Error, Result};
14
15#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
17pub struct Message {
18 pub id: Uuid,
20
21 pub conversation_id: String,
23
24 pub user_id: String,
26
27 pub message_type: MessageType,
29
30 #[validate(length(min = 1, max = 100_000))]
32 pub content: String,
33
34 pub attachments: Vec<Attachment>,
36
37 pub metadata: HashMap<String, serde_json::Value>,
39
40 pub timestamp: DateTime<Utc>,
42
43 pub parent_id: Option<Uuid>,
45
46 pub flags: MessageFlags,
48}
49
50impl Message {
51 #[must_use]
62 pub fn text(content: impl Into<String>) -> Self {
63 Self {
64 id: Uuid::new_v4(),
65 conversation_id: Uuid::new_v4().to_string(),
66 user_id: "anonymous".to_string(),
67 message_type: MessageType::Text,
68 content: content.into(),
69 attachments: Vec::new(),
70 metadata: HashMap::new(),
71 timestamp: Utc::now(),
72 parent_id: None,
73 flags: MessageFlags::default(),
74 }
75 }
76
77 #[must_use]
79 pub fn with_type(content: impl Into<String>, message_type: MessageType) -> Self {
80 let mut message = Self::text(content);
81 message.message_type = message_type;
82 message
83 }
84
85 #[must_use]
87 pub fn with_conversation_id(mut self, id: impl Into<String>) -> Self {
88 self.conversation_id = id.into();
89 self
90 }
91
92 #[must_use]
94 pub fn with_user_id(mut self, id: impl Into<String>) -> Self {
95 self.user_id = id.into();
96 self
97 }
98
99 #[must_use]
101 pub fn with_attachment(mut self, attachment: Attachment) -> Self {
102 self.attachments.push(attachment);
103 self
104 }
105
106 #[must_use]
108 pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
109 self.metadata.insert(key.into(), value);
110 self
111 }
112
113 #[must_use]
115 pub fn with_parent(mut self, parent_id: Uuid) -> Self {
116 self.parent_id = Some(parent_id);
117 self
118 }
119
120 #[must_use]
122 pub fn with_flags(mut self, flags: MessageFlags) -> Self {
123 self.flags = flags;
124 self
125 }
126
127 pub fn validate(&self) -> Result<()> {
133 Validate::validate(self).map_err(|e| Error::Validation(e.to_string()))?;
134
135 if self.content.is_empty() && self.attachments.is_empty() {
137 return Err(Error::InvalidInput(
138 "Message must have content or attachments".to_string(),
139 ));
140 }
141
142 Ok(())
143 }
144
145 #[must_use]
147 pub fn is_system(&self) -> bool {
148 matches!(self.message_type, MessageType::System)
149 }
150
151 #[must_use]
153 pub fn has_attachments(&self) -> bool {
154 !self.attachments.is_empty()
155 }
156
157 #[must_use]
159 pub fn attachment_size(&self) -> usize {
160 self.attachments.iter().map(|a| a.size).sum()
161 }
162}
163
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(rename_all = "lowercase")]
167pub enum MessageType {
168 Text,
170 Command,
172 System,
174 Error,
176 Embed,
178 File,
180 Image,
182 Audio,
184 Video,
186}
187
188impl MessageType {
189 #[must_use]
191 pub fn is_media(&self) -> bool {
192 matches!(self, Self::File | Self::Image | Self::Audio | Self::Video)
193 }
194}
195
196#[derive(Debug, Clone, Default, Serialize, Deserialize)]
198pub struct MessageFlags {
199 pub urgent: bool,
201 pub private: bool,
203 pub ephemeral: bool,
205 pub sensitive: bool,
207 pub bypass_filters: bool,
209 pub no_log: bool,
211}
212
213#[derive(Debug, Clone, Serialize, Deserialize)]
215pub struct Attachment {
216 pub id: Uuid,
218 pub filename: String,
220 pub mime_type: String,
222 pub size: usize,
224 pub url: String,
226 pub thumbnail_url: Option<String>,
228 pub metadata: HashMap<String, serde_json::Value>,
230}
231
232impl Attachment {
233 #[must_use]
235 pub fn new(
236 filename: impl Into<String>,
237 mime_type: impl Into<String>,
238 size: usize,
239 url: impl Into<String>,
240 ) -> Self {
241 Self {
242 id: Uuid::new_v4(),
243 filename: filename.into(),
244 mime_type: mime_type.into(),
245 size,
246 url: url.into(),
247 thumbnail_url: None,
248 metadata: HashMap::new(),
249 }
250 }
251
252 #[must_use]
254 pub fn is_image(&self) -> bool {
255 self.mime_type.starts_with("image/")
256 }
257
258 #[must_use]
260 pub fn is_video(&self) -> bool {
261 self.mime_type.starts_with("video/")
262 }
263
264 #[must_use]
266 pub fn is_audio(&self) -> bool {
267 self.mime_type.starts_with("audio/")
268 }
269}
270
271#[derive(Debug, Clone, Serialize, Deserialize)]
273pub struct Response {
274 pub id: Uuid,
276 pub conversation_id: String,
278 pub content: String,
280 pub response_type: ResponseType,
282 pub error: Option<ResponseError>,
284 pub metadata: HashMap<String, serde_json::Value>,
286 pub timestamp: DateTime<Utc>,
288 pub usage: Option<TokenUsage>,
290 pub flags: ResponseFlags,
292 pub suggestions: Vec<Suggestion>,
294}
295
296impl Response {
297 #[must_use]
299 pub fn text(conversation_id: impl Into<String>, content: impl Into<String>) -> Self {
300 Self {
301 id: Uuid::new_v4(),
302 conversation_id: conversation_id.into(),
303 content: content.into(),
304 response_type: ResponseType::Text,
305 error: None,
306 metadata: HashMap::new(),
307 timestamp: Utc::now(),
308 usage: None,
309 flags: ResponseFlags::default(),
310 suggestions: Vec::new(),
311 }
312 }
313
314 #[must_use]
316 pub fn error(conversation_id: impl Into<String>, error: ResponseError) -> Self {
317 let mut response = Self::text(conversation_id, error.message.clone());
318 response.response_type = ResponseType::Error;
319 response.error = Some(error);
320 response
321 }
322
323 #[must_use]
325 pub fn with_usage(mut self, usage: TokenUsage) -> Self {
326 self.usage = Some(usage);
327 self
328 }
329
330 #[must_use]
332 pub fn with_suggestion(mut self, suggestion: Suggestion) -> Self {
333 self.suggestions.push(suggestion);
334 self
335 }
336
337 #[must_use]
339 pub fn with_flags(mut self, flags: ResponseFlags) -> Self {
340 self.flags = flags;
341 self
342 }
343
344 #[must_use]
346 pub fn is_error(&self) -> bool {
347 self.error.is_some() || matches!(self.response_type, ResponseType::Error)
348 }
349
350 #[must_use]
352 pub fn total_tokens(&self) -> usize {
353 self.usage.as_ref().map_or(0, |u| u.total_tokens)
354 }
355}
356
357#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
359#[serde(rename_all = "lowercase")]
360pub enum ResponseType {
361 Text,
363 Markdown,
365 Html,
367 Json,
369 Error,
371 Stream,
373 Embed,
375}
376
377#[derive(Debug, Clone, Serialize, Deserialize)]
379pub struct ResponseError {
380 pub code: String,
382 pub message: String,
384 pub retryable: bool,
386 pub retry_after: Option<u64>,
388}
389
390impl ResponseError {
391 #[must_use]
393 pub fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
394 Self {
395 code: code.into(),
396 message: message.into(),
397 retryable: false,
398 retry_after: None,
399 }
400 }
401
402 #[must_use]
404 pub fn retryable(mut self, retryable: bool) -> Self {
405 self.retryable = retryable;
406 self
407 }
408
409 #[must_use]
411 pub fn retry_after(mut self, seconds: u64) -> Self {
412 self.retry_after = Some(seconds);
413 self
414 }
415}
416
417#[derive(Debug, Clone, Default, Serialize, Deserialize)]
419pub struct ResponseFlags {
420 pub truncated: bool,
422 pub partial: bool,
424 pub cached: bool,
426 pub sensitive: bool,
428 pub no_cache: bool,
430}
431
432#[derive(Debug, Clone, Serialize, Deserialize)]
434pub struct TokenUsage {
435 pub input_tokens: usize,
437 pub output_tokens: usize,
439 pub total_tokens: usize,
441 pub estimated_cost: f64,
443 pub model: String,
445}
446
447impl TokenUsage {
448 #[must_use]
450 pub fn new(input_tokens: usize, output_tokens: usize, model: impl Into<String>) -> Self {
451 let model_string = model.into();
452 let total_tokens = input_tokens + output_tokens;
453 let estimated_cost = Self::calculate_cost(input_tokens, output_tokens, &model_string);
454
455 Self {
456 input_tokens,
457 output_tokens,
458 total_tokens,
459 estimated_cost,
460 model: model_string,
461 }
462 }
463
464 fn calculate_cost(input_tokens: usize, output_tokens: usize, model: &str) -> f64 {
465 let (input_rate, output_rate) = match model {
467 "anthropic.claude-opus-4-1" => (0.015, 0.075),
468 "anthropic.claude-sonnet-4" => (0.003, 0.015),
469 "anthropic.claude-haiku" => (0.00025, 0.00125),
470 _ => (0.001, 0.002),
471 };
472
473 (input_tokens as f64 / 1000.0)
474 .mul_add(input_rate, output_tokens as f64 / 1000.0 * output_rate)
475 }
476}
477
478#[derive(Debug, Clone, Serialize, Deserialize)]
480pub struct Suggestion {
481 pub text: String,
483 pub action: SuggestionAction,
485 pub icon: Option<String>,
487}
488
489#[derive(Debug, Clone, Serialize, Deserialize)]
491#[serde(rename_all = "snake_case")]
492pub enum SuggestionAction {
493 Message(String),
495 Command(String),
497 Url(String),
499 Custom(serde_json::Value),
501}
502
503#[cfg(test)]
504mod tests {
505 use super::*;
506
507 #[test]
508 fn test_message_creation() {
509 let message = Message::text("Hello, bot!");
510 assert_eq!(message.content, "Hello, bot!");
511 assert_eq!(message.message_type, MessageType::Text);
512 assert!(message.validate().is_ok());
513 }
514
515 #[test]
516 fn test_message_builder() {
517 let attachment = Attachment::new(
518 "image.png",
519 "image/png",
520 1024,
521 "http://example.com/image.png",
522 );
523 let message = Message::text("Check this out")
524 .with_conversation_id("conv-123")
525 .with_user_id("user-456")
526 .with_attachment(attachment)
527 .with_metadata("key", serde_json::json!("value"));
528
529 assert_eq!(message.conversation_id, "conv-123");
530 assert_eq!(message.user_id, "user-456");
531 assert_eq!(message.attachments.len(), 1);
532 assert!(message.metadata.contains_key("key"));
533 }
534
535 #[test]
536 fn test_empty_message_validation() {
537 let mut message = Message::text("");
538 message.content.clear();
539 assert!(message.validate().is_err());
540 }
541
542 #[test]
543 fn test_response_creation() {
544 let response = Response::text("conv-123", "Hello, user!");
545 assert_eq!(response.content, "Hello, user!");
546 assert_eq!(response.conversation_id, "conv-123");
547 assert!(!response.is_error());
548 }
549
550 #[test]
551 fn test_error_response() {
552 let error = ResponseError::new("E001", "Something went wrong")
553 .retryable(true)
554 .retry_after(60);
555 let response = Response::error("conv-123", error);
556
557 assert!(response.is_error());
558 assert!(response.error.is_some());
559
560 let error = response.error.unwrap();
561 assert_eq!(error.code, "E001");
562 assert!(error.retryable);
563 assert_eq!(error.retry_after, Some(60));
564 }
565
566 #[test]
567 fn test_token_usage() {
568 let usage = TokenUsage::new(100, 50, "anthropic.claude-opus-4-1");
569 assert_eq!(usage.total_tokens, 150);
570 assert!(usage.estimated_cost > 0.0);
571 }
572
573 #[test]
574 fn test_attachment_types() {
575 let image = Attachment::new(
576 "photo.jpg",
577 "image/jpeg",
578 2048,
579 "http://example.com/photo.jpg",
580 );
581 assert!(image.is_image());
582 assert!(!image.is_video());
583 assert!(!image.is_audio());
584
585 let video = Attachment::new(
586 "movie.mp4",
587 "video/mp4",
588 1_048_576,
589 "http://example.com/movie.mp4",
590 );
591 assert!(!video.is_image());
592 assert!(video.is_video());
593 assert!(!video.is_audio());
594
595 let audio = Attachment::new(
596 "song.mp3",
597 "audio/mpeg",
598 4096,
599 "http://example.com/song.mp3",
600 );
601 assert!(!audio.is_image());
602 assert!(!audio.is_video());
603 assert!(audio.is_audio());
604 }
605
606 #[cfg(feature = "property-testing")]
607 mod property_tests {
608 use super::*;
609 use proptest::prelude::*;
610
611 proptest! {
612 #[test]
613 fn test_message_id_uniqueness(content in any::<String>()) {
614 let msg1 = Message::text(content.clone());
615 let msg2 = Message::text(content);
616 prop_assert_ne!(msg1.id, msg2.id);
617 }
618
619 #[test]
620 fn test_token_cost_calculation(
621 input in 0usize..100_000,
622 output in 0usize..100_000
623 ) {
624 let usage = TokenUsage::new(input, output, "test-model");
625 prop_assert_eq!(usage.total_tokens, input + output);
626 prop_assert!(usage.estimated_cost >= 0.0);
627 }
628 }
629 }
630}