Coverage for src/appl/core/types/role.py: 100%

1from enum import Enum 

2 

3from .basic import * 

4 

5 

6class MessageRoleType(str, Enum): 

7 """The type of the role.""" 

8 

9 SYSTEM = "system" 

10 USER = "user" 

11 ASSISTANT = "assistant" 

12 TOOL = "tool" 

13 

14 

15class MessageRole(BaseModel): 

16 """The role of the message owner.""" 

17 

18 type: Optional[str] = None 

19 name: Optional[str] = None 

20 

21 def __init__( 

22 self, 

23 type: Optional[Union[str, MessageRoleType]] = None, 

24 name: Optional[str] = None, 

25 ): 

26 """Initialize the MessageRole object. 

27 

28 Args: 

29 type: The type of the role. 

30 name: An optional name for the role, differentiate between roles of the same type." 

31 """ 

32 if isinstance(type, MessageRoleType): 

33 type = type.value 

34 super().__init__(type=type, name=name) 

35 

36 @property 

37 def is_system(self) -> bool: 

38 """Whether the role is a system role.""" 

39 return self.type == MessageRoleType.SYSTEM 

40 

41 @property 

42 def is_user(self) -> bool: 

43 """Whether the role is a user role.""" 

44 return self.type == MessageRoleType.USER 

45 

46 @property 

47 def is_assistant(self) -> bool: 

48 """Whether the role is an assistant role.""" 

49 return self.type == MessageRoleType.ASSISTANT 

50 

51 @property 

52 def is_tool(self) -> bool: 

53 """Whether the role is a tool role.""" 

54 return self.type == MessageRoleType.TOOL 

55 

56 def get_dict(self) -> Dict[str, Any]: 

57 """Get the role as a dictionary.""" 

58 data = {"role": self.type} 

59 if self.name: 

60 data["name"] = self.name 

61 return data 

62 

63 def __str__(self) -> str: 

64 s = str(self.type) 

65 if self.name: 

66 s += f"({self.name})" 

67 return s 

68 

69 def __eq__(self, other: object) -> bool: 

70 if isinstance(other, MessageRole): 

71 return self.type == other.type and self.name == other.name 

72 return False 

73 

74 

75SYSTEM_ROLE = MessageRole(MessageRoleType.SYSTEM) 

76"""The system role with name not specified.""" 

77USER_ROLE = MessageRole(MessageRoleType.USER) 

78"""The user role with name not specified.""" 

79ASSISTANT_ROLE = MessageRole(MessageRoleType.ASSISTANT) 

80"""The assistant role with name not specified.""" 

81TOOL_ROLE = MessageRole(MessageRoleType.TOOL) 

82"""The tool role with name not specified."""