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

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.decoders = {}
        self.closed = False
        self.last_error = 0
        self.queue = asyncio.Queue(maxsize=8)
        self.flush_task = self.loop.create_task(self.flush_phrases())
        self.worker_task = self.loop.create_task(self.transcribe_phrases())

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

    def write(self, user, data):
        if self.closed or user is None or user.bot or not data.opus or data.packet.is_silence():
            return
        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:
                    return
                opus = connection.dave_session.decrypt(user.id, MediaType.audio, 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.add_pcm(user, pcm, time.monotonic())
        except Exception as exc:
            # 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 : %s", 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
        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()
            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)))
            for item in ready:
                try:
                    self.queue.put_nowait(item)
                except asyncio.QueueFull:
                    logger.warning("File vocale pleine : une phrase a été ignorée")

    async def transcribe_phrases(self):
        # Une seule réponse à la fois pour conserver l'ordre des phrases.
        async with httpx.AsyncClient(timeout=60) as http:
            while not self.closed:
                user, pcm = await self.queue.get()
                try:
                    audio = BytesIO()
                    with wave.open(audio, 'wb') as wav:
                        wav.setnchannels(2)
                        wav.setsampwidth(2)
                        wav.setframerate(48000)
                        wav.writeframes(pcm)
                    response = await http.post(
                        self.stt_url,
                        files={'file': ('phrase.wav', audio.getvalue(), 'audio/wav')},
                    )
                    response.raise_for_status()
                    text = response.json()['text'].strip()
                    if text and not self.closed:
                        await self.on_phrase(user, text)
                except asyncio.CancelledError:
                    raise
                except Exception as exc:
                    logger.warning("Traitement vocal impossible : %s", type(exc).__name__)
                finally:
                    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)

    def cleanup(self):
        # La bibliothèque appelle cleanup depuis son thread de réception.
        with self.lock:
            self.closed = True
            self.buffers.clear()
            self.decoders.clear()
        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()