From 8bf92a4a535e12d4e3aa78c602c317f32d10a200 Mon Sep 17 00:00:00 2001 From: LobaDK Date: Sat, 18 May 2024 18:01:02 +0200 Subject: [PATCH 1/5] Moved init_models() to the on_ready() event --- QuantumKat.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/QuantumKat.py b/QuantumKat.py index 7cc84df..7cccc8f 100644 --- a/QuantumKat.py +++ b/QuantumKat.py @@ -24,7 +24,6 @@ async def init_models(): await conn.run_sync(models.Base.metadata.create_all) -# run(init_models()) log_helper = LogHelper() misc_helper = MiscHelper() @@ -166,6 +165,8 @@ async def on_guild_join(guild): @bot.event async def on_ready(): + # await init_models() + # Add all servers the bot is in to the database on startup in case the bot was added while offline for guild in bot.guilds: try: From b572bbdc3fc6ca766926a3b37d61fe7ae154290a Mon Sep 17 00:00:00 2001 From: LobaDK Date: Sat, 18 May 2024 18:02:00 +0200 Subject: [PATCH 2/5] Improved Content-Type header detection and removed TODO's --- cogs/utils/utils.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/cogs/utils/utils.py b/cogs/utils/utils.py index c2d5c38..a44a4a6 100644 --- a/cogs/utils/utils.py +++ b/cogs/utils/utils.py @@ -81,7 +81,7 @@ def header_mime_type(self) -> str: str: The MIME type from the 'Content-Type' header, or None if not present. """ try: - return self.header["Content-Type"] + return self.header["Content-Type"].split(";")[0] except KeyError: return None @@ -189,10 +189,6 @@ def strip_embed_disabler(url: str) -> str: return url.replace("<", "").replace(">", "") -# TODO: It shouldn't necessarily be added here, but add a function to detect and get files attached to a ctx object (ctx.message.attachments) -# TODO: While we're at it with the above, add the ability (unsure if function is necessary) for the bot to include the message a user is replying to i.e. they initiated the chat command while replying to a message -# TODO: Improve file type detection and minimize the amount of data being fetched before the checks (check the Content-Type header of the response) -# TODO: Add some fallbacks where we may still download the file if the Content-Type header is not available, but only the first 1 KB or so def get_image_as_base64(url_or_byte_stream: str | bytes) -> list[str]: """ Converts an image from a URL or byte stream into a base64 encoded string. From c70f539f1e6351478948ea4be39399bdedc768b8 Mon Sep 17 00:00:00 2001 From: LobaDK Date: Sat, 18 May 2024 18:03:24 +0200 Subject: [PATCH 3/5] Fixed attachment URL's being used instead of their byte data --- cogs/Chat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cogs/Chat.py b/cogs/Chat.py index eea6ea7..0c50e8f 100644 --- a/cogs/Chat.py +++ b/cogs/Chat.py @@ -207,7 +207,7 @@ async def initiateChat( for attachment in ctx.message.attachments: try: base64_images.extend( - get_image_as_base64(attachment.url) + get_image_as_base64(await attachment.read()) ) except ( UnsupportedImageFormatError, From b835f16cb7e48ee7f4e1a9bccd09126d4d6288a6 Mon Sep 17 00:00:00 2001 From: LobaDK Date: Sat, 18 May 2024 18:05:41 +0200 Subject: [PATCH 4/5] Include the message the user is replying to as a system message reference, allowing ChatGPT to get additional context --- cogs/Chat.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/cogs/Chat.py b/cogs/Chat.py index 0c50e8f..5620616 100644 --- a/cogs/Chat.py +++ b/cogs/Chat.py @@ -21,6 +21,7 @@ from QuantumKat import log_helper, misc_helper, discord_helper +TOKEN_LIMIT = 1024*2 class Chat(commands.Cog): def __init__(self, bot: commands.Bot): @@ -178,8 +179,11 @@ async def initiateChat( """ if self.FOUND_API_KEY is True: if user_message: - tokens = calculate_tokens(user_message, self.system_message) - if not tokens > 1024: + system_message = self.system_message + if ctx.message.reference: + system_message += f"The user has included this message in their response, which was written by {ctx.message.reference.resolved.author.display_name}: `{ctx.message.reference.resolved.content}`. Use it as context for the response." + tokens = calculate_tokens(user_message, system_message) + if not tokens > TOKEN_LIMIT: command = ctx.invoked_with user_message = ctx.message.content.split( f"{self.bot.command_prefix}{command}", 1 @@ -239,7 +243,7 @@ async def initiateChat( messages = [ { "role": "system", - "content": self.system_message.format( + "content": system_message.format( user=ctx.author.id, version=".".join( str(misc_helper.get_git_commit_count()) @@ -304,7 +308,7 @@ async def initiateChat( ) else: await ctx.reply( - f"Message is too long! Your message is {tokens} tokens long, but the maximum is 1024 tokens.", + f"Message is too long! Your message is {tokens} tokens long, but the maximum is {TOKEN_LIMIT} tokens.", silent=True, ) else: From fdd8d4e256c31682e9de79324210a1182f2fae37 Mon Sep 17 00:00:00 2001 From: LobaDK Date: Sat, 18 May 2024 19:38:04 +0200 Subject: [PATCH 5/5] Fixed is_banned check breaking on DM's --- QuantumKat.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/QuantumKat.py b/QuantumKat.py index 7cccc8f..6a0169b 100644 --- a/QuantumKat.py +++ b/QuantumKat.py @@ -24,7 +24,6 @@ async def init_models(): await conn.run_sync(models.Base.metadata.create_all) - log_helper = LogHelper() misc_helper = MiscHelper() discord_helper = DiscordHelper() @@ -91,7 +90,7 @@ async def setup(bot: commands.Bot): async def is_authenticated(ctx: commands.Context) -> bool: - if ctx.author.id == int(OWNER_ID): + if ctx.author.id in bot.owner_ids: return True if not ctx.command.name.casefold() == "request_auth": authenticated_server_ids = await crud.get_authenticated_servers( @@ -119,7 +118,7 @@ async def is_authenticated(ctx: commands.Context) -> bool: async def is_reboot_scheduled(ctx: commands.Context) -> bool: - if ctx.author.id == int(OWNER_ID): + if ctx.author.id in bot.owner_ids: return True if bot.reboot_scheduled: await ctx.reply( @@ -131,26 +130,28 @@ async def is_reboot_scheduled(ctx: commands.Context) -> bool: async def is_banned(ctx: commands.Context) -> bool: - if ctx.author.id == int(OWNER_ID): + if ctx.author.id in bot.owner_ids: return True user = await crud.get_user( AsyncSessionLocal, schemas.User.Get(user_id=ctx.author.id) ) - server = await crud.get_server( - AsyncSessionLocal, schemas.Server.Get(server_id=ctx.guild.id) - ) if user and user.is_banned: await ctx.reply( "You have been banned from using QuantumKat. Please contact the bot owner for more information.", silent=True, ) return False - if server and server.is_banned: - await ctx.reply( - "This server has been banned from using QuantumKat. Please contact the bot owner for more information.", - silent=True, + + if not discord_helper.is_dm(ctx): + server = await crud.get_server( + AsyncSessionLocal, schemas.Server.Get(server_id=ctx.guild.id) ) - return False + if server and server.is_banned: + await ctx.reply( + "This server has been banned from using QuantumKat. Please contact the bot owner for more information.", + silent=True, + ) + return False return True