From 8c08eed8429fddfcd43b71c89d72b29bc76c4ddb Mon Sep 17 00:00:00 2001 From: crate Date: Sat, 7 Jun 2025 21:11:10 -0500 Subject: [PATCH] small refactor --- garfmain.py | 69 +++++++++++++++++++++++------------------------ garfpy/garfai.py | 43 +++++++++++++++++++---------- garfpy/weather.py | 12 +++++---- 3 files changed, 70 insertions(+), 54 deletions(-) diff --git a/garfmain.py b/garfmain.py index 4eeaa13..efe6ee6 100644 --- a/garfmain.py +++ b/garfmain.py @@ -27,6 +27,7 @@ intents.messages = True intents.message_content = True garfbot = commands.Bot(command_prefix=["garfbot ", "garf", "$"], intents=intents) +garfbot.remove_command("help") garf_respond = GarfbotRespond() garfield = GarfAI() @@ -48,7 +49,7 @@ async def on_ready(): @garfbot.command(name="ping") -async def ping(ctx, *, target): +async def garfbot_ping(ctx, *, target): """Ping a target""" logger.info( f"Ping Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Target: {target}" @@ -57,7 +58,7 @@ async def ping(ctx, *, target): @garfbot.command(name="dns") -async def dns(ctx, *, target): +async def garfbot_dns(ctx, *, target): """DNS lookup for a target""" logger.info( f"NSLookup Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Target: {target}" @@ -66,7 +67,7 @@ async def dns(ctx, *, target): @garfbot.command(name="hack") -async def hack(ctx, *, target): +async def garfbot_hack(ctx, *, target): """Nmap scan a target""" logger.info( f"Nmap Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Target: {target}" @@ -108,13 +109,30 @@ async def garfbot_shop(ctx, *, query): @garfbot.command(name="weather") async def garfbot_weather(ctx, *, location): - embed = await weather.weather(location) - await ctx.send(embed=embed) + await weather.weather(ctx, location) -# @garfbot.command(name="help") -# async def garfbot_help(ctx): -# await help(ctx) +@garfbot.command(name="chat") +async def garfchat(ctx, *, prompt): + answer = await garfield.generate_chat(prompt) + logger.info( + f"Chat Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Prompt: {prompt}" + ) + await ctx.send(answer) + + +@garfbot.command(name="pic") +async def garfpic(ctx, *, prompt): + logger.info( + f"Image Request - User: {ctx.author.name}, Server: {ctx.guild.name}, Prompt: {prompt}" + ) + await ctx.send(f"`Please wait... image generation queued: {prompt}`") + await garfield.garfpic(ctx, prompt) + + +@garfbot.command(name="help") +async def garfbot_help(ctx): + await help(ctx) @garfbot.event @@ -124,41 +142,22 @@ async def on_message(message): content = message.content.strip() lower = content.lower() - user_name = message.author.name - guild_id = message.guild.id - guild_name = message.guild.name if message.guild else "Direct Message" # Chats & pics if lower.startswith("hey garfield") or isinstance( message.channel, discord.DMChannel ): - prompt = content[12:] if lower.startswith("hey garfield") else message.content - answer = await garfield.generate_chat(prompt) - logger.info( - f"Chat Request - User: {user_name}, Server: {guild_name}, Prompt: {prompt}" - ) - await message.channel.send(answer) - - if lower.startswith("garfpic "): - prompt = content[8:] - logger.info( - f"Image Request - User: {user_name}, Server: {guild_name}, Prompt: {prompt}" - ) - await message.channel.send( - f"`Please wait... image generation queued: {prompt}`" - ) - await garfield.garfpic(message, prompt) - - # GarfBot help - elif lower.strip() == "garfbot help": - await help(message) - - # Army of Dawn Server only!! - elif message.guild and message.guild.id == 719605634772893757: - await aod_message(garfbot, message) + ctx = await garfbot.get_context(message) + await garfchat(ctx, prompt=content) # Auto-responses elif message.guild: + guild_id = message.guild.id + + # Army of Dawn Server only!! + if guild_id == 719605634772893757: + await aod_message(garfbot, message) + responses = garf_respond.get_responses(guild_id) if lower.startswith("garfbot response "): diff --git a/garfpy/garfai.py b/garfpy/garfai.py index b3f81ec..f46e2d5 100644 --- a/garfpy/garfai.py +++ b/garfpy/garfai.py @@ -16,17 +16,15 @@ class GarfAI: self.imgmodel = config.IMG_MODEL self.image_request_queue = asyncio.Queue() - async def garfpic(self, message, prompt): - await self.image_request_queue.put({"message": message, "prompt": prompt}) + async def garfpic(self, ctx, prompt): + await self.image_request_queue.put({"ctx": ctx, "prompt": prompt}) async def generate_image(self, prompt): + client = AsyncOpenAI(api_key=self.openaikey) try: - client = AsyncOpenAI(api_key=self.openaikey) response = await client.images.generate( model=self.imgmodel, prompt=prompt, n=1, size="1024x1024" ) - image_url = response.data[0].url - return image_url except openai.BadRequestError as e: return f"`GarfBot Error: ({e.status_code}) - Your request was rejected as a result of our safety system.`" except openai.InternalServerError as e: @@ -35,32 +33,48 @@ class GarfAI: except Exception as e: logger.error(e) return "`GarfBot Error: Lasagna`" + data = getattr(response, "data", None) + if not data: + logger.error("No data in response") + return "`GarfBot Error: No images generated`" + + first_image = data[0] if len(data) > 0 else None + if not first_image: + logger.error("No image in response data") + return "`GarfBot Error: No images generated`" + + image_url = getattr(first_image, "url", None) + if not image_url: + logger.error("No URL in image response") + return "`GarfBot Error: No image URL returned`" + + return image_url async def process_image_requests(self): async with aiohttp.ClientSession() as session: while True: request = await self.image_request_queue.get() - message = request["message"] + ctx = request["ctx"] prompt = request["prompt"] image_url = await self.generate_image(prompt) - if "GarfBot Error" not in image_url: + if image_url and "GarfBot Error" not in image_url: logger.info("Downloading & sending image...") async with session.get(image_url) as resp: if resp.status == 200: image_data = await resp.read() image = io.BytesIO(image_data) image.seek(0) - timestamp = message.created_at.strftime("%Y%m%d%H%M%S") + timestamp = ctx.message.created_at.strftime("%Y%m%d%H%M%S") filename = f"{timestamp}_generated_image.png" sendfile = discord.File(fp=image, filename=filename) try: - await message.channel.send(file=sendfile) + await ctx.send(file=sendfile) except Exception as e: logger.error(e) else: - await message.channel.send("`GarfBot Error: Odie`") + await ctx.send("`GarfBot Error: Odie`") else: - await message.channel.send(image_url) + await ctx.send(image_url) self.image_request_queue.task_done() await asyncio.sleep(2) @@ -78,15 +92,16 @@ class GarfAI: ], max_tokens=400, ) - answer = response.choices[0].message.content + answer = str(response.choices[0].message.content) return answer.replace("an AI language model", "a cartoon animal") except openai.BadRequestError as e: + logger.error(e) return f"`GarfBot Error: {e}`" except openai.APIError as e: - logger.info(e, flush=True) + logger.error(e) return "`GarfBot Error: Monday`" except Exception as e: - logger.info(e, flush=True) + logger.error(e) return "`GarfBot Error: Lasagna`" async def wikisum(self, query): diff --git a/garfpy/weather.py b/garfpy/weather.py index afb87bb..4927a19 100644 --- a/garfpy/weather.py +++ b/garfpy/weather.py @@ -110,7 +110,7 @@ class WeatherAPI: city_name = " ".join(params) return {"q": f"{city_name},US"} - async def get_weather(self, location, units="metric"): + async def get_weather(self, ctx, location, units="metric"): location_params = self.parse_location(location) params = { @@ -126,6 +126,7 @@ class WeatherAPI: return await response.json() except aiohttp.ClientError as e: logger.error(f"Error fetching weather data for '{location}': {e}") + ctx.send(f"`Error fetching weather data for '{location}': {e}`") return None def weather_embed(self, weather_data): @@ -196,7 +197,8 @@ class WeatherAPI: return embed - async def weather(self, location): - weather_data = await self.get_weather(location) - embed = self.weather_embed(weather_data) - return embed + async def weather(self, ctx, location): + weather_data = await self.get_weather(ctx, location) + if weather_data: + embed = self.weather_embed(weather_data) + await ctx.send(embed=embed)