diff --git a/cogs/chatgpt.py b/cogs/chatgpt.py index b336a5e..ede75a7 100644 --- a/cogs/chatgpt.py +++ b/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', @@ -56,6 +59,36 @@ class ChatGPT(commands.Cog): output_cost = cost_table[model]["output_tokens"] * output_tokens 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") @@ -92,6 +125,56 @@ class ChatGPT(commands.Cog): plt.savefig(f"{self.data_dir}costs/{graph_title}_categories.png") 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", @@ -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, @@ -356,11 +441,24 @@ class ChatGPT(commands.Cog): await f.write(await resp.read()) 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", diff --git a/cogs/donate.py b/cogs/donate.py index 8bb60c1..348fb57 100644 --- a/cogs/donate.py +++ b/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.", diff --git a/cogs/stable_diffusion.py b/cogs/stable_diffusion.py index e573cb4..d54e724 100644 --- a/cogs/stable_diffusion.py +++ b/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,