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):
        super().__init__()
        self.on_phrase = on_phrase
        self.on_transcription = on_transcription
        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_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

    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:
                        try:
                            stage = 'affichage'
                            if self.on_transcription:
                                await self.on_transcription(user, text)
                        finally:
                            if not self.closed:
                                self.pending.append((user, text))
                                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()
            self.pending_event.clear()
            batch, self.pending = self.pending, []
            if not batch:
                continue
            started = time.monotonic()
            self.responding = True
            try:
                await self.on_phrase(batch)
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                logger.warning("Réponse au lot vocal impossible : %s", type(exc).__name__)
            finally:
                self.last_reply_seconds = time.monotonic() - started
                self.responding = False

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