Refactor main function into ChaosBot class for improved structure and initialization of components

botbotbot/__init__.py

@@ -1,142 +1,109 @@
 import logging
-import random
 import tomllib
+from typing import Any
 
 import discord
 
 from botbotbot.ai import AIBot
+from botbotbot.nicks import NickBot
 from botbotbot.shuffle import Shuffler
 from botbotbot.text import TextBot
 from botbotbot.tts import CambAI
 from botbotbot.voice import VoiceBot
 from botbotbot.wordlist import Wordlist
 
+logger = logging.getLogger(__name__)
+
+
+class ChaosBot:
+    def __init__(self, config: dict[str, Any]) -> None:
+        self.config = config
+
+        intents = discord.Intents.all()
+        intents.presences = False
+
+        self.bot = discord.Bot(description="Discord Chaos Bot", intents=intents)
+
+        self.init_aibot()
+        self.init_cambai()
+        self.init_guild_ids()
+
+        self.shuffler = Shuffler()
+        self.word_list = Wordlist(self.bot, self.guild_ids)
+        self.text_bot = TextBot(
+            self.bot,
+            self.word_list,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.voice_bot = VoiceBot(
+            self.bot,
+            self.cambai,
+            aibot=self.ai_bot,
+            shuffler=self.shuffler,
+            guild_ids=self.guild_ids,
+        )
+        self.nick_bot = NickBot(
+            self.bot, shuffler=self.shuffler, guild_ids=self.guild_ids
+        )
+
+        self.init_events()
+
+    def init_aibot(self) -> None:
+        self.ai_bot: AIBot | None = None
+
+        if isinstance(key := self.config.get("mistral_api_key"), str):
+            system_prompt = """Tu es une intelligence artificelle qui répond en français.
+            Tu dois faire un commentaire pertinent en lien avec ce qui te sera dit.
+            Ta réponse doit être très courte.
+            Ta réponse doit être une seule phrase.
+            TA RÉPONSE DOIT ÊTRE EN FRANÇAIS !!!"""
+
+            self.ai_bot = AIBot(
+                key,
+                model="mistral-large-latest",
+                system_message=system_prompt,
+            )
+        else:
+            logger.warning("No AI Bot.")
+
+    def init_cambai(self) -> None:
+        self.cambai: CambAI | None = None
+        if isinstance(key := self.config.get("cambai_api_key"), str):
+            self.cambai = CambAI(key)
+        else:
+            logger.warning("No CambAI.")
+
+    def init_guild_ids(self) -> None:
+        guild_ids = self.config.get("guild_ids")
+        if not (
+            isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)
+        ):
+            logger.error("Guild IDs must be a list of integers.")
+            guild_ids = []
+        self.guild_ids = guild_ids
+
+    def init_events(self) -> None:
+        self.bot.add_listener(self.on_ready, "on_ready")
+
+        self.word_list.init_events()
+        self.text_bot.init_events()
+        self.voice_bot.init_events()
+        self.nick_bot.init_events()
+
+    async def on_ready(self) -> None:
+        logger.info(f"We have logged in as {self.bot.user}")
+
+    def run(self) -> None:
+        self.bot.run(self.config.get("token"))
+
 
 def main() -> None:
-    description = """Discord Chaos Bot"""
-
-    logger = logging.getLogger(__name__)
     logging.basicConfig(level=logging.INFO)
 
     with open("config.toml", "rb") as config_file:
         config = tomllib.load(config_file)
 
-    guild_ids = config.get("guild_ids")
+    chaos_bot = ChaosBot(config)
-    if not (isinstance(guild_ids, list) and all(isinstance(i, int) for i in guild_ids)):
+    chaos_bot.run()
-        logger.error("Guild IDs must be a list of integers.")
-        guild_ids = []
-
-    system_prompt = """Tu es une intelligence artificelle qui répond en français.
-    Tu dois faire un commentaire pertinent en lien avec ce qui te sera dit.
-    Ta réponse doit être très courte.
-    Ta réponse doit être une seule phrase.
-    TA RÉPONSE DOIT ÊTRE EN FRANÇAIS !!!"""
-
-    aibot: AIBot | None = None
-
-    if isinstance(key := config.get("mistral_api_key"), str):
-        aibot = AIBot(
-            key,
-            model="mistral-large-latest",
-            system_message=system_prompt,
-        )
-    else:
-        logger.warning("No AI Bot.")
-
-    cambai: CambAI | None = None
-    if isinstance(key := config.get("cambai_api_key"), str):
-        cambai = CambAI(key)
-    else:
-        logger.warning("No CambAI.")
-
-    intents = discord.Intents.default()
-    intents.members = True
-    intents.message_content = True
-    intents.reactions = True
-    intents.voice_states = True
-
-    bot = discord.Bot(description=description, intents=intents)
-    shf = Shuffler()
-    wl = Wordlist(bot, guild_ids)
-    wl.init_events()
-
-    text_bot = TextBot(bot, wl, aibot=aibot, shuffler=shf)
-    text_bot.init_events()
-
-    voice_bot = VoiceBot(bot, cambai, aibot=aibot, shuffler=shf, guild_ids=guild_ids)
-    voice_bot.init_events()
-
-    @bot.listen("on_ready")
-    async def on_ready() -> None:
-        logger.info(f"We have logged in as {bot.user}")
-
-    @bot.slash_command(
-        name="alea", guild_ids=guild_ids, description="Modifier les pseudos"
-    )
-    async def alea(ctx: discord.ApplicationContext) -> None:
-        logger.info(f"ALEA {ctx.author}")
-        await ctx.defer()
-        if await shf.try_shuffle(ctx.guild):
-            embed = discord.Embed(title="ALEA", color=discord.Colour.green())
-            await ctx.respond(embed=embed, ephemeral=True, delete_after=30)
-            logger.info("FIN ALEA")
-        else:
-            embed = discord.Embed(title="ERRE ALEA", color=discord.Colour.red())
-            await ctx.respond(embed=embed)
-            logger.info("ERRE ALEA")
-
-    @bot.slash_command(
-        name="indu", guild_ids=guild_ids, description="Poser une question à MistralAI"
-    )
-    async def indu(ctx: discord.ApplicationContext, prompt: str) -> None:
-        if aibot is None:
-            return
-        logger.info(f"INDU {ctx.author} {prompt}")
-        await ctx.defer()
-        res_stream = await aibot.get_response_stream(prompt)
-
-        embed = discord.Embed(
-            title=prompt,
-            description="",
-            thumbnail="https://mistral.ai/images/favicon/favicon-32x32.png",
-            color=discord.Colour.orange(),
-        )
-        message = await ctx.respond(embed=embed)
-
-        async for chunk in res_stream:
-            if chunk.data.choices[0].delta.content is not None:
-                embed.description += chunk.data.choices[0].delta.content
-                await message.edit(embed=embed)
-
-        embed.colour = None
-        await message.edit(embed=embed)
-        logger.info("FIN INDU")
-
-    @bot.slash_command(
-        name="chan", guild_ids=guild_ids, description="Donner de nouveaux pseudos"
-    )
-    async def chan(ctx: discord.ApplicationContext, file: discord.Attachment) -> None:
-        logger.info(f"CHAN {ctx.author}")
-        await ctx.defer()
-
-        members = ctx.guild.members
-        members.remove(ctx.guild.owner)
-
-        nicks = (await file.read()).decode().splitlines()
-        if len(nicks) < len(members):
-            embed = discord.Embed(title="ERRE CHAN", color=discord.Colour.red())
-            await ctx.respond(embed=embed)
-            return
-
-        nicks = random.choices(nicks, k=len(members))
-        for member, nick in zip(members, nicks):
-            logger.info(member, nick)
-            await member.edit(nick=nick)
-
-        embed = discord.Embed(
-            title="CHAN", description="\n".join(nicks), color=discord.Colour.green()
-        )
-        await ctx.respond(embed=embed)
-        logger.info("FIN CHAN")
-
-    bot.run(config.get("token"))

botbotbot/nicks.py

@@ -0,0 +1,73 @@
+import logging
+import random
+
+import discord
+
+from botbotbot.shuffle import Shuffler
+
+logger = logging.getLogger(__name__)
+
+
+class NickBot:
+    def __init__(
+        self,
+        bot: discord.Bot,
+        shuffler: Shuffler,
+        guild_ids: list[int] = [],
+    ) -> None:
+        self.bot = bot
+        self.shf = shuffler
+        self.guild_ids = guild_ids
+
+    def init_events(self) -> None:
+        self.bot.add_application_command(
+            discord.SlashCommand(
+                self.alea,
+                name="alea",
+                description="Modifier les pseudos",
+                guild_ids=self.guild_ids,
+            )
+        )
+        self.bot.add_application_command(
+            discord.SlashCommand(
+                self.chan,
+                name="chan",
+                description="Donner de nouveaux pseudos",
+                guild_ids=self.guild_ids,
+            )
+        )
+
+    async def alea(self, ctx: discord.ApplicationContext) -> None:
+        logger.info(f"ALEA <{ctx.author}>.")
+        await ctx.defer()
+        if await self.shf.try_shuffle(ctx.guild):
+            await ctx.respond("ALEA", ephemeral=True, delete_after=30)
+            logger.info("FIN ALEA")
+        else:
+            await ctx.respond("ERRE ALEA", ephemeral=True, delete_after=30)
+            logger.info("ERRE ALEA")
+
+    async def chan(
+        self, ctx: discord.ApplicationContext, file: discord.Attachment
+    ) -> None:
+        logger.info(f"CHAN <{ctx.author}>.")
+        await ctx.defer()
+
+        members = ctx.guild.members
+        members.remove(ctx.guild.owner)
+
+        nicks = (await file.read()).decode().splitlines()
+        if len(nicks) < len(members):
+            await ctx.respond("ERRE CHAN", ephemeral=True, delete_after=30)
+            return
+
+        nicks = random.choices(nicks, k=len(members))
+        for member, nick in zip(members, nicks):
+            logger.info(member, nick)
+            await member.edit(nick=nick)
+
+        embed = discord.Embed(
+            title="CHAN", description="\n".join(nicks), color=discord.Colour.green()
+        )
+        await ctx.respond(embed=embed, ephemeral=True, delete_after=30)
+        logger.info("FIN CHAN")

botbotbot/text.py

@@ -21,6 +21,7 @@ class TextBot:
         wordlist: Wordlist,
         aibot: AIBot | None = None,
         shuffler: Shuffler | None = None,
+        guild_ids: list[int] = [],
         rnd_weights: list[float] = [10, 5, 10],
     ) -> None:
         self.bot = bot
@@ -28,12 +29,21 @@ class TextBot:
         self.wl = wordlist
         self._rnd_weights = rnd_weights
         self.shf = shuffler
+        self.guild_ids = guild_ids
 
     def init_events(self) -> None:
         self.bot.add_listener(self.on_message, "on_message")
         self.bot.add_listener(self.add_more_reaction, "on_reaction_add")
         self.bot.add_listener(self.react_message_edit, "on_message_edit")
         self.bot.add_listener(self.rando_shuffle, "on_message")
+        self.bot.add_application_command(
+            discord.SlashCommand(
+                self.indu,
+                name="indu",
+                guild_ids=self.guild_ids,
+                description="Poser une question à MistralAI",
+            )
+        )
 
     @property
     def rnd_weights(self) -> list[float]:
@@ -167,7 +177,7 @@ class TextBot:
             and after.author != self.bot.user
             and random.random() < 20 / 100
         ):
-            logger.info(f"React to edit from {after.author}")
+            logger.info(f"React to edit from {after.author}.")
             await after.add_reaction("👀")
 
     async def rando_shuffle(self, message: discord.Message) -> None:
@@ -179,3 +189,27 @@ class TextBot:
         ):
             logger.info(f"Message shuffle after message from {message.author}")
             await self.shf.try_shuffle(message.guild)
+
+    async def indu(self, ctx: discord.ApplicationContext, prompt: str) -> None:
+        if self.aibot is None:
+            return
+        logger.info(f"INDU {ctx.author} {prompt}")
+        await ctx.defer()
+        res_stream = await self.aibot.get_response_stream(prompt)
+
+        embed = discord.Embed(
+            title=prompt,
+            description="",
+            thumbnail="https://mistral.ai/images/favicon/favicon-32x32.png",
+            color=discord.Colour.orange(),
+        )
+        message = await ctx.respond(embed=embed)
+
+        async for chunk in res_stream:
+            if chunk.data.choices[0].delta.content is not None:
+                embed.description += chunk.data.choices[0].delta.content
+                await message.edit(embed=embed)
+
+        embed.colour = None
+        await message.edit(embed=embed)
+        logger.info("FIN INDU")