Merge pull request #32 from phixxy/main

Add Image logging

cogs/chatgpt.py

@@ -14,12 +14,15 @@ class ChatGPT(commands.Cog):
 
     def __init__(self, bot):
         self.bot = bot
+        self.admin_id = 242018983241318410
         self.API_KEY = os.getenv("openai.api_key")
+        self.default_budget = 20
         self.working_dir = "tmp/chatgpt/"
         self.data_dir = "data/chatgpt/"
         self.folder_setup()
         self.remind_me_loop.start()
         self.http_session = self.create_aiohttp_session()
+        self.dalle_budget = self.get_budget()
         self.logger = logging.getLogger("bot")
         self.headers = {
             'Content-Type': 'application/json',
@@ -57,6 +60,36 @@ class ChatGPT(commands.Cog):
         cost = input_cost + output_cost
         return cost
     
+    def get_budget(self):
+        month = time.strftime("%B")
+        year = time.strftime("%Y")
+        key = f"{month}_{year}"
+        budget_file = f"{self.data_dir}budget.json"
+        if not os.path.exists(budget_file):
+            with open(budget_file, "w") as f:
+                json.dump({key:self.default_budget},f)
+        with open(budget_file, "r") as f:
+            budget_dict = json.loads(f.readline())
+        if key not in budget_dict:
+            budget_dict[key] = self.default_budget
+            with open(budget_file, "w") as f:
+                json.dump(budget_dict,f)
+            return self.default_budget
+        else:
+            return budget_dict[key]
+        
+    def budget_add(self, amount):
+        month = time.strftime("%B")
+        year = time.strftime("%Y")
+        key = f"{month}_{year}"
+        budget_file = f"{self.data_dir}budget.json"
+        with open(budget_file, "r") as f:
+            budget_dict = json.loads(f.readline())
+        budget_dict[key] += amount
+        with open(budget_file, "w") as f:
+            json.dump(budget_dict,f)
+            self.dalle_budget += amount
+
     def add_cost(self, category: str, cost: float):
         day = time.strftime("%d")
         month = time.strftime("%B")
@@ -93,6 +126,56 @@ class ChatGPT(commands.Cog):
         plt.close()
         return f"{self.data_dir}costs/{graph_title}_categories.png"
     
+    async def get_monthly_cost(self, month=time.strftime("%B"), year=time.strftime("%Y")):
+        total_cost = 0
+        for x in range(1,32):
+            if x < 10:
+                x = f"0{x}"
+            filepath = f"{self.data_dir}costs/{month}_{x}_{year}.json"
+            if os.path.exists(filepath):
+                with open(filepath, "r") as f:
+                    costs_dict = json.loads(f.readline())
+                for category, cost in costs_dict.items():
+                    total_cost += cost
+        return total_cost
+    
+    async def moderation_check(self, prompt):
+        data = {
+            "input": prompt
+        }
+
+        url = "https://api.openai.com/v1/moderations"
+
+
+        async with self.http_session.post(url, json=data, headers=self.headers) as resp:
+            response_data = await resp.json()
+            flagged = response_data['results'][0]['flagged']
+            categories = response_data['results'][0]['categories']
+            category_scores = response_data['results'][0]['category_scores']
+            return (flagged, categories, category_scores)
+        
+    @commands.command()
+    async def budget(self, ctx, command=None, budget=None):
+        try:
+            if ctx.author.id == self.admin_id:
+                if command == "add" and budget!= None:
+                    self.budget_add(float(budget))
+                    await ctx.send(f"Budget increased by {budget}")
+                elif command == "remove" and budget!= None:
+                    self.budget_add(-float(budget))
+                    await ctx.send(f"Budget decreased by {budget}")
+                elif command == "set" and budget!= None:
+                    self.budget_add(float(budget) - self.dalle_budget)
+                    await ctx.send(f"Budget set to {budget}")
+                else:
+                    await ctx.send(f"The current budget is {self.dalle_budget}")
+            else:
+                await ctx.send(f"The current budget is {self.dalle_budget}")
+        except Exception as e:
+            self.logger.exception(f"Budget command failed: {e}")
+            await ctx.send(f"Budget command failed")
+
+
     @commands.command(
         name="costs",
         brief="Get the costs for a given month and year.",
@@ -101,7 +184,6 @@ class ChatGPT(commands.Cog):
         usage="costs [month] [year]",
     )
     async def costs(self, ctx, month=time.strftime("%B"), year=time.strftime("%Y")):
-        print('working')
         total_cost = 0
         cost_per_day = {}
         for x in range(1,32):
@@ -332,6 +414,9 @@ class ChatGPT(commands.Cog):
             await ctx.send(chunk)
     
     async def dalle_api_call(self, prompt: str, model: str="dall-e-2", quality: str="standard", size: str="1024x1024") -> tuple:
+        if self.dalle_budget <= await self.get_monthly_cost():
+            self.logger.info("DALL-E API call failed due to budget")
+            return (1337, "DALL-E API call failed due to budget. Consider using !donate to fund the bot.")
         data = {
             "model": model,
             "prompt": prompt,
@@ -357,10 +442,23 @@ class ChatGPT(commands.Cog):
                 await f.close()
             return resp.status
         
+    async def send_moderation_message(self, command, user_id, username, prompt, categories, category_scores):
+        categories = [k for k, v in categories.items() if v]
+        category_scores = {k: v for k, v in category_scores.items() if v > 0.5}
+        embed = discord.Embed(title="Moderation", description=f"Command: {command}\nUsername: {username}\nUser ID: {user_id}\nPrompt: {prompt}\nCategories: {categories}\nCategory Scores: {category_scores}", color=0x00ff00)
+        embed.set_footer(text="Moderation")
+        user = self.bot.get_user(self.admin_id)
+        await user.send(embed=embed)
 
     async def generate_dalle_image(self, ctx, model, quality="standard", size="1024x1024") -> None:
         prompt = ctx.message.content.split(" ", maxsplit=1)[1]
         await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.")
+        flagged, categories, category_scores = await self.moderation_check(prompt)
+        if flagged:
+            self.logger.info(f"Prompt {prompt} was flagged for inappropriate content.")
+            await ctx.send(f"This prompt {prompt} was flagged for inappropriate content. This has been reported.")
+            await self.send_moderation_message("dalle", ctx.author.id, ctx.author.name, prompt, categories, category_scores)
+            return
         resp_status, resp = await self.dalle_api_call(prompt, model=model, quality=quality, size=size)
         if resp_status != 200:
             await ctx.send(f"Error generating image: {resp_status}: {resp}")
@@ -376,6 +474,53 @@ class ChatGPT(commands.Cog):
             log_filepath.writelines(log_data)
         await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f)
 
+    @commands.command(
+        description="Big Spenders",
+        help="Generate a list of the biggest spenders. Usage: !bigspenders",
+        brief="Generate list of big spenders"
+        )
+    async def bigspenders(self, ctx):
+        filenames = os.listdir(self.data_dir + "logs/")
+        user_cost_dict = {}
+        for filename in filenames:
+            if ".log" in filename:
+                with open(f"{self.data_dir}logs/{filename}", 'r', encoding="utf-8") as f:
+                    for line in f:
+                        try:
+                            if "!dalle3hd" in line:
+                                cost = 0.08
+                                username = line[0:line.index(':')]
+                                if " " in username:
+                                    break
+                                if username not in user_cost_dict:
+                                    user_cost_dict[username] = 0
+                                user_cost_dict[username] += cost
+                            elif "!dalle2" in line:
+                                cost = 0.02
+                                username = line[0:line.index(':')]
+                                if " " in username:
+                                    break
+                                if username not in user_cost_dict:
+                                    user_cost_dict[username] = 0
+                                user_cost_dict[username] += cost
+                            if "!dalle" in line or "!dalle3" in line:
+                                cost = 0.04
+                                username = line[0:line.index(':')]
+                                if " " in username:
+                                    break
+                                if username not in user_cost_dict:
+                                    user_cost_dict[username] = 0
+                                user_cost_dict[username] += cost
+                            else:
+                                pass
+                        except:
+                            pass
+        message = "Big Spenders:\n"
+        sorted_dictionary = sorted(user_cost_dict.items(), key=lambda x: x[1], reverse=True)
+        for user in sorted_dictionary:
+            message += f"{user[0]}: ${user[1]:.2f}\n"
+        await ctx.send(message)
+
 
     @commands.command(
         description="Dalle 2", 

cogs/donate.py

@@ -10,6 +10,7 @@ class Donate(commands.Cog):
 
     def __init__(self, bot):
         self.bot = bot
+        self.admin_id = 242018983241318410
         self.donation_link = "https://patreon.com/phixxy"
         self.data_dir = "data/donate/"
         self.donor_file = "data/donate/supporters.txt"
@@ -31,6 +32,26 @@ class Donate(commands.Cog):
     async def donate(self, ctx):
         await ctx.send(f"If you would like to support the future of this bot please consider donating here: {self.donation_link}")
 
+    @commands.command()
+    async def add_supporter(self, ctx, username: str):
+        if ctx.author.id != self.admin_id:
+            return
+        with open(self.donor_file, 'a') as f:
+            f.write(f"{username}\n")
+        await ctx.send(f"Added {username} to supporters")
+
+    @commands.command()
+    async def remove_supporter(self, ctx, username: str):
+        if ctx.author.id != self.admin_id:
+            return
+        with open(self.donor_file, 'r') as f:
+            supporters = f.readlines()
+        with open(self.donor_file, 'w') as f:
+            for line in supporters:
+                if line!= username:
+                    f.write(line)
+        await ctx.send(f"Removed {username} from supporters")
+
     @commands.command(
         description="Supporters", 
         help="Get information on who supports this bot.", 

cogs/stable_diffusion.py

@@ -300,7 +300,7 @@ class StableDiffusion(commands.Cog):
         key_value_pairs = self.get_kv_from_ctx(ctx)
         negative_prompt = self.get_negative_prompt()
         if negative_prompt != self.default_neg_prompt:
-            ctx.send(f"Using non-default negative prompt: {negative_prompt}")
+            await ctx.send(f"Using non-default negative prompt: {negative_prompt}")
         headers = {'Content-Type': 'application/json'}  
         payload = {
             "prompt": prompt,