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, on_transcription=None, on_batch_status=None):
        super().__init__()
        self.on_phrase = on_phrase
        self.on_transcription = on_transcription
        self.on_batch_status = on_batch_status
        self.pending_receipts = []
        self.pending = []
        self.pending_event = asyncio.Event()
        self.stt_url = stt_url
        self.loop = asyncio.get_running_loop()
        self.lock = threading.Lock()
        self.buffers = {}
        self.decoders = {}
        self.closed = False
        self.responding = False
        self.last_stt_seconds = None
        self.last_reply_seconds = None
        self.last_voice = 0.0
        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()
        # File audio bornée si Whisper ne suit plus le rythme du vocal.
        self.queue = asyncio.Queue(maxsize=32)
        self.flush_task = self.loop.create_task(self.flush_phrases())
        self.worker_task = self.loop.create_task(self.transcribe_phrases())
        self.reply_task = self.loop.create_task(self.respond_batches())
        self.reply_task.add_done_callback(self.task_finished)
        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')
        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
            self.last_voice = 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:
                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)))
                        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):
        # La transcription continue pendant les réponses du LLM.
        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()
                    self.last_stt_seconds = time.monotonic() - started
                    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:
                        receipt = None
                        try:
                            stage = 'affichage'
                            if self.on_transcription:
                                receipt = await self.on_transcription(user, text)
                        finally:
                            if not self.closed:
                                self.pending.append((user, text))
                                self.pending_receipts.append(receipt)
                                self.pending_event.set()
                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:
                    self.queue.task_done()

    async def respond_batches(self):
        while not self.closed:
            await self.pending_event.wait()
            # L'affichage STT continue ; seule la réponse attend le silence global.
            while not self.closed:
                await self.queue.join()
                with self.lock:
                    quiet = time.monotonic() - self.last_voice > 2.0
                    buffered = bool(self.buffers)
                if quiet and not buffered and self.queue.empty():
                    break
                await asyncio.sleep(0.1)
            if self.closed:
                return
            self.pending_event.clear()
            batch, self.pending = self.pending, []
            receipts, self.pending_receipts = self.pending_receipts, []
            if not batch:
                continue
            started = time.monotonic()
            self.responding = True
            try:
                await self.report_status(receipts, 'processing')
                await self.on_phrase(batch)
                await self.report_status(receipts, 'done')
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                await self.report_status(receipts, 'error')
                logger.warning("Réponse au lot vocal impossible : %s", type(exc).__name__)
            finally:
                self.last_reply_seconds = time.monotonic() - started
                self.responding = False

    async def report_status(self, receipts, state):
        if self.on_batch_status:
            try:
                await self.on_batch_status(receipts, state)
            except Exception as exc:
                logger.warning('Mise à jour du statut impossible : %s', type(exc).__name__)

    @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)

    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.pending.clear()
            self.pending_receipts.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)
            self.loop.call_soon_threadsafe(self.reply_task.cancel)

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