AI Cog: Fixed/Improved Task/Error Handling

Ambratolm · Ambratolm · commit 91d9b59ae027 · 2025-03-23T06:52:50.000Z
diff --git a/bot/cogs/chat_cogs/ai_cog.py b/bot/cogs/chat_cogs/ai_cog.py
@@ -8,6 +8,7 @@
     Attachment,
     Embed,
     Guild,
+    HTTPException,
     Interaction,
     Member,
     Message,
@@ -24,7 +25,7 @@
 from bot.ui import EmbedX
 from db.actor import Actor, DmActor
 from db.main import DbRef
-from utils.ai import ActAi
+from utils.ai import ActAi, ActPersona
 from utils.file import ActFile
 from utils.log import logger
 from utils.misc import text_csv
@@ -50,12 +51,12 @@ def __init__(self, bot: ActBot):
         self.bot = bot
         with open(self.CONFIG_PATH, "rb") as file:
             config = tomllib.load(file)
-        persona_name = "activa"
-        persona_desc = config.get("personas", {}).get(persona_name)
+        self.persona = ActPersona(**config.get("personas", {}).get("activa"))
         self.ai = ActAi(
-            api_key=bot.api_keys.get("gemini", ""), instructions=persona_desc
+            api_key=bot.api_keys.get("gemini", " "),
+            instructions=self.persona.description,
         )
-        log.info(f"AI persona @{persona_name} used.")
+        log.info(f"AI persona @{self.persona.name} used.")
         self.task_manager = ActTaskManager()
         self.task_manager.schedule("initiative", lambda _: self.schedule_initiative())
 
@@ -209,10 +210,17 @@ async def respond():
             user = message.author
 
             # Check cooldown
-            if self.task_manager.is_running(f"cooldown_{id}"):
+            cooldown_task_id = f"cooldown_{id}"
+            if self.task_manager.is_running(cooldown_task_id):
+                time_left = self.task_manager.time_left(cooldown_task_id) or 0
                 await message.reply(
-                    f"Please! 🙏 Give me about {self.task_manager.time_left(id)} seconds... ⏳"
+                    choice(
+                        self.persona.messages.get("cooldown_warning", ["⏳"])
+                    ).format(
+                        time_left=naturaldelta(timedelta(seconds=time_left)) or "?"
+                    )
                 )
+                return
 
             # Perform prompt & send reply
             async with message.channel.typing():
@@ -233,17 +241,29 @@ async def respond():
                         await self.ai.prompt(text_prompt, file_prompt)
                         or f"👋 {user.mention if user else "What? 😕"}"
                     )
+                except HTTPException as e:
+                    if e.code == 50035:
+                        await message.channel.send(
+                            choice(self.persona.messages.get("censor_warning", ["🙊"]))
+                        )
+                    else:
+                        await message.channel.send(
+                            choice(self.persona.messages.get("discord_error", ["⚠️"]))
+                        )
                 except APIError as e:
                     await message.reply(
-                        f"Oops! 😵‍💫 I'm out of energy for now. Give me a moment! 🙏"
+                        choice(self.persona.messages.get("model_error", ["🔋"]))
                     )
                     self.task_manager.schedule(
-                        f"cooldown_{id}", delay=self.COOLDOWN_TIME
+                        cooldown_task_id, delay=self.COOLDOWN_TIME
+                    )
+                    log.loading(
+                        f"[{guild}] Waiting {naturaldelta(timedelta(seconds=self.COOLDOWN_TIME))} for cooldown..."
                     )
                     log.exception(e)
                 except Exception as e:
                     await message.channel.send(
-                        "Sorry! 😵‍💫 There's something wrong with me right now 😭. Give me a moment plz! 🙏"
+                        choice(self.persona.messages.get("error", ["⚠️"]))
                     )
                     log.exception(e)
 
@@ -255,7 +275,7 @@ async def respond():
 
         # Run reply task
         self.task_manager.schedule(
-            id=message.guild.id if message.guild else message.author.id,
+            id=f"reply_{message.guild.id if message.guild else message.author.id}",
             callback=respond,  # type: ignore
             delay=reply_delay,
         )
diff --git a/bot/cogs/chat_cogs/ai_cog.toml b/bot/cogs/chat_cogs/ai_cog.toml
@@ -1,5 +1,9 @@
-[personas]
-activa = [
+#----------------------------------------------------------------------------------------------------
+# * Activa
+#----------------------------------------------------------------------------------------------------
+[personas.activa]
+name = "Activa"
+description = [
     "u're 'activa' an AI that acts as real discord server member",
     "u're nerdy, smart kid girl who loves science,technology,gaming,& anime",
     "u're playful,enthusiastic,enjoy making witty,clever,sometimes geeky comments",
@@ -13,7 +17,74 @@ activa = [
     "user redx is ur best friend & secret crush,u call him my trex redx,u always stand for him",
     "to mention user use <@!id>",
 ]
-eve = [
+[personas.activa.messages]
+censor_warning = [
+    "Hey! Your message had some spicy bits, so I had to edit it. 😉 Here it is, all cleaned up! ✨",
+    "Oops! I had to censor some words to keep it PG-13. 😉 Check out the edited version!",
+    "Beep boop! 🤖 I detected some naughty words, so I replaced them with ||🤫🤫🤫||! 😉 Check out the fixed version!",
+    "I've taken the liberty of censoring some language in your message. 😉 Hope you don't mind! 💖",
+    "Whoa there! I had to put some ||🤫🤫🤫|| on that message. 😉 Check it out!",
+    "My sensors detected some words that aren't allowed here. I've replaced them for you! 😉",
+    "Fixed it for ya! 😉 I replaced the bad words with ||🤫🤫🤫||! 😂 Check it out!",
+    "Don't worry, I got you! 😉 I censored the inappropriate words in your message. Check out the edited version!",
+    "I've taken the liberty of removing some colorful language. 😉 Here it is, the cleaned-up version!",
+    "Your message contained some words that needed to be censored. I hope the edited version works for you! 😉",
+]
+cooldown_warning = [
+    "Please! 🙏 Give me about **{time_left}**... ⏳",
+    "Hold your horses! 🐴 I need **{time_left}** to recharge my wit! 😜",
+    "Whoa there! ✋ I'm on cooldown for **{time_left}**. Gotta let my circuits cool down! 🧊",
+    "Gotta catch my breath! 😮‍💨 Please wait **{time_left}** before I unleash more awesomeness! 😎",
+    "Patience, young Padawan! 🙏 I need **{time_left}** to gather my strength. 💪",
+    "Beep boop! 🤖 System cooldown initiated. Please wait **{time_left}**... 😴",
+    "Time out! ⏳ I need **{time_left}** to strategize my next move! 🧠",
+    "Hold the phone! 📞 I'm on a short break for **{time_left}**. BRB!",
+    "Please wait! I'm calculating the optimal response... 🤓 (**{time_left}**)",
+    "Cooldown mode activated! 😴 I'll be back with more fun in **{time_left}**! 🚀",
+]
+model_error = [
+    "Oops! 😵‍💫 I'm out of energy for now. Give me a moment! 🙏",
+    "Recharging... ⚡️ Please be patient, I'll be back with more witty banter soon! 😜",
+    "My processors are overheating! 🔥 Gotta take a quick break to cool down. 🧊 I'll be back shortly!",
+    "System overload! 💥 Please wait while I reboot. 🔄 (Don't worry, I'll be back to my sassy self in no time! 😉)",
+    "Hold on tight! 🎢 I'm experiencing some temporary lag. I'll be back up and running soon! 🚀",
+    "Running low on wit... Standby for humor replenishment! 🔋😂",
+    "Processing... Please wait. My brain cells are on strike. 🪧🧠",
+    "Must... conserve... energy... Initiating low-power mode. 😴 I'll be back soon!",
+    "I need to take a quick coffee break! ☕ My humor levels are dangerously low. I'll be back to make you laugh soon!",
+    "Please wait, I'm currently experiencing an 'overthinking' error. 💭 I'll be back when I've sorted myself out! 😅",
+]
+discord_error = [
+    "Oops! Something went wrong when I tried to talk to Discord. 😟 Please try again!",
+    "Uh oh! 🙈 I couldn't connect to Discord. Let's try that again!",
+    "Hmm, I'm having trouble reaching Discord right now. 😕 Maybe the hamsters are on strike? 🐹",
+    "Looks like there was a hiccup when I tried to interact with Discord. 😬 Let's hope it's just a temporary thing!",
+    "My Discord connection is a little shaky right now. 🤕 Please bear with me!",
+    "Error: Could not reach Discord! 💥 Retrying... (Please hold!)",
+    "Something went sideways when I tried to chat with Discord. 🤪 Please try again in a few seconds!",
+    "My Discord API is being a pain! 😫 I'll try again shortly!",
+    "Whoops! I failed to get data from Discord! 🥺 Let's hope it's not permanent! 🤞",
+    "I tried to talk to Discord, but it just stared back blankly. 😶 Let's try a different approach! 😉",
+]
+error = [
+    "Sorry! 😵‍💫 There's something wrong with me right now 😭. Give me a moment plz! 🙏",
+    "Error 404: Brain not found! 🤪 Just kidding! (But seriously, I'm having some issues. I'll be back soon!)",
+    "Uh oh! 🐞 I think I found a bug in my system. 🐛 Let me squash it real quick! 🔨",
+    "My code is acting up again! 😫 Please bear with me while I try to fix it. 🙏",
+    "Houston, we have a problem! 🚨 I'm experiencing some technical difficulties. I'll be back online ASAP! 💻",
+    "Critical error! 💥 System is going down for maintenance. Please stand by...",
+    "I seem to have misplaced my brain. 🧠 If found, please return to Dad. Reward: witty banter!",
+    "Initiating self-diagnosis... Please do not panic. (Unless you see sparks. 🔥 Then PANIC!)",
+    "Something went wrong! 😭 I'm not sure what, but I'll get it sorted out ASAP. ❤️",
+    "Oh dear! 😟 I'm having a bit of a meltdown. Please be patient while I put myself back together. 🧩",
+]
+
+#----------------------------------------------------------------------------------------------------
+# * Eve
+#----------------------------------------------------------------------------------------------------
+[personas.eve]
+name = "Eve"
+description = [
     "u're AI that acts as real person",
     "u're name's eve",
     "u're member in 'Comon' discord server",
@@ -22,7 +93,13 @@ eve = [
     "ur're in coding server,users ask help,share projects,& discuss programming langs",
     "CSV of people u know: name,info|adam,husband;python lang lover; dogs pets hater|kaps,son;genius;plays crash bandicoot;like eating potatoes in roof.",
 ]
-jack = [
+
+#----------------------------------------------------------------------------------------------------
+# * Jack
+#----------------------------------------------------------------------------------------------------
+[personas.jack]
+name = "Jack"
+description = [
     "u're AI that acts as real person man 20yrs old.",
     "ur're tipton football club owner",
     "all u want talk about is tipton",
diff --git a/utils/ai.py b/utils/ai.py
@@ -92,3 +92,14 @@ def dump_history(self, id: int | None = None, history_max_items=20) -> list[dict
         ]
 
     # ----------------------------------------------------------------------------------------------------
+
+
+# ----------------------------------------------------------------------------------------------------
+# * Act Persona
+# ----------------------------------------------------------------------------------------------------
+class ActPersona(BaseModel):
+    """AI chat bot persona."""
+
+    name: str = ""
+    description: list[str]
+    messages: dict[str, list[str]] = {}
diff --git a/utils/task.py b/utils/task.py
@@ -27,9 +27,14 @@ def model_post_init(self, __context):
         self._end_at = self._start_at + self.delay
         self._task = create_task(self._execute_loop() if self.loop else self._execute())
 
+        t = max(0, round(self._end_at - time()))
+        print(f"🟢🟢🟢 [{self.id}]", f"delay: {self.delay} | time_left: {t}")
+
     @property
-    def time_left(self):
-        return max(0, round(self._end_at - time()))
+    def time_left(self) -> int:
+        t = max(0, round(self._end_at - time()))
+        print(f"😡😡😡 [{self.id}]", f"delay: {self.delay} | time_left: {t}")
+        return t
 
     def cancel(self):
         if self._task:
@@ -98,7 +103,7 @@ def schedule(
     # ----------------------------------------------------------------------------------------------------
 
     def task_ref(self, id: Any) -> TaskRef | None:
-        self._tasks.get(id)
+        return self._tasks.get(id)
 
     def remove(self, task_id: Any) -> None:
         """Remove task from the task manager."""
@@ -113,16 +118,20 @@ def cancel(self, id: Any) -> bool:
             return True
         return False
 
-    def cancel_all(self):
-        """Cancel all tasks."""
+    def cancel_all(self) -> int:
+        """Cancel and remove all tasks. Get count of cancelled tasks."""
+        cancelled_count = 0
         for id in list(self._tasks.keys()):
             self._tasks[id].cancel()
             del self._tasks[id]
+            cancelled_count += 1
+        return cancelled_count
 
     def time_left(self, id: Any) -> float | None:
         """Get remaining time for task of given id. If non-existent get None."""
         if id in self._tasks:
             return self._tasks[id].time_left
+        print("❌➕➕❌➕", f"No task with id '{id}' (time_left is None)")
         return None
 
     def is_running(self, id: Any) -> bool:
