batch-mode-batch-schema 0.2.4

Defines the schema and structures for batch processing, including batch choices, message roles, error handling, and token management, with support for JSON deserialization and validation.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

// ---------------- [ File: batch-mode-batch-schema/src/message_role.rs ]
crate::ix!();

#[derive(Default,Debug,Clone,PartialEq,Eq,Hash)]
pub enum MessageRole {
    #[default]
    Assistant,
    User,
    System,
    Tool,
    Function,
    Unknown(String),
}

impl<'de> Deserialize<'de> for MessageRole {
    fn deserialize<D>(deserializer: D) -> Result<MessageRole, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        match s.as_str() {
            "assistant" => Ok(MessageRole::Assistant),
            "user" => Ok(MessageRole::User),
            "system" => Ok(MessageRole::System),
            "tool" => Ok(MessageRole::Tool),
            "function" => Ok(MessageRole::Function),
            other => Ok(MessageRole::Unknown(other.to_string())),
        }
    }
}

impl Serialize for MessageRole {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let s = match self {
            MessageRole::Assistant => "assistant",
            MessageRole::User => "user",
            MessageRole::System => "system",
            MessageRole::Tool => "tool",
            MessageRole::Function => "function",
            MessageRole::Unknown(other) => other.as_str(),
        };
        serializer.serialize_str(s)
    }
}

#[cfg(test)]
mod message_role_tests {
    use super::*;

    // Test suite for MessageRole
    #[test]
    fn test_message_role_deserialization() {
        // Known roles
        let roles = vec!["assistant", "user", "system", "tool", "function"];
        let expected_roles = vec![
            MessageRole::Assistant,
            MessageRole::User,
            MessageRole::System,
            MessageRole::Tool,
            MessageRole::Function,
        ];

        for (role_str, expected_role) in roles.iter().zip(expected_roles.iter()) {
            let json = format!("\"{}\"", role_str);
            let role: MessageRole = serde_json::from_str(&json).unwrap();
            pretty_assert_eq!(&role, expected_role);
        }

        // Unknown role
        let json = "\"unknown_role\"";
        let role: MessageRole = serde_json::from_str(json).unwrap();
        pretty_assert_eq!(role, MessageRole::Unknown("unknown_role".to_string()));

        // Empty string as role
        let json = "\"\"";
        let role: MessageRole = serde_json::from_str(json).unwrap();
        pretty_assert_eq!(role, MessageRole::Unknown("".to_string()));

        // Invalid role (non-string)
        let json = "123";
        let result: Result<MessageRole, _> = serde_json::from_str(json);
        assert!(result.is_err());

        // Null role
        let json = "null";
        let result: Result<MessageRole, _> = serde_json::from_str(json);
        assert!(result.is_err());
    }
}