messageforge/
chat_message.rs

1use crate::prelude::*;
2use derive_base_message::BaseMessage;
3
4#[derive(BaseMessage, Debug, Serialize, Deserialize)]
5pub struct ChatMessage {
6    role: String,
7    #[serde(flatten)]
8    base: BaseMessageFields,
9}
10
11#[cfg(test)]
12mod tests {
13    use super::*;
14    use serde_json;
15    use std::collections::HashMap;
16
17    #[test]
18    fn test_chat_message_serialization_with_empty_fields() {
19        let chat_message = ChatMessage::new(
20            "Test message",
21            "User".to_string(), // role
22        );
23
24        let serialized = serde_json::to_string(&chat_message).expect("Serialization failed");
25
26        let expected =
27            r#"{"role":"User","content":"Test message","example":false,"message_type":"Chat"}"#;
28        assert_eq!(serialized, expected);
29    }
30
31    #[test]
32    fn test_chat_message_serialization_with_values() {
33        let mut additional_kwargs = HashMap::new();
34        additional_kwargs.insert("key1".to_string(), "value1".to_string());
35
36        let mut response_metadata = HashMap::new();
37        response_metadata.insert("meta_key".to_string(), "meta_value".to_string());
38
39        let mut chat_message = ChatMessage::new("Test message", "User".to_string());
40
41        chat_message.base.example = true;
42        chat_message.base.additional_kwargs = additional_kwargs.clone();
43        chat_message.base.response_metadata = response_metadata.clone();
44        chat_message.base.id = Some("1234".to_string());
45        chat_message.base.name = Some("Test Name".to_string());
46
47        let serialized = serde_json::to_string(&chat_message).expect("Serialization failed");
48
49        // The expected output should not have an extra closing brace
50        let expected = r#"{"role":"User","content":"Test message","example":true,"message_type":"Chat","additional_kwargs":{"key1":"value1"},"response_metadata":{"meta_key":"meta_value"},"id":"1234","name":"Test Name"}"#;
51        assert_eq!(serialized, expected);
52    }
53
54    #[test]
55    fn test_chat_message_deserialization() {
56        let json_data = r#"
57        {
58            "role": "User",
59            "content": "Test message",
60            "example": true,
61            "message_type": "Chat",
62            "additional_kwargs": {
63                "key1": "value1"
64            },
65            "response_metadata": {
66                "meta_key": "meta_value"
67            },
68            "id": "1234",
69            "name": "Test Name"
70        }
71        "#;
72
73        let chat_message: ChatMessage =
74            serde_json::from_str(json_data).expect("Deserialization failed");
75
76        assert_eq!(chat_message.role, "User");
77        assert_eq!(chat_message.base.content, "Test message");
78        assert!(chat_message.base.example);
79        assert_eq!(chat_message.base.message_type, MessageType::Chat);
80        assert_eq!(
81            chat_message.base.additional_kwargs.get("key1").unwrap(),
82            "value1"
83        );
84        assert_eq!(
85            chat_message.base.response_metadata.get("meta_key").unwrap(),
86            "meta_value"
87        );
88        assert_eq!(chat_message.base.id.as_deref(), Some("1234"));
89        assert_eq!(chat_message.base.name.as_deref(), Some("Test Name"));
90    }
91
92    #[test]
93    fn test_chat_message_serialization_with_partial_values() {
94        let mut additional_kwargs = HashMap::new();
95        additional_kwargs.insert("key2".to_string(), "value2".to_string());
96
97        let mut chat_message = ChatMessage::new("Partial message", "User".to_string());
98
99        chat_message.base.additional_kwargs = additional_kwargs.clone();
100        chat_message.base.id = Some("5678".to_string());
101
102        let serialized = serde_json::to_string(&chat_message).expect("Serialization failed");
103
104        let expected = r#"{"role":"User","content":"Partial message","example":false,"message_type":"Chat","additional_kwargs":{"key2":"value2"},"id":"5678"}"#;
105        assert_eq!(serialized, expected);
106    }
107}