satoko-chatbot/voice_stt.py

"""Réception Discord : une phrase par personne, envoyée au serveur Whisper."""
import asyncio
import audioop
import logging
import threading
import time
import wave
from io import BytesIO
from collections import Counter

import discord
import httpx
from davey import MediaType
from discord.ext import voice_recv

logger = logging.getLogger(__name__)


class STTSink(voice_recv.AudioSink):
    def __init__(self, on_phrase, stt_url):
        super().__init__()
        self.on_phrase = on_phrase
        self.stt_url = stt_url
        self.loop = asyncio.get_running_loop()
        self.lock = threading.Lock()
        self.buffers = {}
        self.busy = False
        self.ignored_speakers = {}
        self.decoders = {}
        self.closed = False
        self.last_error = 0
        self.stats_lock = threading.Lock()
        self.stats = Counter(audio=0, rtcp=0, pcm=0, audible=0, phrases=0, stt=0)
        self.last_diagnostic = time.monotonic()
        # Uniquement le transfert vers le worker, aucune phrase en attente.
        self.queue = asyncio.Queue(maxsize=1)
        self.flush_task = self.loop.create_task(self.flush_phrases())
        self.worker_task = self.loop.create_task(self.transcribe_phrases())
        self.flush_task.add_done_callback(self.task_finished)
        self.worker_task.add_done_callback(self.task_finished)

    def task_finished(self, task):
        if not task.cancelled() and task.exception() is not None:
            logger.error("Tâche vocale arrêtée : %s erreur=%s",
                         task.get_coro().__name__, type(task.exception()).__name__)

    def wants_opus(self):
        # L'extension ne déchiffre pas encore DAVE : le faire avant Opus.
        return True

    def count(self, name):
        with self.stats_lock:
            self.stats[name] += 1

    def log_diagnostic(self, reason='10s'):
        with self.stats_lock:
            stats = dict(self.stats)
        voice = getattr(self, 'voice_client', None)
        connection = getattr(voice, '_connection', None)
        session = getattr(connection, 'dave_session', None)
        logger.info(
            "Diagnostic vocal (%s, compteurs cumulés) : %s file=%d dave_version=%s dave_ready=%s",
            reason, stats, self.queue.qsize(),
            getattr(connection, 'dave_protocol_version', None), getattr(session, 'ready', False),
        )

    @voice_recv.AudioSink.listener()
    def on_rtcp_packet(self, packet, guild):
        self.count('rtcp')

    @voice_recv.AudioSink.listener()
    def on_voice_member_speaking_start(self, member):
        self.count('speaking_start')

    def receiver_stopped(self, error):
        if error:
            logger.error("Thread de réception vocale arrêté : %s", type(error).__name__)
        else:
            logger.info("Thread de réception vocale arrêté normalement")

    def write(self, user, data):
        if self.closed:
            return
        self.count('audio')
        for ignored, reason in ((user is None, 'unknown_user'),
                                (getattr(user, 'bot', False), 'bot'),
                                (not data.opus, 'empty_opus'),
                                (data.packet.is_silence(), 'silence_packet')):
            if ignored:
                self.count(reason)
                return
        stage = 'DAVE'
        try:
            connection = self.voice_client._connection
            opus = bytes(data.opus)
            if connection.dave_protocol_version:
                if connection.dave_session is None or not connection.dave_session.ready:
                    self.count('dave_not_ready')
                    return
                opus = connection.dave_session.decrypt(user.id, MediaType.audio, opus)
            stage = 'Opus'
            with self.lock:
                if self.closed:
                    return
                decoder = self.decoders.get(user.id)
                if decoder is None:
                    decoder = self.decoders[user.id] = discord.opus.Decoder()
                pcm = decoder.decode(opus, fec=False)
                self.count('pcm')
                self.add_pcm(user, pcm, time.monotonic())
        except Exception as exc:
            self.count(stage + '_errors')
            # Ne jamais journaliser les paquets audio ou les clés de session.
            if time.monotonic() - self.last_error > 30:
                logger.warning("Réception vocale impossible à l'étape %s : %s", stage, type(exc).__name__)
                self.last_error = time.monotonic()

    def add_pcm(self, user, pcm, now):
        """Appelé sous le verrou ; PCM Discord = 48 kHz, stéréo, 16 bits."""
        audible = audioop.rms(pcm, 2) >= 200
        self.count('audible' if audible else 'below_threshold')
        last_ignored = self.ignored_speakers.get(user.id)
        if self.busy or (last_ignored is not None and now - last_ignored < 0.8):
            if audible:
                self.ignored_speakers[user.id] = now
            self.count('busy_ignored')
            return
        self.ignored_speakers.pop(user.id, None)
        if user.id not in self.buffers:
            if not audible:
                return
            self.buffers[user.id] = [user, bytearray(), now]
        buffer = self.buffers[user.id]
        # Plafond de 20 secondes même si la boucle asyncio est occupée.
        available = 48000 * 2 * 2 * 20 - len(buffer[1])
        buffer[1].extend(pcm[:available])
        if audible:
            buffer[2] = now

    async def flush_phrases(self):
        while True:
            await asyncio.sleep(0.1)
            now = time.monotonic()
            if now - self.last_diagnostic >= 10:
                self.log_diagnostic()
                self.last_diagnostic = now
            ready = []
            with self.lock:
                if self.busy:
                    continue
                for user_id, (user, pcm, last_voice) in list(self.buffers.items()):
                    if now - last_voice >= 0.8 or len(pcm) >= 48000 * 4 * 20:
                        del self.buffers[user_id]
                        if len(pcm) >= 48000 * 4 * 0.25:
                            ready.append((user, bytes(pcm)))
                            self.busy = True
                            # Oublier aussi les phrases déjà commencées par les autres.
                            self.ignored_speakers.update(
                                {uid: buffer[2] for uid, buffer in self.buffers.items()})
                            self.buffers.clear()
                            break
                        else:
                            self.count('too_short')
            for item in ready:
                try:
                    self.queue.put_nowait(item)
                    self.count('phrases')
                    logger.info("Phrase vocale sélectionnée : durée=%.2fs transfert=%d",
                                len(item[1]) / (48000 * 4), self.queue.qsize())
                except asyncio.QueueFull:
                    self.count('queue_full')
                    logger.warning("File vocale pleine : une phrase a été ignorée")

    async def transcribe_phrases(self):
        # Rester occupé jusqu'à la fin de la transcription ET de la réponse.
        async with httpx.AsyncClient(timeout=60) as http:
            while not self.closed:
                user, pcm = await self.queue.get()
                stage = 'STT'
                started = time.monotonic()
                try:
                    audio = BytesIO()
                    with wave.open(audio, 'wb') as wav:
                        wav.setnchannels(2)
                        wav.setsampwidth(2)
                        wav.setframerate(48000)
                        wav.writeframes(pcm)
                    self.count('stt')
                    logger.info("Envoi STT : durée audio=%.2fs", len(pcm) / (48000 * 4))
                    response = await http.post(
                        self.stt_url,
                        files={'file': ('phrase.wav', audio.getvalue(), 'audio/wav')},
                    )
                    response.raise_for_status()
                    text = response.json()['text'].strip()
                    logger.info("Retour STT : HTTP=%d durée=%.2fs caractères=%d",
                                response.status_code, time.monotonic() - started, len(text))
                    if not text:
                        self.count('empty_transcription')
                    if text and not self.closed:
                        stage = 'réponse'
                        await self.on_phrase(user, text)
                except asyncio.CancelledError:
                    raise
                except Exception as exc:
                    self.count(stage + '_errors')
                    status = exc.response.status_code if isinstance(exc, httpx.HTTPStatusError) else None
                    logger.warning("Traitement vocal impossible : étape=%s erreur=%s HTTP=%s durée=%.2fs",
                                   stage, type(exc).__name__, status, time.monotonic() - started)
                finally:
                    with self.lock:
                        self.busy = False
                    self.queue.task_done()

    @voice_recv.AudioSink.listener()
    def on_voice_member_disconnect(self, member, ssrc):
        if member:
            with self.lock:
                self.buffers.pop(member.id, None)
                self.decoders.pop(member.id, None)
                self.ignored_speakers.pop(member.id, None)

    def cleanup(self):
        # La bibliothèque appelle cleanup depuis son thread de réception.
        with self.lock:
            if self.closed:
                return
            self.closed = True
            self.buffers.clear()
            self.decoders.clear()
            self.ignored_speakers.clear()
        self.log_diagnostic('arrêt')
        if not self.loop.is_closed():
            self.loop.call_soon_threadsafe(self.flush_task.cancel)
            self.loop.call_soon_threadsafe(self.worker_task.cancel)

    async def close(self):
        self.cleanup()
        self.flush_task.cancel()
        self.worker_task.cancel()
        await asyncio.gather(self.flush_task, self.worker_task, return_exceptions=True)
        while not self.queue.empty():
            self.queue.get_nowait()
            self.queue.task_done()