Coverage for chatgpt_proxy / db / models.py: 99%

1# MIT License 

2# 

3# Copyright (c) 2025 Tuomo Kriikkula 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23"""Read-only models for chatgpt_proxy's 'adhoc ORM'. 

24 

25TODO: consider replacing these with asyncpg.Record 

26 derived classes? 

27""" 

28 

29import ast 

30import datetime 

31import ipaddress 

32from dataclasses import dataclass 

33from enum import StrEnum 

34 

35max_ast_literal_eval_size = 1000 

36 

37 

38# Mirrored in ChatGPTBotsMutator.uc! 

39class SayType(StrEnum): 

40 ALL = "0" 

41 TEAM = "1" 

42 

43 

44# TODO: for LLM prompting purposes, we could have some 

45# type to refer to US Army, PAVN, NLF, etc.! 

46class Team(StrEnum): 

47 North = "0" 

48 South = "1" 

49 Neutral = "3" 

50 

51 

52@dataclass(slots=True, frozen=True) 

53class Game: 

54 id: str 

55 level: str 

56 start_time: datetime.datetime 

57 game_server_address: ipaddress.IPv4Address 

58 game_server_port: int 

59 stop_time: datetime.datetime | None = None 

60 openai_previous_response_id: str | None = None 

61 

62 

63@dataclass(slots=True, frozen=True) 

64class GamePlayer: 

65 game_id: str 

66 id: int 

67 name: str 

68 team: Team 

69 score: int 

70 

71 def as_markdown_dict(self) -> dict: 

72 return { 

73 "Name": self.name, 

74 "Team:": self.team.name, 

75 "Score:": self.score, 

76 } 

77 

78 def wire_format(self) -> str: 

79 return f"{self.name}\n{self.team}\n{self.score}" 

80 

81 

82@dataclass(slots=True, frozen=True) 

83class OpenAIQuery: 

84 time: datetime.datetime 

85 game_id: str 

86 game_server_address: ipaddress.IPv4Address 

87 game_server_port: int 

88 request_length: int 

89 response_length: int 

90 openai_response_id: str 

91 

92 

93@dataclass(slots=True, frozen=True) 

94class GameChatMessage: 

95 id: int 

96 message: str 

97 game_id: str 

98 send_time: datetime.datetime 

99 sender_name: str 

100 sender_team: Team 

101 channel: SayType 

102 

103 def as_markdown_dict(self) -> dict: 

104 return { 

105 "Message": self.message, 

106 "Sender:": self.sender_name, 

107 "Team:": self.sender_team, 

108 "Channel:": self.channel, 

109 } 

110 

111 def wire_format(self) -> str: 

112 return f"{self.sender_name}\n{self.sender_team}\n{self.channel}\n{self.message}" 

113 

114 

115@dataclass(slots=True, frozen=True) 

116class GameObjective: 

117 name: str 

118 team_state: Team 

119 

120 # TODO: this is wrong. 

121 # - maybe add custom __str__ 

122 def wire_format(self) -> str: 

123 return f"('{self.name}',{int(self.team_state)})" 

124 

125 

126@dataclass(slots=True, frozen=True) 

127class GameObjectiveState: 

128 game_id: str 

129 objectives: list[GameObjective] 

130 

131 def wire_format(self) -> str: 

132 # TODO: too many quotes with this method! 

133 return str([(obj.name, int(obj.team_state)) for obj in self.objectives]) 

134 

135 @staticmethod 

136 def from_wire_format( 

137 game_id: str, 

138 wire_format_data: str, 

139 ) -> "GameObjectiveState": 

140 if len(wire_format_data) > max_ast_literal_eval_size: 

141 raise ValueError("wire_format_data too long") 

142 

143 raw_objs: list[tuple[str, int]] = ast.literal_eval(wire_format_data) 

144 

145 t = type(raw_objs) 

146 if t is not list: 

147 raise ValueError(f"objs: expected list type, got {t}") 

148 

149 objs = [] 

150 for obj_name, obj_state in raw_objs: 

151 if type(obj_name) is not str: 

152 raise ValueError(f"obj_name: expected str type, got {type(obj_name)}") 

153 if type(obj_state) is not int: 

154 raise ValueError(f"obj_state: expected int type, got {type(obj_state)}") 

155 

156 obj_state_enum = Team(str(obj_state)) 

157 

158 objs.append(GameObjective( 

159 name=obj_name, 

160 team_state=obj_state_enum, 

161 )) 

162 

163 return GameObjectiveState( 

164 game_id=game_id, 

165 objectives=objs, 

166 ) 

167 

168 

169@dataclass(slots=True, frozen=True) 

170class GameKill: 

171 id: int 

172 game_id: str 

173 kill_time: datetime.datetime 

174 killer_name: str 

175 victim_name: str 

176 killer_team: Team 

177 victim_team: Team 

178 damage_type: str 

179 kill_distance_m: float 

180 

181 def as_markdown_dict(self) -> dict: 

182 # TODO: also do this for other well known damage type prefixes? 

183 dmg_type = self.damage_type.replace("RODmgType_", "") 

184 return { 

185 "Killer": self.killer_name, 

186 "Victim": self.victim_name, 

187 "Killer Team:": self.killer_team, 

188 "Victim Team:": self.victim_team, 

189 "Damage Type:": dmg_type, 

190 "Kill Distance (m):": round(self.kill_distance_m, 1), 

191 }