Coverage for baibu/chat/models.py: 100%

1import uuid 

2 

3from django.conf import settings 

4from django.db import models 

5from django.utils import timezone 

6from django.utils.translation import gettext_lazy as _ 

7 

8from baibu.users.models import ConsentRecord 

9from baibu.users.models import ImmutableRecordError 

10 

11 

12class ModelVariant(models.Model): 

13 """A language model the assistant can use, reached through LiteLLM. 

14 

15 ``model`` is any model name LiteLLM accepts. The special name ``mock`` 

16 answers locally without a network call, for development and tests. 

17 """ 

18 

19 name = models.CharField(max_length=128, unique=True) 

20 model = models.CharField(max_length=255, help_text=_("LiteLLM model name, or 'mock'.")) 

21 api_base = models.CharField(max_length=512, blank=True, help_text=_("Endpoint URL, if the model needs one.")) 

22 api_key_env = models.CharField( 

23 max_length=128, 

24 blank=True, 

25 help_text=_("Name of the environment variable that holds the API key. Keys are never stored here."), 

26 ) 

27 parameters = models.JSONField( 

28 default=dict, blank=True, help_text=_('Extra parameters for every call, e.g. {"temperature": 0.7}.') 

29 ) 

30 is_default = models.BooleanField(default=False, help_text=_("Used for new replies. Only one can be the default.")) 

31 created_at = models.DateTimeField(auto_now_add=True) 

32 updated_at = models.DateTimeField(auto_now=True) 

33 

34 class Meta: 

35 ordering = ["name"] 

36 constraints = [ 

37 models.UniqueConstraint( 

38 fields=["is_default"], condition=models.Q(is_default=True), name="chat_one_default_model" 

39 ), 

40 ] 

41 verbose_name = _("Model variant") 

42 verbose_name_plural = _("Model variants") 

43 

44 def __str__(self) -> str: 

45 return self.name 

46 

47 def validate_constraints(self, exclude=None): 

48 # The admin moves the default flag when saving, so do not refuse a second default here. 

49 super().validate_constraints(exclude={*(exclude or ()), "is_default"}) 

50 

51 

52class Prompt(models.Model): 

53 """One version of a named prompt. Versions are never edited once saved.""" 

54 

55 name = models.CharField(max_length=64, default="system") 

56 version = models.PositiveIntegerField(editable=False) 

57 body = models.TextField( 

58 help_text=_("A Django template. Available: platform, user_name, language_code, language_name.") 

59 ) 

60 active = models.BooleanField(default=False) 

61 notes = models.CharField(max_length=255, blank=True) 

62 created_by = models.ForeignKey( 

63 settings.AUTH_USER_MODEL, on_delete=models.SET_NULL, null=True, blank=True, related_name="+" 

64 ) 

65 created_at = models.DateTimeField(auto_now_add=True) 

66 

67 class Meta: 

68 ordering = ["name", "-version"] 

69 constraints = [ 

70 models.UniqueConstraint(fields=["name", "version"], name="chat_prompt_name_version"), 

71 models.UniqueConstraint( 

72 fields=["name"], condition=models.Q(active=True), name="chat_one_active_prompt_per_name" 

73 ), 

74 ] 

75 verbose_name = _("Prompt") 

76 verbose_name_plural = _("Prompts") 

77 

78 def __str__(self) -> str: 

79 return f"{self.name} v{self.version}" 

80 

81 def save(self, *args, **kwargs): 

82 if self._state.adding: 

83 latest = Prompt.objects.filter(name=self.name).aggregate(models.Max("version"))["version__max"] 

84 self.version = (latest or 0) + 1 

85 elif Prompt.objects.filter(pk=self.pk).exclude(body=self.body).exists(): 

86 msg = "Prompt versions cannot be edited; save a new version instead." 

87 raise ImmutableRecordError(msg) 

88 super().save(*args, **kwargs) 

89 

90 def validate_constraints(self, exclude=None): 

91 # The admin moves the active flag when saving, so do not refuse a second active version here. 

92 super().validate_constraints(exclude={*(exclude or ()), "active"}) 

93 

94 

95class Conversation(models.Model): 

96 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

97 user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, related_name="conversations") 

98 title = models.CharField(max_length=120, blank=True) 

99 language_code = models.CharField(max_length=24, blank=True) 

100 # Consent under the chat scope when the conversation started. 

101 consent_tier = models.CharField(max_length=32, choices=ConsentRecord.Tier.choices) 

102 created_at = models.DateTimeField(auto_now_add=True) 

103 last_activity_at = models.DateTimeField(default=timezone.now, db_index=True) 

104 

105 class Meta: 

106 ordering = ["-last_activity_at"] 

107 verbose_name = _("Conversation") 

108 verbose_name_plural = _("Conversations") 

109 

110 def __str__(self) -> str: 

111 return self.title or f"Conversation {self.pk}" 

112 

113 @property 

114 def active_run(self): 

115 return self.runs.filter(status__in=Run.ACTIVE_STATUSES).first() 

116 

117 

118class Message(models.Model): 

119 class Role(models.TextChoices): 

120 USER = "user", _("User") 

121 ASSISTANT = "assistant", _("Assistant") 

122 

123 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

124 conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="messages") 

125 role = models.CharField(max_length=16, choices=Role.choices) 

126 content = models.TextField() 

127 # Set on user messages: the same key twice means the same send. NULL on 

128 # assistant messages, so the unique constraint ignores them. 

129 idempotency_key = models.CharField(max_length=64, null=True, blank=True) # noqa: DJ001 

130 # Set on assistant messages: the run that wrote it. 

131 run = models.OneToOneField("chat.Run", on_delete=models.SET_NULL, null=True, blank=True, related_name="reply") 

132 # Set on assistant messages that used tools: [{"title", "url"}]. 

133 sources = models.JSONField(default=list, blank=True) 

134 created_at = models.DateTimeField(auto_now_add=True) 

135 

136 class Meta: 

137 ordering = ["created_at", "id"] 

138 constraints = [ 

139 models.UniqueConstraint(fields=["conversation", "idempotency_key"], name="chat_message_idempotency"), 

140 ] 

141 verbose_name = _("Message") 

142 verbose_name_plural = _("Messages") 

143 

144 def __str__(self) -> str: 

145 return f"{self.role} message {self.pk}" 

146 

147 

148class Run(models.Model): 

149 """One attempt to generate the assistant's reply to a user message.""" 

150 

151 class Status(models.TextChoices): 

152 QUEUED = "queued", _("Queued") 

153 RUNNING = "running", _("Running") 

154 COMPLETED = "completed", _("Completed") 

155 FAILED = "failed", _("Failed") 

156 

157 ACTIVE_STATUSES = (Status.QUEUED, Status.RUNNING) 

158 

159 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

160 conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="runs") 

161 triggering_message = models.ForeignKey(Message, on_delete=models.CASCADE, related_name="runs") 

162 status = models.CharField(max_length=16, choices=Status.choices, default=Status.QUEUED, db_index=True) 

163 idempotency_key = models.CharField(max_length=64, unique=True) 

164 attempt = models.PositiveSmallIntegerField(default=1) 

165 model_variant = models.ForeignKey(ModelVariant, on_delete=models.SET_NULL, null=True, blank=True) 

166 prompt = models.ForeignKey(Prompt, on_delete=models.SET_NULL, null=True, blank=True) 

167 # What was actually used, kept even if the variant or prompt is deleted. 

168 model_name = models.CharField(max_length=255, blank=True) 

169 prompt_version = models.CharField(max_length=80, blank=True) 

170 usage = models.JSONField(default=dict, blank=True) 

171 error = models.JSONField(default=dict, blank=True) 

172 latency_ms = models.PositiveIntegerField(null=True, blank=True) 

173 created_at = models.DateTimeField(auto_now_add=True) 

174 started_at = models.DateTimeField(null=True, blank=True) 

175 completed_at = models.DateTimeField(null=True, blank=True) 

176 

177 class Meta: 

178 ordering = ["created_at", "id"] 

179 verbose_name = _("Run") 

180 verbose_name_plural = _("Runs") 

181 

182 def __str__(self) -> str: 

183 return f"Run {self.pk} ({self.status})" 

184 

185 

186class ToolInvocation(models.Model): 

187 """One tool call made by the model during a run.""" 

188 

189 class Status(models.TextChoices): 

190 COMPLETED = "completed", _("Completed") 

191 FAILED = "failed", _("Failed") 

192 

193 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

194 run = models.ForeignKey(Run, on_delete=models.CASCADE, related_name="tool_invocations") 

195 tool_name = models.CharField(max_length=64, db_index=True) 

196 provider = models.CharField(max_length=64, blank=True) 

197 call_id = models.CharField(max_length=255) 

198 status = models.CharField(max_length=16, choices=Status.choices) 

199 arguments = models.JSONField(default=dict, blank=True) 

200 result_count = models.PositiveIntegerField(null=True, blank=True) 

201 error = models.CharField(max_length=500, blank=True) 

202 latency_ms = models.PositiveIntegerField(default=0) 

203 created_at = models.DateTimeField(auto_now_add=True) 

204 

205 class Meta: 

206 ordering = ["created_at", "id"] 

207 verbose_name = _("Tool invocation") 

208 verbose_name_plural = _("Tool invocations") 

209 

210 def __str__(self) -> str: 

211 return f"{self.tool_name} ({self.status})" 

212 

213 

214class ChatFlag(models.Model): 

215 """A user's report about an assistant reply, for staff to review once.""" 

216 

217 class Reason(models.TextChoices): 

218 HARMFUL = "harmful", _("Harmful or offensive") 

219 INCORRECT = "incorrect", _("Wrong or misleading") 

220 OTHER = "other", _("Something else") 

221 

222 class Status(models.TextChoices): 

223 OPEN = "open", _("Open") 

224 CONFIRMED = "confirmed", _("Problem confirmed") 

225 DISMISSED = "dismissed", _("No problem found") 

226 

227 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

228 conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="flags") 

229 message = models.ForeignKey(Message, on_delete=models.CASCADE, related_name="flags") 

230 reported_by = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, related_name="+") 

231 reason = models.CharField(max_length=16, choices=Reason.choices) 

232 note = models.CharField(max_length=500, blank=True) 

233 status = models.CharField(max_length=16, choices=Status.choices, default=Status.OPEN, db_index=True) 

234 decided_by = models.ForeignKey( 

235 settings.AUTH_USER_MODEL, on_delete=models.SET_NULL, null=True, blank=True, related_name="+" 

236 ) 

237 decided_at = models.DateTimeField(null=True, blank=True) 

238 decision_note = models.TextField(blank=True) 

239 created_at = models.DateTimeField(auto_now_add=True) 

240 

241 class Meta: 

242 ordering = ["created_at"] 

243 constraints = [ 

244 models.UniqueConstraint(fields=["message", "reported_by"], name="chat_flag_once_per_message"), 

245 ] 

246 verbose_name = _("Reported reply") 

247 verbose_name_plural = _("Reported replies") 

248 

249 def __str__(self) -> str: 

250 return f"{self.get_reason_display()} ({self.get_status_display()})" 

251 

252 

253class AudioClip(models.Model): 

254 """A voice message: the recording in private storage and its transcription. 

255 

256 The user's message is created empty when the audio arrives; the 

257 transcript becomes its content once transcription succeeds. 

258 """ 

259 

260 class Status(models.TextChoices): 

261 PENDING = "pending", _("Pending") 

262 TRANSCRIBING = "transcribing", _("Transcribing") 

263 TRANSCRIBED = "transcribed", _("Transcribed") 

264 FAILED = "failed", _("Failed") 

265 

266 ACTIVE_STATUSES = (Status.PENDING, Status.TRANSCRIBING) 

267 

268 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

269 conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="audio_clips") 

270 message = models.OneToOneField(Message, on_delete=models.CASCADE, related_name="audio") 

271 storage_key = models.CharField(max_length=255) 

272 content_type = models.CharField(max_length=64) 

273 size_bytes = models.PositiveIntegerField() 

274 duration_seconds = models.FloatField(null=True, blank=True) 

275 status = models.CharField(max_length=16, choices=Status.choices, default=Status.PENDING, db_index=True) 

276 attempt = models.PositiveSmallIntegerField(default=1) 

277 provider = models.CharField(max_length=128, blank=True) 

278 language = models.CharField(max_length=24, blank=True) 

279 error = models.JSONField(default=dict, blank=True) 

280 created_at = models.DateTimeField(auto_now_add=True) 

281 queued_at = models.DateTimeField(default=timezone.now) 

282 started_at = models.DateTimeField(null=True, blank=True) 

283 completed_at = models.DateTimeField(null=True, blank=True) 

284 

285 class Meta: 

286 ordering = ["created_at", "id"] 

287 verbose_name = _("Audio clip") 

288 verbose_name_plural = _("Audio clips") 

289 

290 def __str__(self) -> str: 

291 return f"Audio {self.pk} ({self.status})" 

292 

293 

294class Topic(models.Model): 

295 """A conversation topic, defined by the deployment, for optional tagging.""" 

296 

297 slug = models.SlugField(max_length=64, unique=True) 

298 label = models.CharField(max_length=128) 

299 description = models.CharField( 

300 max_length=500, blank=True, help_text=_("Helps the tagger decide; not shown to users.") 

301 ) 

302 active = models.BooleanField(default=True) 

303 

304 class Meta: 

305 ordering = ["label"] 

306 verbose_name = _("Topic") 

307 verbose_name_plural = _("Topics") 

308 

309 def __str__(self) -> str: 

310 return self.label 

311 

312 

313class ConversationTag(models.Model): 

314 """The latest tagging of a conversation: topics, language and intent.""" 

315 

316 conversation = models.OneToOneField(Conversation, on_delete=models.CASCADE, related_name="tag") 

317 topics = models.ManyToManyField(Topic, blank=True, related_name="tags") 

318 language = models.CharField(max_length=24, blank=True) 

319 intent = models.CharField(max_length=32, blank=True) 

320 tagger = models.CharField(max_length=255, blank=True) 

321 # Conversation activity covered by this tagging; newer activity means re-tag. 

322 tagged_activity_at = models.DateTimeField() 

323 tagged_at = models.DateTimeField(auto_now=True) 

324 error = models.CharField(max_length=500, blank=True) 

325 

326 class Meta: 

327 verbose_name = _("Conversation tag") 

328 verbose_name_plural = _("Conversation tags") 

329 

330 def __str__(self) -> str: 

331 return f"Tags for {self.conversation_id}" 

332 

333 

334class ConversationEvent(models.Model): 

335 """Audit trail: what happened in a conversation and when. Append-only.""" 

336 

337 id = models.UUIDField(primary_key=True, default=uuid.uuid7, editable=False) 

338 conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name="events") 

339 run = models.ForeignKey(Run, on_delete=models.SET_NULL, null=True, blank=True, related_name="events") 

340 message = models.ForeignKey(Message, on_delete=models.SET_NULL, null=True, blank=True, related_name="events") 

341 type = models.CharField(max_length=64, db_index=True) 

342 payload = models.JSONField(default=dict, blank=True) 

343 created_at = models.DateTimeField(auto_now_add=True) 

344 

345 class Meta: 

346 ordering = ["created_at", "id"] 

347 verbose_name = _("Conversation event") 

348 verbose_name_plural = _("Conversation events") 

349 

350 def __str__(self) -> str: 

351 return f"{self.type} ({self.created_at:%Y-%m-%d %H:%M})" 

352 

353 def save(self, *args, **kwargs): 

354 if not self._state.adding: 

355 msg = "Conversation events are append-only." 

356 raise ImmutableRecordError(msg) 

357 super().save(*args, **kwargs)