1use serde::{Deserialize, Serialize};
 2use std::fmt::{self, Display};
 3
 4#[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
 5#[serde(rename_all = "lowercase")]
 6pub enum Role {
 7    User,
 8    Assistant,
 9    System,
10}
11
12impl Role {
13    pub fn from_proto(role: i32) -> Role {
14        match proto::LanguageModelRole::from_i32(role) {
15            Some(proto::LanguageModelRole::LanguageModelUser) => Role::User,
16            Some(proto::LanguageModelRole::LanguageModelAssistant) => Role::Assistant,
17            Some(proto::LanguageModelRole::LanguageModelSystem) => Role::System,
18            Some(proto::LanguageModelRole::LanguageModelTool) => Role::System,
19            None => Role::User,
20        }
21    }
22
23    pub fn to_proto(&self) -> proto::LanguageModelRole {
24        match self {
25            Role::User => proto::LanguageModelRole::LanguageModelUser,
26            Role::Assistant => proto::LanguageModelRole::LanguageModelAssistant,
27            Role::System => proto::LanguageModelRole::LanguageModelSystem,
28        }
29    }
30
31    pub fn cycle(self) -> Role {
32        match self {
33            Role::User => Role::Assistant,
34            Role::Assistant => Role::System,
35            Role::System => Role::User,
36        }
37    }
38}
39
40impl Display for Role {
41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            Role::User => write!(f, "user"),
44            Role::Assistant => write!(f, "assistant"),
45            Role::System => write!(f, "system"),
46        }
47    }
48}
49
50impl From<Role> for ollama::Role {
51    fn from(val: Role) -> Self {
52        match val {
53            Role::User => ollama::Role::User,
54            Role::Assistant => ollama::Role::Assistant,
55            Role::System => ollama::Role::System,
56        }
57    }
58}
59
60impl From<Role> for open_ai::Role {
61    fn from(val: Role) -> Self {
62        match val {
63            Role::User => open_ai::Role::User,
64            Role::Assistant => open_ai::Role::Assistant,
65            Role::System => open_ai::Role::System,
66        }
67    }
68}