diff --git a/reginaldCog/reginald.py b/reginaldCog/reginald.py index 930ce46..619b2f8 100644 --- a/reginaldCog/reginald.py +++ b/reginaldCog/reginald.py @@ -1,6 +1,11 @@ import openai import os import random +import requests +import base64 +import aiohttp +from io import BytesIO +from PIL import Image from redbot.core import Config, commands class ReginaldCog(commands.Cog): @@ -14,6 +19,19 @@ class ReginaldCog(commands.Cog): openai_api_key=None ) + def has_kanium_role(): + async def predicate(ctx): + kanium_role_id = 280260875678515200 + return any(role.id == kanium_role_id for role in ctx.author.roles) + return commands.check(predicate) + + def has_janitor_role(): + async def predicate(ctx): + janitor_role_id = 672156832323600396 + return any(role.id == janitor_role_id for role in ctx.author.roles) + return commands.check(predicate) + + @commands.guild_only() @commands.has_permissions(manage_guild=True) @commands.command(help="Set the OpenAI API key") @@ -22,6 +40,7 @@ class ReginaldCog(commands.Cog): await ctx.send("OpenAI API key set successfully.") @commands.guild_only() + @has_kanium_role() @commands.command(help="Ask Reginald a question") @commands.cooldown(1, 60, commands.BucketType.user) # 1-minute cooldown per user async def reginald(self, ctx, *, prompt=None): @@ -82,6 +101,73 @@ class ReginaldCog(commands.Cog): response_text = response_text[split_index:].strip() chunks.append(response_text) return chunks + + @commands.guild_only() + @has_janitor_role() + @commands.command(help="Ask Reginald to generate an image based on a prompt") + @commands.cooldown(1, 300, commands.BucketType.user) # 5-minute cooldown per user + async def reginaldimagine(self, ctx, *, prompt=None): + ignored_user_id = 138125632876838912 + if ctx.author.id == ignored_user_id: + return + + if prompt is None: + await ctx.author.send("Please provide a prompt for the image.") + return + + api_key = await self.config.guild(ctx.guild).openai_api_key() + if api_key is None: + await ctx.author.send('OpenAI API key not set. Please use the "!setreginaldcogapi" command to set the key.') + return + + try: + image_url = await self.generate_image(api_key, prompt) + if image_url: + async with ctx.typing(): + async with aiohttp.ClientSession() as session: + async with session.get(image_url) as resp: + image_data = await resp.read() + image = Image.open(BytesIO(image_data)) + image.save("image.png") + await ctx.send(file=discord.File("image.png")) + os.remove("image.png") + else: + await ctx.author.send("I apologize, but I am unable to generate an image based on the provided prompt.") + except Exception as e: + await ctx.author.send(f"I apologize, but I am unable to generate an image at this time. Error message: {str(e)}") + + async def generate_image(self, api_key, prompt): + openai.api_key = api_key + model = "image-alpha-001" # Replace this with the appropriate DALL-E model name + response = openai.Image.create( + model=model, + prompt=prompt, + n=1, + size="512x512", + response_format="url" + ) + + if response and response.choices and response.choices[0].url: + return response.choices[0].url + return None + + @reginald.error + async def reginald_error(ctx, error): + if isinstance(error, commands.BadArgument): + await ctx.send("I'm sorry, but I couldn't understand your input. Please check your message and try again.") + elif isinstance(error, commands.CheckFailure): + await ctx.send("You do not have the required role to use this command.") + else: + await ctx.send(f"An unexpected error occurred: {error}") + + @reginaldimagine.error + async def reginaldimagine_error(ctx, error): + if isinstance(error, commands.BadArgument): + await ctx.send("I'm sorry, but I couldn't understand your input. Please check your message and try again.") + elif isinstance(error, commands.CheckFailure): + await ctx.send("You do not have the required role to use this command.") + else: + await ctx.send(f"An unexpected error occurred: {error}") def setup(bot): cog = ReginaldCog(bot)