openai_tools/common/
role.rs

1use serde::{Deserialize, Serialize};
2use strum_macros::Display;
3
4#[derive(Display, Debug, Clone, Deserialize, Serialize, PartialEq)]
5pub enum Role {
6    #[serde(rename = "system")]
7    #[strum(to_string = "system")]
8    System,
9    #[serde(rename = "user")]
10    #[strum(to_string = "user")]
11    User,
12    #[serde(rename = "assistant")]
13    #[strum(to_string = "assistant")]
14    Assistant,
15    #[serde(rename = "function")]
16    #[strum(to_string = "function")]
17    Function,
18    #[serde(rename = "tool")]
19    #[strum(to_string = "tool")]
20    Tool,
21}
22
23impl TryFrom<String> for Role {
24    type Error = &'static str;
25
26    fn try_from(role: String) -> Result<Self, Self::Error> {
27        let role = role.to_lowercase();
28        match role.as_str() {
29            "system" => Ok(Role::System),
30            "user" => Ok(Role::User),
31            "assistant" => Ok(Role::Assistant),
32            "function" => Ok(Role::Function),
33            "tool" => Ok(Role::Tool),
34            _ => Err("Unknown role"),
35        }
36    }
37}
38
39impl Role {
40    pub fn as_str(&self) -> &str {
41        match self {
42            Role::System => "system",
43            Role::User => "user",
44            Role::Assistant => "assistant",
45            Role::Function => "function",
46            Role::Tool => "tool",
47        }
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54    #[test]
55    fn test_role_conversion() {
56        assert_eq!(Role::try_from("system".to_string()).unwrap(), Role::System);
57        assert_eq!(Role::try_from("user".to_string()).unwrap(), Role::User);
58        assert_eq!(Role::try_from("assistant".to_string()).unwrap(), Role::Assistant);
59        assert_eq!(Role::try_from("function".to_string()).unwrap(), Role::Function);
60        assert_eq!(Role::try_from("tool".to_string()).unwrap(), Role::Tool);
61        assert!(Role::try_from("unknown".to_string()).is_err());
62    }
63
64    #[test]
65    fn test_role_as_str() {
66        assert_eq!(Role::System.as_str(), "system");
67        assert_eq!(Role::User.as_str(), "user");
68        assert_eq!(Role::Assistant.as_str(), "assistant");
69        assert_eq!(Role::Function.as_str(), "function");
70        assert_eq!(Role::Tool.as_str(), "tool");
71    }
72}
openai_tools/common/role.rs

openai_tools/common/
role.rs
