import logging import wave import asyncio import assemblyai as aai import discord from discord.ext import commands, voice_recv, tasks from discord.opus import Decoder as OpusDecoder import time # Replace with your API key aai.settings.api_key = "aa9962f0088a449a9c4ab2361e96cc08" 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 logger = logging.getLogger("discord") location = "/home/pi/Conjurer/transcripts/" class CommunicationObject: def __init__(self, msg_type, user, data): 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): self.queue = queue self.user = user_id self.username = str(user_id.id) #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 = 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 = 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 = 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): 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): super().__init__() self.queue = queue self.wavewriter = {} def on_user_connect(self, username): self.wavewriter[str(username.id)] = [WaveWriter(username, self.queue), 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: logger.info("Cleanup for SRBuffer") for item in self.wavewriter.values(): item[0].cleanup() except Exception: 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) self.worker = None self.mc = None self.vc = None @commands.hybrid_command(name="transcribe") 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): logger.debug("Already transcribing") else: logger.debug("Already connected with other client") else: logger.debug("Connected") vc = await ctx.author.voice.channel.connect(cls=voice_recv.VoiceRecvClient) 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 > 3000149433 and timediff_write > 500014943 and item[4]: logger.info("File rotation time since last write %s %s", timediff_write, timediff_write/1e9) 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: logger.debug("Ignoring self") return if before.channel is None: logger.info("User %s connected to channel %s", user, after.channel.name) self.wsink.on_user_connect(user) elif after.channel is None: 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: logger.debug("User VC status changed %s", user.id) logger.debug("Before %s", before) logger.debug("After %s", after) @commands.command(name="stop_transcribe") 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): logger.info("Transcript uploader start") config = aai.TranscriptionConfig(language_code="pl") transcriber = aai.Transcriber() logger.debug("Transcriber queue id: %s",id(self.comm_queue)) while True: logger.debug("waiting for tasks") item = await self.comm_queue.get() logger.debug("Got %s", item) if "STOP" in item.type: logger.debug("Queue ended") break elif "send_file" in item.type: 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: logger.warn("Exceptiom occured %s", e) await self.mc.send(f"{item.user} : {transcript.text}") if transcript.error: logger.error(transcript.error) raise AssertionError elif "user_cleanup" in item.type: logger.info("User %s disconnected - cleanup action") else: logger.warn("Something went wrong with object %s", item) logger.info("Transcript uploader stoped") async def setup(bot): logger = logging.getLogger("discord") await bot.add_cog(Transcriber(bot)) logger.info("Loading voice transcribed module done")