Refactor API and utils modules, add cooldown checks

.github/workflows/format.yaml

@@ -0,0 +1,12 @@
+on: [push, pull_request]
+
+jobs:
+  python-black:
+    name: Python Black
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Python Black
+        uses: cytopia/[email protected]
+        with:
+          path: '.'

api.py

@@ -8,14 +8,16 @@
 def home():
     return "Hello, I am alive"
 
+
 @app.route("/health")
 def health():
     return jsonify({"status": "ok", "message": "I am alive"}), 200
 
+
 def run():
     app.run(host="0.0.0.0", port=8080)
 
 
 def keep_alive():
     t = Thread(target=run)
-    t.start()
+    t.start()

cogs/imagine_cog.py

@@ -150,6 +150,7 @@ async def model_autocomplete(self,
 
     @app_commands.command(name="imagine", description="Imagine a prompt")
     @app_commands.autocomplete(model=model_autocomplete)
+    @app_commands.checks.cooldown(1, 15)
     @app_commands.describe(prompt="Imagine a prompt", height="Height of the image", width="Width of the image", negative="The things not to include in the image", cached="Removes the image seed", nologo="Remove the logo", enhance="Disables Prompt enhancing if set to False", private="Only you can see the generated Image if set to True")
     async def imagine_command(self, interaction, prompt:str, model: str = "Dreamshaper", width:int = 1000, height:int = 1000, negative:str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True, private:bool = False):
         await interaction.response.send_message(embed=discord.Embed(title="Generating Image", description="Please wait while we generate your image", color=discord.Color.blurple()), ephemeral=True)
@@ -192,6 +193,7 @@ async def imagine_command(self, interaction, prompt:str, model: str = "Dreamshap
         return
 
     @app_commands.command(name="multi-imagine", description="Imagine multiple prompts")
+    @app_commands.checks.cooldown(1, 30)
     @app_commands.describe(prompt="Imagine a prompt", height="Height of the image", width="Width of the image", negative="The things not to include in the image", cached="Removes the image seed", nologo="Remove the logo", enhance="Disables Prompt enhancing if set to False", private="Only you can see the generated Image if set to True")
     async def multiimagine_command(self, interaction, prompt:str, width:int = 1000, height:int = 1000, negative:str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True, private:bool = False):
 
@@ -241,6 +243,55 @@ async def multiimagine_command(self, interaction, prompt:str, width:int = 1000,
 
         return
 
+    @imagine_command.error
+    async def imagine_command_error(
+        self, interaction: discord.Interaction, error: app_commands.AppCommandError
+    ):
+        if isinstance(error, app_commands.CommandOnCooldown):
+            end_time = datetime.datetime.now() + datetime.timedelta(
+                seconds=error.retry_after
+            )
+            end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S UTC")
+            end_time_ts = f"<t:{int(end_time.timestamp())}>"
+
+            hours, remainder = divmod(error.retry_after, 3600)
+            minutes, seconds = divmod(remainder, 60)
+            seconds = round(seconds)
+            time_left = f"{ hours + ' hour, ' if not hours<1 else ''}{int(minutes)} minute{'s' if minutes != 1 else ''} and {seconds} second{'s' if seconds != 1 else ''}"
+
+            embed = discord.Embed(
+                title="⏳ Cooldown",
+                description=f"You have to wait until **{end_time_ts}** ({time_left}) before using the </imagine:1123582901544558612> command again.",
+                color=discord.Color.red(),
+            )
+
+            await interaction.response.send_message(embed=embed, ephemeral=True)
+
+    @multiimagine_command.error
+    async def multiimagine_command_error(
+        self, interaction: discord.Interaction, error: app_commands.AppCommandError
+    ):
+        if isinstance(error, app_commands.CommandOnCooldown):
+            end_time = datetime.datetime.now() + datetime.timedelta(
+                seconds=error.retry_after
+            )
+            end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S UTC")
+            end_time_ts = f"<t:{int(end_time.timestamp())}>"
+
+            hours, remainder = divmod(error.retry_after, 3600)
+            minutes, seconds = divmod(remainder, 60)
+            seconds = round(seconds)
+            time_left = f"{ hours + ' hour, ' if not hours<1 else ''}{int(minutes)} minute{'s' if minutes != 1 else ''} and {seconds} second{'s' if seconds != 1 else ''}"
+
+            embed = discord.Embed(
+                title="⏳ Cooldown",
+                description=f"You have to wait until **{end_time_ts}** ({time_left}) before using the <multi-imagine:1187375074722975837> again.",
+                color=discord.Color.red(),
+            )
+
+            await interaction.response.send_message(embed=embed, ephemeral=True)
+
+
 async def setup(bot):
     await bot.add_cog(Imagine(bot))
     print("Imagine cog loaded")

main.py

@@ -20,6 +20,31 @@
 start_time = None
 latencies = []
 
+commands_ = {
+    "</imagine:1123582901544558612> 🎨": """Generates AI Images based on your prompts
+- **prompt** 🗣️ : Your prompt for the Image to be generated
+- **model** 🤖 : The model to be used for generating the Image
+- **width** ↔️ : The width of your prompted Image
+- **height** ↕️ : The height of your prompted Image
+- **cached** : specifies whether to return a cached image
+- **negative** ❎ : Specifies what not to be in the generated images
+- **nologo** 🚫 : Specifies whether to remove the logo from the generated images (deafault False)
+- **enhance** 🖼️ : Specifies whether to enhance the image prompt or not (default True)
+- **private** 🔒 : when set to True the generated Image will only be visible to you
+""",
+    "</multi-imagine:1187375074722975837> 🎨": """Generates AI Images using all available models
+- **prompt** 🗣️ : Your prompt for the Image to be generated
+- **width** ↔️ : The width of your prompted Image
+- **height** ↕️ : The height of your prompted Image
+- **cached** : specifies whether to return a cached image
+- **negative** ❎ : Specifies what not to be in the generated images
+- **nologo** 🚫 : Specifies whether to remove the logo from the generated images (deafault False)
+- **enhance** 🖼️ : Specifies whether to enhance the image prompt or not (default True)
+- **private** 🔒 : when set to True the generated Image will only be visible to you
+""",
+    "</help:1125407202388230155> ❓": "Displays this",
+}
+
 
 class pollinationsBot(commands.Bot):
     def __init__(self):
@@ -53,15 +78,24 @@ async def on_ready(self):
 
 
 async def load():
-    for filename in os.listdir('./cogs'):
-        if filename.endswith('.py'):
-            await bot.load_extension(f'cogs.{filename[:-3]}')
+    for filename in os.listdir("./cogs"):
+        if filename.endswith(".py"):
+            await bot.load_extension(f"cogs.{filename[:-3]}")
+
 
 @bot.event
 async def on_message(message):
     if message.author == bot.user:
         return
 
+    if bot.user in message.mentions:
+        embed = discord.Embed(
+            description="Hello, I am the Pollinations.ai Bot. I am here to help you with your AI needs. Type `!help` or click </help:1125407202388230155> to get started.",
+            color=discord.Color.og_blurple(),
+        )
+
+        await message.reply(embed=embed)
+
     await bot.process_commands(message)
 
 
@@ -102,16 +136,16 @@ async def ping(ctx):
 
         latency = (end - ctx.start) * 1000
 
+        embed.add_field(name="Ping", value=f"{bot.latency * 1000:.2f} ms", inline=False)
         embed.add_field(name="Message Latency", value=f"{latency:.2f} ms", inline=False)
-        embed.add_field(
-            name="Websocket Latency", value=f"{bot.latency * 1000:.2f} ms", inline=False
-        )
 
         # Calculate the average ping of the bot in the last 10 minutes
         if latencies:
             average_ping = statistics.mean(latencies)
             embed.add_field(
-                name="Average Ping", value=f"{average_ping:.2f} ms", inline=False
+                name="Average Message Latency",
+                value=f"{average_ping:.2f} ms",
+                inline=False,
             )
 
         global start_time
@@ -141,31 +175,30 @@ async def ping(ctx):
         print(e, file=sys.stdout)
 
 
-# @bot.hybrid_command(name="help", description="View the various commands of this server")
-# async def help(ctx):
-#     user = bot.get_user(1168874960081649684)
-#     profilePicture = user.avatar.url
+@bot.hybrid_command(name="help", description="View the various commands of this server")
+async def help(ctx):
+    user = bot.get_user(1123551005993357342)
+    profilePicture = user.avatar.url
 
-#     embed = discord.Embed(
-#         title="SecretSanctuary Bot Commands",
-#         url=APP_URI,
-#         description="Here is the list of the available commands:",
-#         color=discord.Color.og_blurple(),
-#     )
+    embed = discord.Embed(
+        title="Pollinations.ai Bot Commands",
+        url=APP_URI,
+        description="Here is the list of the available commands:",
+        color=discord.Color.og_blurple(),
+    )
 
-#     embed.set_thumbnail(url=profilePicture)
-#     for i in commands_.keys():
-#         embed.add_field(name=i, value=commands_[i], inline=False)
+    embed.set_thumbnail(url=profilePicture)
+    for i in commands_.keys():
+        embed.add_field(name=i, value=commands_[i], inline=False)
 
-#     embed.set_footer(
-#         text="Information requested by: {}".format(ctx.author.name),
-#         icon_url=ctx.author.avatar.url,
-#     )
+    embed.set_footer(
+        text="Information requested by: {}".format(ctx.author.name),
+        icon_url=ctx.author.avatar.url,
+    )
 
-#     await ctx.send(embed=embed)
+    await ctx.send(embed=embed)
 
 
 if __name__ == "__main__":
-    # run_server()
     keep_alive()
     bot.run(token=TOKEN)

utils.py

@@ -18,36 +18,51 @@
 db = client["pollinations"]
 collection = db["prompts"]
 
+
 def get_prompt_data(message_id: int):
     try:
         return collection.find_one({"_id": message_id})
     except Exception as e:
         print(e)
         return None
 
+
 def save_prompt_data(message_id: int, data: dict):
     try:
         collection.insert_one(data)
     except Exception as e:
         print(e)
 
+
 def update_prompt_data(message_id: int, data: dict):
     try:
         collection.update_one({"_id": message_id}, {"$set": data})
     except Exception as e:
         print(e)
 
+
 def delete_prompt_data(message_id: int):
     try:
         collection.delete_one({"_id": message_id})
     except Exception as e:
         print(e)
 
-async def generate_image(prompt: str, width: int = 500, height: int = 500, model: str = "turbo", negative: str|None = None, cached:bool = False, nologo:bool = False, enhance:bool = True):
-
+
+async def generate_image(
+    prompt: str,
+    width: int = 500,
+    height: int = 500,
+    model: str = "turbo",
+    negative: str | None = None,
+    cached: bool = False,
+    nologo: bool = False,
+    enhance: bool = True,
+):
     model = model.lower()
 
-    print(f"Generating image with prompt: {prompt}, width: {width}, height: {height}, model: {model}, negative: {negative}, cached: {cached}, nologo: {nologo}, enhance: {enhance}")
+    print(
+        f"Generating image with prompt: {prompt}, width: {width}, height: {height}, model: {model}, negative: {negative}, cached: {cached}, nologo: {nologo}, enhance: {enhance}"
+    )
 
     seed = str(random.randint(0, 1000000000))
 
@@ -69,7 +84,7 @@ async def generate_image(prompt: str, width: int = 500, height: int = 500, model
         "cached": cached,
         "nologo": nologo,
         "enhance": enhance,
-        "bookmark_url": quote(url, safe=':/&=?'),
+        "bookmark_url": quote(url, safe=":/&=?"),
     }
 
     dic["seed"] = seed if not cached else None