Coverage for src/appl/role_changer.py: 94%

1from typing import Any, Optional 

2 

3from typing_extensions import override 

4 

5from .core.context import PromptContext 

6from .core.modifiers import PrinterModifier, PrinterPush 

7from .core.types import MessageRole, MessageRoleType 

8 

9 

10class RoleChanger(PrinterModifier): 

11 """The contextual role changer of the prompts.""" 

12 

13 _new_role: Optional[MessageRole] = None 

14 

15 def __init__( 

16 self, role: Optional[MessageRole] = None, _ctx: Optional[PromptContext] = None 

17 ): 

18 """Initialize the RoleChanger object. 

19 

20 Args: 

21 role: The new role of the prompts. Defaults to None. 

22 _ctx: The prompt context filled automatically by the APPL function. 

23 """ 

24 super().__init__(_ctx) 

25 if role is not None: 

26 self._new_role = role 

27 

28 @override 

29 @property 

30 def push_args(self) -> PrinterPush: 

31 return PrinterPush(new_role=self._new_role) 

32 

33 

34class SystemRole(RoleChanger): 

35 """Change the role of the prompts to system.""" 

36 

37 def __init__(self, name: Optional[str] = None, **kwargs: Any): 

38 """Initialize the SystemRole object. 

39 

40 Args: 

41 name: The name of the system role. Defaults to None. 

42 **kwargs: The keyword arguments to pass to the RoleChanger constructor. 

43 """ 

44 role = MessageRole(MessageRoleType.SYSTEM, name=name) 

45 super().__init__(role=role, **kwargs) 

46 

47 

48class UserRole(RoleChanger): 

49 """Change the role of the prompts to user.""" 

50 

51 def __init__(self, name: Optional[str] = None, **kwargs: Any): 

52 """Initialize the UserRole object. 

53 

54 Args: 

55 name: The name of the user role. Defaults to None. 

56 **kwargs: The keyword arguments to pass to the RoleChanger constructor. 

57 """ 

58 role = MessageRole(MessageRoleType.USER, name=name) 

59 super().__init__(role=role, **kwargs) 

60 

61 

62class AIRole(RoleChanger): 

63 """Change the role of the prompts to assistant.""" 

64 

65 def __init__(self, name: Optional[str] = None, **kwargs: Any): 

66 """Initialize the AIRole object. 

67 

68 Args: 

69 name: The name of the assistant role. Defaults to None. 

70 **kwargs: The keyword arguments to pass to the Role 

71 """ 

72 role = MessageRole(MessageRoleType.ASSISTANT, name=name) 

73 super().__init__(role=role, **kwargs) 

74 

75 

76class ToolRole(RoleChanger): 

77 """Change the role of the prompts to tool.""" 

78 

79 def __init__(self, name: Optional[str] = None, **kwargs: Any): 

80 """Initialize the ToolRole object. 

81 

82 Args: 

83 name: The name of the tool role. Defaults to None. 

84 **kwargs: The keyword arguments to pass to the RoleChanger constructor. 

85 """ 

86 role = MessageRole(MessageRoleType.TOOL, name=name) 

87 super().__init__(role=role, **kwargs)