universal_bot_core/
message.rs

1//! Message and response types for Universal Bot
2//!
3//! This module defines the core message structures used for communication
4//! between the bot and its users, as well as internal message passing.
5
6use 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/// A message sent to the bot
16#[derive(Debug, Clone, Serialize, Deserialize, Validate)]
17pub struct Message {
18    /// Unique message ID
19    pub id: Uuid,
20
21    /// Conversation ID for context tracking
22    pub conversation_id: String,
23
24    /// User ID who sent the message
25    pub user_id: String,
26
27    /// Message type
28    pub message_type: MessageType,
29
30    /// Message content
31    #[validate(length(min = 1, max = 100_000))]
32    pub content: String,
33
34    /// Optional attachments
35    pub attachments: Vec<Attachment>,
36
37    /// Message metadata
38    pub metadata: HashMap<String, serde_json::Value>,
39
40    /// Timestamp when the message was created
41    pub timestamp: DateTime<Utc>,
42
43    /// Optional parent message ID for threading
44    pub parent_id: Option<Uuid>,
45
46    /// Message flags
47    pub flags: MessageFlags,
48}
49
50impl Message {
51    /// Create a new text message
52    ///
53    /// # Example
54    ///
55    /// ```rust
56    /// use universal_bot_core::Message;
57    ///
58    /// let message = Message::text("Hello, bot!");
59    /// assert_eq!(message.content, "Hello, bot!");
60    /// ```
61    #[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    /// Create a new message with a specific type
78    #[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    /// Set the conversation ID
86    #[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    /// Set the user ID
93    #[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    /// Add an attachment
100    #[must_use]
101    pub fn with_attachment(mut self, attachment: Attachment) -> Self {
102        self.attachments.push(attachment);
103        self
104    }
105
106    /// Add metadata
107    #[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    /// Set the parent message ID for threading
114    #[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    /// Set message flags
121    #[must_use]
122    pub fn with_flags(mut self, flags: MessageFlags) -> Self {
123        self.flags = flags;
124        self
125    }
126
127    /// Validate the message
128    ///
129    /// # Errors
130    ///
131    /// Returns an error if validation fails.
132    pub fn validate(&self) -> Result<()> {
133        Validate::validate(self).map_err(|e| Error::Validation(e.to_string()))?;
134
135        // Additional validation
136        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    /// Check if this is a system message
146    #[must_use]
147    pub fn is_system(&self) -> bool {
148        matches!(self.message_type, MessageType::System)
149    }
150
151    /// Check if this message has attachments
152    #[must_use]
153    pub fn has_attachments(&self) -> bool {
154        !self.attachments.is_empty()
155    }
156
157    /// Get the total size of attachments in bytes
158    #[must_use]
159    pub fn attachment_size(&self) -> usize {
160        self.attachments.iter().map(|a| a.size).sum()
161    }
162}
163
164/// Type of message
165#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
166#[serde(rename_all = "lowercase")]
167pub enum MessageType {
168    /// Plain text message
169    Text,
170    /// Command to the bot
171    Command,
172    /// System message
173    System,
174    /// Error message
175    Error,
176    /// Embedded content
177    Embed,
178    /// File attachment
179    File,
180    /// Image attachment
181    Image,
182    /// Audio attachment
183    Audio,
184    /// Video attachment
185    Video,
186}
187
188impl MessageType {
189    /// Check if this is a media type
190    #[must_use]
191    pub fn is_media(&self) -> bool {
192        matches!(self, Self::File | Self::Image | Self::Audio | Self::Video)
193    }
194}
195
196/// Message flags for special handling
197#[derive(Debug, Clone, Default, Serialize, Deserialize)]
198pub struct MessageFlags {
199    /// Message is urgent
200    pub urgent: bool,
201    /// Message is private
202    pub private: bool,
203    /// Message should be ephemeral
204    pub ephemeral: bool,
205    /// Message contains sensitive content
206    pub sensitive: bool,
207    /// Message should bypass filters
208    pub bypass_filters: bool,
209    /// Message should not be logged
210    pub no_log: bool,
211}
212
213/// An attachment to a message
214#[derive(Debug, Clone, Serialize, Deserialize)]
215pub struct Attachment {
216    /// Unique attachment ID
217    pub id: Uuid,
218    /// Filename
219    pub filename: String,
220    /// MIME type
221    pub mime_type: String,
222    /// Size in bytes
223    pub size: usize,
224    /// URL or path to the attachment
225    pub url: String,
226    /// Optional thumbnail URL
227    pub thumbnail_url: Option<String>,
228    /// Attachment metadata
229    pub metadata: HashMap<String, serde_json::Value>,
230}
231
232impl Attachment {
233    /// Create a new attachment
234    #[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    /// Check if this is an image attachment
253    #[must_use]
254    pub fn is_image(&self) -> bool {
255        self.mime_type.starts_with("image/")
256    }
257
258    /// Check if this is a video attachment
259    #[must_use]
260    pub fn is_video(&self) -> bool {
261        self.mime_type.starts_with("video/")
262    }
263
264    /// Check if this is an audio attachment
265    #[must_use]
266    pub fn is_audio(&self) -> bool {
267        self.mime_type.starts_with("audio/")
268    }
269}
270
271/// A response from the bot
272#[derive(Debug, Clone, Serialize, Deserialize)]
273pub struct Response {
274    /// Unique response ID
275    pub id: Uuid,
276    /// Conversation ID
277    pub conversation_id: String,
278    /// Response content
279    pub content: String,
280    /// Response type
281    pub response_type: ResponseType,
282    /// Optional error information
283    pub error: Option<ResponseError>,
284    /// Response metadata
285    pub metadata: HashMap<String, serde_json::Value>,
286    /// Timestamp
287    pub timestamp: DateTime<Utc>,
288    /// Token usage information
289    pub usage: Option<TokenUsage>,
290    /// Response flags
291    pub flags: ResponseFlags,
292    /// Optional suggested actions
293    pub suggestions: Vec<Suggestion>,
294}
295
296impl Response {
297    /// Create a new text response
298    #[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    /// Create an error response
315    #[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    /// Add token usage information
324    #[must_use]
325    pub fn with_usage(mut self, usage: TokenUsage) -> Self {
326        self.usage = Some(usage);
327        self
328    }
329
330    /// Add a suggestion
331    #[must_use]
332    pub fn with_suggestion(mut self, suggestion: Suggestion) -> Self {
333        self.suggestions.push(suggestion);
334        self
335    }
336
337    /// Set response flags
338    #[must_use]
339    pub fn with_flags(mut self, flags: ResponseFlags) -> Self {
340        self.flags = flags;
341        self
342    }
343
344    /// Check if this response contains an error
345    #[must_use]
346    pub fn is_error(&self) -> bool {
347        self.error.is_some() || matches!(self.response_type, ResponseType::Error)
348    }
349
350    /// Get the total tokens used
351    #[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/// Type of response
358#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
359#[serde(rename_all = "lowercase")]
360pub enum ResponseType {
361    /// Plain text response
362    Text,
363    /// Markdown formatted response
364    Markdown,
365    /// HTML formatted response
366    Html,
367    /// JSON structured response
368    Json,
369    /// Error response
370    Error,
371    /// Streaming response chunk
372    Stream,
373    /// Response with embedded content
374    Embed,
375}
376
377/// Response error information
378#[derive(Debug, Clone, Serialize, Deserialize)]
379pub struct ResponseError {
380    /// Error code
381    pub code: String,
382    /// Error message
383    pub message: String,
384    /// Whether the error is retryable
385    pub retryable: bool,
386    /// Optional retry after duration in seconds
387    pub retry_after: Option<u64>,
388}
389
390impl ResponseError {
391    /// Create a new response error
392    #[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    /// Set whether the error is retryable
403    #[must_use]
404    pub fn retryable(mut self, retryable: bool) -> Self {
405        self.retryable = retryable;
406        self
407    }
408
409    /// Set retry after duration
410    #[must_use]
411    pub fn retry_after(mut self, seconds: u64) -> Self {
412        self.retry_after = Some(seconds);
413        self
414    }
415}
416
417/// Response flags
418#[derive(Debug, Clone, Default, Serialize, Deserialize)]
419pub struct ResponseFlags {
420    /// Response was truncated
421    pub truncated: bool,
422    /// Response is partial (streaming)
423    pub partial: bool,
424    /// Response was cached
425    pub cached: bool,
426    /// Response contains sensitive content
427    pub sensitive: bool,
428    /// Response should not be cached
429    pub no_cache: bool,
430}
431
432/// Token usage information
433#[derive(Debug, Clone, Serialize, Deserialize)]
434pub struct TokenUsage {
435    /// Input tokens used
436    pub input_tokens: usize,
437    /// Output tokens generated
438    pub output_tokens: usize,
439    /// Total tokens (input + output)
440    pub total_tokens: usize,
441    /// Estimated cost in USD
442    pub estimated_cost: f64,
443    /// Model used
444    pub model: String,
445}
446
447impl TokenUsage {
448    /// Create new token usage information
449    #[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        // Cost per 1K tokens (example rates)
466        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/// A suggestion for follow-up actions
479#[derive(Debug, Clone, Serialize, Deserialize)]
480pub struct Suggestion {
481    /// Suggestion text
482    pub text: String,
483    /// Action to take if selected
484    pub action: SuggestionAction,
485    /// Optional icon
486    pub icon: Option<String>,
487}
488
489/// Action for a suggestion
490#[derive(Debug, Clone, Serialize, Deserialize)]
491#[serde(rename_all = "snake_case")]
492pub enum SuggestionAction {
493    /// Send a message
494    Message(String),
495    /// Execute a command
496    Command(String),
497    /// Open a URL
498    Url(String),
499    /// Custom action
500    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}