runabot/Christmas/Cogs/Event.py

import random
import wavelink

import onnxruntime as rt

from discord import ApplicationContext, DiscordException, Game, Guild
from discord.ext.commands import Cog

from discord.ext.commands import CommandOnCooldown
from Christmas.UI.Embed import Default_Embed
from Christmas.config import ChristmasConfig

model = rt.InferenceSession("Christmas/Tagging/model.onnx", provider_options="CPUExecutionProvider")

class Event(Cog):
    def __init__(self, bot):
        self.bot = bot
        self.config = ChristmasConfig()
    
    @Cog.listener()
    async def on_application_command_error(self, ctx: ApplicationContext, exception: DiscordException) -> None:
        if isinstance(exception, CommandOnCooldown):
            await ctx.respond(embed=Default_Embed.cooldown(exception.retry_after), ephemeral=True)
        
    @Cog.listener()
    async def on_ready(self) -> None:
        print("Ready!")
        await self.bot.change_presence(activity=Game(name="크리스마스에 함께!"))
        global model
        model = rt.InferenceSession("Christmas/Tagging/model.onnx", provider_options="CPUExecutionProvider")
        print("Model Loaded!")
    
    # connect wavelink
    @Cog.listener()
    async def on_connect(self) -> None:
        await self.bot.wait_until_ready()
        nodes = []
        for node in self.config.LAVALINK:
            nodes.append(wavelink.Node(uri=node["HOST"], password=node["PASSWORD"], identifier=node["IDENTIFIER"]))
        await wavelink.Pool.connect(nodes=nodes, client=self.bot)
    
    @Cog.listener()
    async def on_guild_join(self, guild: Guild) -> None:
        if guild.system_channel is not None:
            await guild.system_channel.send(embed=Default_Embed.guild_join())
        else:
            await random.choice(guild.text_channels).send(embed=Default_Embed.guild_join())
    
    #@Cog.listener()
    #async def on_wavelink_node_ready(self, node: wavelink.Node):
    #    print(f"Node {node.id} is ready!")
        
def setup(bot):
    bot.add_cog(Event(bot))