Implement donation levels

Author: wookie184

bot/exts/smart_eval/_cog.py

@@ -1,3 +1,4 @@
+import asyncio
 import random
 import re
 
@@ -8,6 +9,15 @@
 from bot.bot import SirRobin
 from bot.exts.smart_eval._smart_eval_rules import DEFAULT_RESPONSES, RULES
 
+DONATION_LEVELS = {
+    # Number of donations: (response time, intelligence level)
+    0: (15, 0),
+    10: (10, 1),
+    20: (8, 2),
+    30: (6, 3),
+    40: (5, 4),
+    50: (4, 5),
+}
 
 class SmartEval(commands.Cog):
     """Cog that handles all Smart Eval functionality."""
@@ -20,9 +30,51 @@ def __init__(self, bot: SirRobin):
 
     async def cog_load(self) -> None:
         """Run startup tasks needed when cog is first loaded."""
-        self.total_donations = await self.smarte_donation_cache.length()
+
+    async def get_gpu_capabilities(self) -> tuple[int, int]:
+        """Get the GPU capabilites based on the number of donated GPUs."""
+        total_donations = await self.total_donations()
+        response_time, intelligence_level = DONATION_LEVELS[0]
+        for donation_level, (time, max_response) in DONATION_LEVELS.items():
+            if total_donations >= donation_level:
+                response_time = time
+                intelligence_level = max_response
+            else:
+                break
+
+        return response_time, intelligence_level
+
+    @commands.command()
+    async def donations(self, ctx: commands.Context) -> None:
+        """Display the number of donations recieved so far."""
+        total_donations = await self.total_donations()
+        response_time, intelligence_level = await self.get_gpu_capabilities()
+        msg = (
+            f"Currently, I have received {total_donations} GPU donations, "
+            f"and am at intelligence level {intelligence_level}! "
+        )
+
+        # Calculate donations needed to reach next intelligence level
+        donations_needed = 0
+        for donation_level in DONATION_LEVELS:
+            if donation_level > total_donations:
+                donations_needed = donation_level - total_donations
+                break
+
+        if donations_needed:
+            msg += (
+                f"\n\nTo reach the next intelligence level, I need {donations_needed} more donations! "
+                f"Please consider donating your GPU to help me out. "
+            )
+
+        await ctx.reply(msg)
+
+    async def total_donations(self) -> int:
+        """Get the total number of donations."""
+        return await self.smarte_donation_cache.length()
 
     @commands.command(aliases=[])
+    @commands.max_concurrency(1, commands.BucketType.user)
     async def donate(self, ctx: commands.Context, *, hardware: str | None = None) -> None:
         """
         Donate your GPU to help power our Smart Eval command.
@@ -31,7 +83,10 @@ async def donate(self, ctx: commands.Context, *, hardware: str | None = None) ->
         """
         if await self.smarte_donation_cache.contains(ctx.author.id):
             stored_hardware = await self.smarte_donation_cache.get(ctx.author.id)
-            await ctx.reply(f"Thank you for donating your {stored_hardware} to our Smart Eval command.")
+            await ctx.reply(
+                "I can only take one donation per person. "
+                f"Thank you for donating your {stored_hardware} to our Smart Eval command."
+            )
             return
 
         if hardware is None:
@@ -47,17 +102,19 @@ async def donate(self, ctx: commands.Context, *, hardware: str | None = None) ->
         fake_hardware = ... # Do some regex to pull out a semi-matching type of GPU and insert something else
         await self.smarte_donation_cache.set(ctx.author.id, hardware)
 
-        self.total_donations = await self.smarte_donation_cache.length()
         await ctx.reply(
             "Thank you for donating your GPU to our Smart Eval command!"
             f" I did decide that instead of {hardware}, it would be better if you donated {fake_hardware}."
-            " So I've recorded that GPU donaton instead."
+            " So I've recorded that GPU donation instead."
             " It will be used wisely and definitely not for shenanigans."
         )
 
     @commands.command(aliases=["smarte"])
+    @commands.max_concurrency(1, commands.BucketType.user)
     async def smart_eval(self, ctx: commands.Context, *, code: str) -> None:
         """Evaluate your Python code with PyDis's newest chatbot."""
+        response_time, intelligence_level = await self.get_gpu_capabilities()
+
         if match := FORMATTED_CODE_REGEX.match(code):
             code = match.group("code")
         else:
@@ -81,4 +138,12 @@ async def smart_eval(self, ctx: commands.Context, *, code: str) -> None:
             matching_responses = DEFAULT_RESPONSES
         final_response = random.choice(matching_responses)
 
-        await ctx.reply(final_response)
+        async with ctx.typing():
+            await asyncio.sleep(response_time)
+
+            if len(final_response) <= 1000:
+                await ctx.reply(final_response)
+            else:
+                await ctx.reply(
+                    "There's definitely something wrong but I'm just not sure how to put it concisely into words."
+                )