Coverage for baibu/chat/services.py: 100%
365 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-10-06 13:07 +0000
1"""Conversations, sends, runs and the audit trail.
3A send creates the user's message and a queued Run in one transaction; the
4Run is executed by a Celery task once committed. Sends are idempotent: the
5same key returns the same message and run. A conversation has at most one
6active run at a time. A failed run can be retried, which creates a new run
7for the same message, up to ``CHAT_MAX_ATTEMPTS`` per message.
9A voice send stores the recording, creates the user's message empty and
10queues transcription; the transcript then becomes the message and the reply
11is queued as for a typed message. While transcription is pending the
12conversation is busy, as it is while a reply is being written.
13"""
15import json
16import logging
17import time
18import uuid
19from datetime import timedelta
21from django.conf import settings
22from django.db import transaction
23from django.template import Context
24from django.template import Template
25from django.template.loader import get_template
26from django.utils import timezone
27from django.utils.translation import get_language_info
29from baibu.core import storage
30from baibu.users.consent import consent_state
32from . import llm
33from . import speech
34from . import tools
35from .models import AudioClip
36from .models import Conversation
37from .models import ConversationEvent
38from .models import Message
39from .models import ModelVariant
40from .models import Prompt
41from .models import Run
42from .models import ToolInvocation
44logger = logging.getLogger(__name__)
47class ChatError(Exception):
48 """A request the user can correct; the message is safe to show."""
51class ConversationBusyError(ChatError):
52 pass
55class RetryNotAllowedError(ChatError):
56 pass
59BUSY_MESSAGE = "Please wait for the reply to your last message."
60AUDIO_AREA = "chat/audio"
63def record_event(*, conversation, type: str, run=None, message=None, **payload) -> ConversationEvent: # noqa: A002
64 return ConversationEvent.objects.create(
65 conversation=conversation, run=run, message=message, type=type, payload=payload
66 )
69def start_conversation(*, user, language_code: str = "") -> Conversation:
70 tier = consent_state(user=user, scope=settings.CHAT_CONSENT_SCOPE).tier
71 conversation = Conversation.objects.create(user=user, language_code=language_code, consent_tier=tier)
72 record_event(conversation=conversation, type="conversation_started", consent_tier=tier)
73 return conversation
76def _enqueue(run: Run) -> None:
77 from .tasks import execute_run
79 transaction.on_commit(lambda: execute_run.delay(str(run.pk)))
82def is_busy(conversation: Conversation) -> bool:
83 """A reply is being written, or a voice message is being transcribed."""
84 return (
85 conversation.runs.filter(status__in=Run.ACTIVE_STATUSES).exists()
86 or conversation.audio_clips.filter(status__in=AudioClip.ACTIVE_STATUSES).exists()
87 )
90def _queue_reply(conversation: Conversation, message: Message, *, via: str = "text") -> Run:
91 """Queue the first run for a user message whose content is final. Call
92 inside a transaction holding the conversation's row lock."""
93 run = Run.objects.create(
94 conversation=conversation, triggering_message=message, idempotency_key=f"{message.idempotency_key}:1"
95 )
96 record_event(
97 conversation=conversation, type="message_sent", message=message, characters=len(message.content), via=via
98 )
99 record_event(conversation=conversation, type="run_queued", run=run, message=message, attempt=1)
100 if not conversation.title:
101 conversation.title = message.content.splitlines()[0][:80]
102 conversation.last_activity_at = timezone.now()
103 conversation.save(update_fields=["title", "last_activity_at"])
104 _enqueue(run)
105 return run
108def send_message(*, conversation: Conversation, content: str, idempotency_key: str) -> tuple[Message, Run]:
109 """Add the user's message and queue a reply. Safe to call twice with one key."""
110 content = content.strip()
111 with transaction.atomic():
112 conversation = Conversation.objects.select_for_update().get(pk=conversation.pk)
113 existing = conversation.messages.filter(idempotency_key=idempotency_key).first()
114 if existing is not None:
115 return existing, existing.runs.order_by("-created_at").first()
116 if is_busy(conversation):
117 raise ConversationBusyError(BUSY_MESSAGE)
118 message = Message.objects.create(
119 conversation=conversation, role=Message.Role.USER, content=content, idempotency_key=idempotency_key
120 )
121 run = _queue_reply(conversation, message)
122 return message, run
125def retry_run(run: Run) -> Run:
126 """Queue a new attempt at a failed run's reply."""
127 with transaction.atomic():
128 conversation = Conversation.objects.select_for_update().get(pk=run.conversation_id)
129 run = Run.objects.get(pk=run.pk)
130 if run.status != Run.Status.FAILED:
131 msg = "Only a failed reply can be retried."
132 raise RetryNotAllowedError(msg)
133 latest = run.triggering_message.runs.order_by("-attempt").first()
134 if latest.pk != run.pk:
135 return latest
136 if run.attempt >= settings.CHAT_MAX_ATTEMPTS:
137 msg = "This message has been tried too many times."
138 raise RetryNotAllowedError(msg)
139 if is_busy(conversation):
140 raise ConversationBusyError(BUSY_MESSAGE)
141 attempt = run.attempt + 1
142 new_run = Run.objects.create(
143 conversation=conversation,
144 triggering_message=run.triggering_message,
145 attempt=attempt,
146 idempotency_key=f"{run.idempotency_key.rsplit(':', 1)[0]}:{attempt}",
147 )
148 record_event(
149 conversation=conversation, type="run_retried", run=new_run, retried_run=str(run.pk), attempt=attempt
150 )
151 conversation.last_activity_at = timezone.now()
152 conversation.save(update_fields=["last_activity_at"])
153 _enqueue(new_run)
154 return new_run
157# --- Prompt and model -------------------------------------------------------
160def active_prompt(name: str = "system") -> Prompt | None:
161 return Prompt.objects.filter(name=name, active=True).first()
164def render_system_prompt(conversation: Conversation, prompt: Prompt | None) -> str:
165 from baibu.core.context_processors import platform
167 language_code = conversation.language_code or settings.LANGUAGE_CODE
168 try:
169 language_name = get_language_info(language_code)["name"]
170 except KeyError:
171 language_name = language_code
172 context = {
173 **platform(None),
174 "user_name": conversation.user.name,
175 "language_code": language_code,
176 "language_name": language_name,
177 }
178 if prompt is None:
179 return get_template("chat/system_prompt.txt").render(context).strip()
180 return Template(prompt.body).render(Context(context, autoescape=False)).strip()
183def default_variant() -> ModelVariant | None:
184 return ModelVariant.objects.filter(is_default=True).first()
187def build_messages(run: Run, system_prompt: str) -> list[dict]:
188 """System prompt, then recent turns up to and including the triggering message."""
189 history = list(
190 run.conversation.messages.filter(created_at__lte=run.triggering_message.created_at)
191 # Voice messages without a transcript have no content to send.
192 .exclude(audio__status__in=[*AudioClip.ACTIVE_STATUSES, AudioClip.Status.FAILED])
193 .order_by("-created_at", "-id")
194 .values("role", "content")[: settings.CHAT_CONTEXT_MESSAGES]
195 )
196 history.reverse()
197 return [{"role": "system", "content": system_prompt}, *history]
200# --- Execution --------------------------------------------------------------
203def claim(run_id) -> Run | None:
204 """Move a queued run to running. Returns None if someone else has it."""
205 now = timezone.now()
206 claimed = Run.objects.filter(pk=run_id, status=Run.Status.QUEUED).update(status=Run.Status.RUNNING, started_at=now)
207 if not claimed:
208 return None
209 run = Run.objects.select_related("conversation__user", "triggering_message").get(pk=run_id)
210 record_event(conversation=run.conversation, type="run_started", run=run)
211 return run
214def execute(run_id) -> Run | None:
215 run = claim(run_id)
216 if run is None:
217 return None
218 prompt = active_prompt()
219 variant = default_variant()
220 config = llm.ModelConfig.from_variant(variant)
221 run.prompt = prompt
222 run.model_variant = variant
223 run.prompt_version = str(prompt) if prompt else "built-in"
224 try:
225 messages = build_messages(run, render_system_prompt(run.conversation, prompt))
226 completion = generate(run, config, messages)
227 except llm.ModelError as exc:
228 return fail(run, code=exc.code, message=str(exc))
229 except Exception as exc:
230 logger.exception("Chat run %s failed", run.pk)
231 return fail(run, code="internal_error", message=type(exc).__name__)
232 return complete(run, completion)
235def generate(run: Run, config: llm.ModelConfig, messages: list[dict]) -> llm.Completion:
236 """Ask the model for the reply, running any tools it calls.
238 The model may call tools for up to ``CHAT_MAX_TOOL_ROUNDS`` rounds; the
239 last call offers no tools, so it has to answer. Sources from tools are
240 collected on ``run.sources`` for the reply.
241 """
242 run.model_name = config.model
243 run.sources = []
244 available = {tool.name: tool for tool in tools.available_tools()}
245 messages = list(messages)
246 for round_number in range(settings.CHAT_MAX_TOOL_ROUNDS + 1):
247 offer = [tool.schema() for tool in available.values()] if round_number < settings.CHAT_MAX_TOOL_ROUNDS else []
248 completion = llm.complete(config, messages, tools=offer or None)
249 if not completion.tool_calls:
250 return completion
251 messages.append(completion.raw_message)
252 for call in completion.tool_calls:
253 messages.append({"role": "tool", "tool_call_id": call.id, "content": _run_tool(run, available, call)})
254 return completion # pragma: no cover - the last round offers no tools
257def _run_tool(run: Run, available: dict, call: llm.ToolCall) -> str:
258 started = time.monotonic()
259 tool = available.get(call.name)
260 invocation = ToolInvocation(run=run, tool_name=call.name[:64], call_id=call.id[:255], arguments=call.arguments)
261 try:
262 if tool is None:
263 msg = f"Unknown tool {call.name!r}."
264 raise tools.ToolError(msg)
265 result = tool.run(call.arguments)
266 except tools.ToolError as exc:
267 invocation.status = ToolInvocation.Status.FAILED
268 invocation.error = str(exc)[:500]
269 content = json.dumps({"error": invocation.error})
270 else:
271 invocation.status = ToolInvocation.Status.COMPLETED
272 invocation.provider = result.provider
273 invocation.result_count = len(result.sources)
274 run.sources.extend(s for s in result.sources if s not in run.sources)
275 content = result.content
276 invocation.latency_ms = int((time.monotonic() - started) * 1000)
277 invocation.save()
278 record_event(
279 conversation=run.conversation,
280 type=f"tool_{invocation.status}",
281 run=run,
282 tool=invocation.tool_name,
283 invocation=str(invocation.pk),
284 )
285 return content
288def complete(run: Run, completion: llm.Completion) -> Run:
289 content = completion.content.strip()
290 if not content:
291 return fail(run, code="empty_response", message="The model returned an empty reply.")
292 with transaction.atomic():
293 locked = Run.objects.select_for_update().get(pk=run.pk)
294 if locked.status != Run.Status.RUNNING:
295 # Timed out and failed by the sweeper meanwhile; keep that outcome.
296 return locked
297 reply = Message.objects.create(
298 conversation=run.conversation,
299 role=Message.Role.ASSISTANT,
300 content=content,
301 run=run,
302 sources=getattr(run, "sources", [])[: settings.CHAT_SEARCH_MAX_RESULTS * settings.CHAT_MAX_TOOL_ROUNDS],
303 )
304 run.status = Run.Status.COMPLETED
305 run.model_name = completion.model or run.model_name
306 run.usage = completion.usage
307 run.latency_ms = completion.latency_ms
308 run.completed_at = timezone.now()
309 run.save()
310 record_event(
311 conversation=run.conversation,
312 type="run_completed",
313 run=run,
314 message=reply,
315 model=run.model_name,
316 prompt=run.prompt_version,
317 latency_ms=completion.latency_ms,
318 )
319 Conversation.objects.filter(pk=run.conversation_id).update(last_activity_at=timezone.now())
320 return run
323def fail(run: Run, *, code: str, message: str) -> Run:
324 with transaction.atomic():
325 locked = Run.objects.select_for_update().get(pk=run.pk)
326 if locked.status in (Run.Status.COMPLETED, Run.Status.FAILED):
327 return locked
328 run.status = Run.Status.FAILED
329 run.error = {"code": code, "message": message[:500]}
330 run.completed_at = timezone.now()
331 run.save()
332 record_event(conversation=run.conversation, type="run_failed", run=run, code=code)
333 logger.warning("Chat run %s failed: %s", run.pk, code)
334 return run
337def sweep_runs(now=None) -> dict:
338 """Fail runs stuck too long and re-send queued runs whose task was lost."""
339 now = now or timezone.now()
340 timeout = timedelta(seconds=settings.CHAT_RUN_TIMEOUT_SECONDS)
341 stats = {"failed": 0, "requeued": 0}
342 stuck = Run.objects.filter(status=Run.Status.RUNNING, started_at__lt=now - timeout).select_related("conversation")
343 for run in stuck:
344 if fail(run, code="timeout", message="No reply in time.").error.get("code") == "timeout":
345 stats["failed"] += 1
346 lost = Run.objects.filter(status=Run.Status.QUEUED, created_at__lt=now - timeout).select_related("conversation")
347 for run in lost:
348 if run.created_at < now - 3 * timeout:
349 fail(run, code="not_started", message="The reply was never started.")
350 stats["failed"] += 1
351 else:
352 from .tasks import execute_run
354 execute_run.delay(str(run.pk))
355 stats["requeued"] += 1
356 return stats
359# --- Voice messages ---------------------------------------------------------
362def _enqueue_transcription(clip: AudioClip) -> None:
363 from .tasks import transcribe_audio
365 transaction.on_commit(lambda: transcribe_audio.delay(str(clip.pk)))
368def send_voice(*, conversation: Conversation, audio: bytes, content_type: str, idempotency_key: str) -> Message:
369 """Store a recording, add an empty user message and queue transcription.
371 Safe to call twice with one key. Storage errors are let through and
372 nothing is recorded then.
373 """
374 with transaction.atomic():
375 conversation = Conversation.objects.select_for_update().get(pk=conversation.pk)
376 existing = conversation.messages.filter(idempotency_key=idempotency_key).first()
377 if existing is not None:
378 return existing
379 if is_busy(conversation):
380 raise ConversationBusyError(BUSY_MESSAGE)
381 clip_id = uuid.uuid7()
382 key = storage.save_bytes(
383 storage.build_key(AUDIO_AREA, f"{clip_id}.{speech.extension_for(content_type)}"), audio
384 )
385 try:
386 with transaction.atomic():
387 message = Message.objects.create(
388 conversation=conversation, role=Message.Role.USER, content="", idempotency_key=idempotency_key
389 )
390 clip = AudioClip.objects.create(
391 id=clip_id,
392 conversation=conversation,
393 message=message,
394 storage_key=key,
395 content_type=content_type,
396 size_bytes=len(audio),
397 )
398 record_event(
399 conversation=conversation,
400 type="audio_received",
401 message=message,
402 clip=str(clip.pk),
403 content_type=content_type,
404 bytes=len(audio),
405 )
406 conversation.last_activity_at = timezone.now()
407 conversation.save(update_fields=["last_activity_at"])
408 except Exception:
409 storage.delete(key)
410 raise
411 _enqueue_transcription(clip)
412 return message
415def retry_transcription(clip: AudioClip) -> AudioClip:
416 """Queue another attempt at a failed transcription."""
417 with transaction.atomic():
418 conversation = Conversation.objects.select_for_update().get(pk=clip.conversation_id)
419 clip = AudioClip.objects.get(pk=clip.pk)
420 if clip.status != AudioClip.Status.FAILED:
421 msg = "Only a failed transcription can be retried."
422 raise RetryNotAllowedError(msg)
423 if clip.attempt >= settings.CHAT_MAX_ATTEMPTS:
424 msg = "This message has been tried too many times."
425 raise RetryNotAllowedError(msg)
426 if is_busy(conversation):
427 raise ConversationBusyError(BUSY_MESSAGE)
428 clip.status = AudioClip.Status.PENDING
429 clip.attempt += 1
430 clip.error = {}
431 clip.queued_at = timezone.now()
432 clip.started_at = clip.completed_at = None
433 clip.save()
434 record_event(
435 conversation=conversation,
436 type="transcription_retried",
437 message=clip.message,
438 clip=str(clip.pk),
439 attempt=clip.attempt,
440 )
441 conversation.last_activity_at = timezone.now()
442 conversation.save(update_fields=["last_activity_at"])
443 _enqueue_transcription(clip)
444 return clip
447def transcribe(clip_id) -> AudioClip | None:
448 """Transcribe a pending clip; on success queue the reply."""
449 claimed = AudioClip.objects.filter(pk=clip_id, status=AudioClip.Status.PENDING).update(
450 status=AudioClip.Status.TRANSCRIBING, started_at=timezone.now()
451 )
452 if not claimed:
453 return None
454 clip = AudioClip.objects.select_related("conversation", "message").get(pk=clip_id)
455 provider = speech.get_provider()
456 try:
457 if provider is None:
458 msg = "Voice input is not configured."
459 raise speech.TranscriptionError(msg, code="not_configured")
460 audio = storage.read_bytes(clip.storage_key)
461 if audio is None:
462 msg = "The recording is missing from storage."
463 raise speech.TranscriptionError(msg, code="missing_audio")
464 transcript = provider.transcribe(
465 audio, content_type=clip.content_type, language_code=clip.conversation.language_code
466 )
467 except speech.TranscriptionError as exc:
468 return fail_transcription(clip, code=exc.code, message=str(exc))
469 except Exception as exc:
470 logger.exception("Transcription of %s failed", clip.pk)
471 return fail_transcription(clip, code="internal_error", message=type(exc).__name__)
472 text = " ".join(transcript.text.split())[: settings.CHAT_MAX_MESSAGE_CHARACTERS]
473 if not text:
474 return fail_transcription(clip, code="no_speech", message="No speech was recognised.")
475 with transaction.atomic():
476 conversation = Conversation.objects.select_for_update().get(pk=clip.conversation_id)
477 locked = AudioClip.objects.select_for_update().get(pk=clip.pk)
478 if locked.status != AudioClip.Status.TRANSCRIBING:
479 # Timed out and failed by the sweeper meanwhile; keep that outcome.
480 return locked
481 clip.status = AudioClip.Status.TRANSCRIBED
482 clip.provider = (transcript.provider or getattr(provider, "name", ""))[:128]
483 clip.language = transcript.language[:24]
484 clip.duration_seconds = transcript.duration
485 clip.completed_at = timezone.now()
486 clip.save()
487 message = clip.message
488 message.content = text
489 message.save(update_fields=["content"])
490 record_event(
491 conversation=conversation,
492 type="transcription_completed",
493 message=message,
494 clip=str(clip.pk),
495 provider=clip.provider,
496 characters=len(text),
497 duration_seconds=clip.duration_seconds,
498 )
499 _queue_reply(conversation, message, via="voice")
500 return clip
503def fail_transcription(clip: AudioClip, *, code: str, message: str) -> AudioClip:
504 with transaction.atomic():
505 locked = AudioClip.objects.select_for_update().get(pk=clip.pk)
506 if locked.status in (AudioClip.Status.TRANSCRIBED, AudioClip.Status.FAILED):
507 return locked
508 clip.status = AudioClip.Status.FAILED
509 clip.error = {"code": code, "message": message[:500]}
510 clip.completed_at = timezone.now()
511 clip.save()
512 record_event(
513 conversation=clip.conversation,
514 type="transcription_failed",
515 message=clip.message,
516 clip=str(clip.pk),
517 code=code,
518 )
519 logger.warning("Transcription of %s failed: %s", clip.pk, code)
520 return clip
523def sweep_transcriptions(now=None) -> dict:
524 """Fail transcriptions stuck too long and re-send pending ones whose task was lost."""
525 now = now or timezone.now()
526 timeout = timedelta(seconds=settings.CHAT_RUN_TIMEOUT_SECONDS)
527 stats = {"failed": 0, "requeued": 0}
528 stuck = AudioClip.objects.filter(status=AudioClip.Status.TRANSCRIBING, started_at__lt=now - timeout)
529 for clip in stuck.select_related("conversation"):
530 if fail_transcription(clip, code="timeout", message="No transcript in time.").error.get("code") == "timeout":
531 stats["failed"] += 1
532 lost = AudioClip.objects.filter(status=AudioClip.Status.PENDING, queued_at__lt=now - timeout)
533 for clip in lost.select_related("conversation"):
534 if clip.queued_at < now - 3 * timeout:
535 fail_transcription(clip, code="not_started", message="The transcription was never started.")
536 stats["failed"] += 1
537 else:
538 from .tasks import transcribe_audio
540 transcribe_audio.delay(str(clip.pk))
541 stats["requeued"] += 1
542 return stats
545def new_idempotency_key() -> str:
546 return uuid.uuid4().hex
549# --- Reports ----------------------------------------------------------------
552def report_message(*, message: Message, user, reason: str, note: str = ""):
553 """Flag an assistant reply for staff. Reporting the same reply twice is a no-op."""
554 from .models import ChatFlag
556 if message.role != Message.Role.ASSISTANT or message.conversation.user_id != user.pk:
557 msg = "Only replies in your own conversations can be reported."
558 raise ChatError(msg)
559 flag, created = ChatFlag.objects.get_or_create(
560 message=message,
561 reported_by=user,
562 defaults={"conversation": message.conversation, "reason": reason, "note": note.strip()[:500]},
563 )
564 if created:
565 record_event(conversation=message.conversation, type="reply_reported", message=message, reason=reason)
566 return flag
569def decide_flag(flag, status: str, *, reviewer, note: str = "") -> bool:
570 """Record the one staff decision on a report. Returns False if already decided."""
571 from django.urls import reverse
573 from baibu.notifications.services import notify
575 from .models import ChatFlag
577 if status not in {ChatFlag.Status.CONFIRMED, ChatFlag.Status.DISMISSED}:
578 msg = f"Not a decision: {status}"
579 raise ValueError(msg)
580 with transaction.atomic():
581 locked = ChatFlag.objects.select_for_update().get(pk=flag.pk)
582 if locked.status != ChatFlag.Status.OPEN:
583 return False
584 locked.status = status
585 locked.decided_by = reviewer
586 locked.decided_at = timezone.now()
587 locked.decision_note = note.strip()
588 locked.save()
589 record_event(conversation=locked.conversation, type=f"report_{status}", message=locked.message)
590 notify(
591 locked.reported_by,
592 kind="chat_report_reviewed",
593 link=reverse("chat:conversation", args=[locked.conversation_id]),
594 )
595 flag.refresh_from_db()
596 return True