conjurer/voice_recognition_commands.py

import asyncio
import logging
import time
import wave

import assemblyai as aai
import discord
from discord.ext import commands, tasks, voice_recv
from discord.opus import Decoder as OpusDecoder

from constants import ASSEMBLYAI_API_KEY, TRANSCRIPTS_PATH

# Credentials come from constants (env ASSEMBLYAI_API_KEY, or the 'assemblyai'
# machine in the netrc at CONJURER_NETRC_FILE). Raising here means the guarded
# extension loader logs the reason and disables ONLY this cog.
if not ASSEMBLYAI_API_KEY:
    raise RuntimeError(
        "AssemblyAI API key not configured (netrc machine 'assemblyai' or "
        "ASSEMBLYAI_API_KEY env) - voice recognition stays disabled"
    )
aai.settings.api_key = ASSEMBLYAI_API_KEY

discord.opus._load_default()
CHANNELS = OpusDecoder.CHANNELS
SAMPLE_WIDTH = OpusDecoder.SAMPLE_SIZE // OpusDecoder.CHANNELS
SAMPLING_RATE = OpusDecoder.SAMPLING_RATE

# rotate file after there is 0.5s between last received pcm for user.
# delete messsages after user disconnect

LOCATION = TRANSCRIPTS_PATH


class CommunicationObject:
    def __init__(self, msg_type, user, data):
        self.logger = logging.getLogger("discord")
        self.type = msg_type
        self.user = user
        self.data = data

    def __repr__(self):
        return f"{self.type} : {self.user} : {self.data}"


class WaveWriter:
    def __init__(self, user_id, queue, location):
        self.queue = queue
        self.user = user_id
        self.username = str(user_id.id)
        self.location = location
        # here we can add hashing function to make transcription files not possible to be connected with discord id
        self.present_file_id = 0
        self.file_name_past = None
        self.file_name_present = (
            self.location + self.username + "_" + str(self.present_file_id) + ".mp3"
        )
        self.transcript_file_present: wave.Wave_write = wave.open(
            self.file_name_present, "wb"
        )
        self.transcript_file_present.setnchannels(CHANNELS)
        self.transcript_file_present.setsampwidth(SAMPLE_WIDTH)
        self.transcript_file_present.setframerate(SAMPLING_RATE)
        self.file_name_future = (
            self.location + self.username + "_" + str(self.present_file_id + 1) + ".mp3"
        )
        self.transcript_file_future: wave.Wave_write = wave.open(
            self.file_name_future, "wb"
        )
        self.transcript_file_future.setnchannels(CHANNELS)
        self.transcript_file_future.setsampwidth(SAMPLE_WIDTH)
        self.transcript_file_future.setframerate(SAMPLING_RATE)

    async def rotate(self):

        self.transcript_file_present.close()
        self.file_name_past = self.file_name_present
        if self.present_file_id > 10:
            self.present_file_id = 0
        else:
            self.present_file_id += 1
        self.file_name_present = self.file_name_future
        self.transcript_file_present = self.transcript_file_future
        self.file_name_future = (
            self.location + self.username + "_" + str(self.present_file_id + 1) + ".mp3"
        )
        self.transcript_file_future: wave.Wave_write = wave.open(
            self.file_name_future, "wb"
        )
        self.transcript_file_future.setnchannels(CHANNELS)
        self.transcript_file_future.setsampwidth(SAMPLE_WIDTH)
        self.transcript_file_future.setframerate(SAMPLING_RATE)
        operation = CommunicationObject(
            msg_type="send_file", user=self.user, data=[self.file_name_past]
        )

        await self.queue.put(operation)

    def writeframes(self, pcmdata):
        self.transcript_file_present.writeframes(pcmdata)

    def cleanup(self):
        self.logger.info("Cleanup for user %s", self.username)
        self.transcript_file_present.close()
        self.transcript_file_future.close()


class SRBuffer(voice_recv.AudioSink):
    """Endpoint AudioSink that generates a wav file.
    Best used in conjunction with a silence generating sink. (TBD)
    """

    # on member join dodajemytypa do listy
    # on member disconnect - dropujemy go
    def __init__(self, queue, location):
        super().__init__()
        self.queue = queue
        self.wavewriter = {}
        self.logger = logging.getLogger("discord")
        self.location = location

    def on_user_connect(self, username):
        self.wavewriter[str(username.id)] = [
            WaveWriter(username, self.queue, self.location),
            time.time_ns(),
            time.time_ns(),
            False,
            False,
        ]

    def on_user_disconnect(self, username):
        self.wavewriter[str(username.id)].cleanup()
        self.wavewriter.pop(str(username.id))

    def wants_opus(self) -> bool:
        return False

    def write(self, user, data) -> None:
        # logger.info("DAta write")
        if user:
            self.wavewriter[str(user.id)][0].writeframes(data.pcm)
            self.wavewriter[str(user.id)][1] = time.time_ns()  # time from last write
            self.wavewriter[str(user.id)][3] = True  # data written in last iter
            self.wavewriter[str(user.id)][4] = True  # data in buff
            time.sleep(0.001)

    def cleanup(self) -> None:
        try:
            self.logger.info("Cleanup for SRBuffer")
            for item in self.wavewriter.values():
                item[0].cleanup()
        except Exception:
            self.logger.warning(
                "WaveSink got error closing file on cleanup", exc_info=True
            )


class Transcriber(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.threads = []
        self.comm_queue = asyncio.Queue()
        self.wsink = SRBuffer(self.comm_queue, LOCATION)
        self.worker = None
        self.mc = None
        self.vc = None
        self.logger = logging.getLogger("discord")

    @commands.hybrid_command(name="transcribe")
    @commands.has_any_role('Nocna Zmiana', 'Jarl', 'Thane' , 'Bartender')
    async def test(self, ctx):
        if self.vc:
            vc = self.vc  # to juz powinien byc voice channel z funkcja conenct
        else:
            vc = None
        self.mc = ctx.message.channel
        if self.bot.voice_clients:
            if isinstance(self.bot.voice_clients[0], voice_recv.VoiceRecvClient):
                self.logger.debug("Already transcribing")
            else:
                self.logger.debug("Already connected with other client")
        else:
            self.logger.debug("Connected")
            vc = await ctx.author.voice.channel.connect(cls=voice_recv.VoiceRecvClient)
            self.logger.info(self.bot.voice_clients)
            self.check_data.start()
            self.worker = asyncio.create_task(self.transcribe_output_queue())
            vc.listen(self.wsink)

    @tasks.loop(seconds=0.5)
    async def check_data(self):
        for item in self.wsink.wavewriter.values():
            if not item[3]:
                timediff_rotation = time.time_ns() - item[2]
                timediff_write = time.time_ns() - item[1]
                if (
                    timediff_rotation > 13000149433
                    and timediff_write > 500014943
                    and item[4]
                ):
                    self.logger.info(
                        "File rotation time since last write %s %s",
                        timediff_write,
                        timediff_write / 1e9,
                    )
                    self.logger.info(
                        "File rotation time since last rotation %s %s",
                        timediff_rotation,
                        timediff_rotation / 1e9,
                    )
                    item[2] = time.time_ns()
                    item[4] = False
                    await item[0].rotate()
            item[3] = False

    @commands.Cog.listener()
    async def on_voice_state_update(self, user, before, after):
        if user == self.bot.user:
            self.logger.debug("Ignoring self")
            return

        if before.channel is None:
            self.logger.info(
                "User %s connected to channel %s", user, after.channel.name
            )
            self.wsink.on_user_connect(user)
        elif after.channel is None:
            self.logger.info(
                "User %s disconnected from channel %s", user, before.channel.name
            )
            operation = CommunicationObject(
                msg_type="user_cleanup", user=user, data=None
            )
            await self.comm_queue.put(operation)
        else:
            self.logger.debug("User VC status changed %s", user.id)
            self.logger.debug("Before %s", before)
            self.logger.debug("After %s", after)

    @commands.command(name="stop_transcribe")
    @commands.has_any_role('Nocna Zmiana', 'Jarl', 'Thane' , 'Bartender')
    async def stop(self, ctx):
        self.check_data.stop()
        stop_token = CommunicationObject("STOP", None, None)

        await self.comm_queue.put(stop_token)
        await ctx.voice_client.disconnect()

    async def transcribe_output_queue(self):
        self.logger.info("Transcript uploader start")
        config = aai.TranscriptionConfig(language_code="pl")
        transcriber = aai.Transcriber()
        self.logger.debug("Transcriber queue id: %s", id(self.comm_queue))

        while True:
            self.logger.debug("waiting for tasks")
            item = await self.comm_queue.get()
            self.logger.debug("Got %s", item)
            if "STOP" in item.type:
                self.logger.debug("Queue ended")
                break
            elif "send_file" in item.type:
                self.logger.debug("Sending file for transcription")
                transcript = None
                for iter in item.data:
                    try:
                        coro = asyncio.to_thread(
                            transcriber.transcribe, iter, config=config
                        )
                        transcript = await coro
                    except Exception as e:
                        self.logger.warn("Exceptiom occured %s", e)
                    await self.mc.send(f"{item.user} : {transcript.text}")
                    if transcript.error:
                        self.logger.error(transcript.error)
                        raise AssertionError
            elif "user_cleanup" in item.type:
                self.logger.info("User %s disconnected - cleanup action")
            else:
                self.logger.warn("Something went wrong with object %s", item)
        self.logger.info("Transcript uploader stoped")


async def setup(bot):
    logger = logging.getLogger("discord")
    await bot.add_cog(Transcriber(bot))
    logger.info("Loading voice transcriber module done")