diff --git a/cogs/chatgpt.py b/cogs/chatgpt.py index cf9164c..01ab713 100644 --- a/cogs/chatgpt.py +++ b/cogs/chatgpt.py @@ -8,6 +8,7 @@ import aiofiles import aiohttp import discord from discord.ext import commands, tasks +import matplotlib.pyplot as plt class ChatGPT(commands.Cog): @@ -46,7 +47,17 @@ class ChatGPT(commands.Cog): except Exception as e: self.logger.exception(f"ChatGPT failed to make directories: {e}") - def add_cost(self, category, cost): + def text_cost_calc(self, model, input_tokens, output_tokens): + cost_table = {"gpt-3.5-turbo":{"input_tokens":0.0000005,"output_tokens":0.0000015}, + "gpt-4-turbo-preview":{"input_tokens":0.00001,"output_tokens":0.00003}, + "gpt-4-vision-preview":{"input_tokens":0.00001,"output_tokens":0.00003} + } + input_cost = cost_table[model]["input_tokens"] * input_tokens + output_cost = cost_table[model]["output_tokens"] * output_tokens + cost = input_cost + output_cost + return cost + + def add_cost(self, category: str, cost: float): day = time.strftime("%d") month = time.strftime("%B") year = time.strftime("%Y") @@ -63,25 +74,89 @@ class ChatGPT(commands.Cog): with open(filepath, "w") as f: json.dump(costs_dict,f) - @commands.command() - async def costs(self, ctx): - day = time.strftime("%d") - month = time.strftime("%B") - year = time.strftime("%Y") - filepath = f"{self.data_dir}costs/{month}_{day}_{year}.json" - if not os.path.exists(filepath): - with open(filepath, "w") as f: - json.dump({},f) - with open(filepath, "r") as f: - costs_dict = json.loads(f.readline()) - message = "" + async def graph_cost(self, graph_title, costs_per_day): + plt.plot(costs_per_day.values()) + plt.title(f"{graph_title} Costs") + plt.xlabel("Day") + plt.ylabel("Cost") + plt.savefig(f"{self.data_dir}costs/{graph_title}_costs.png") + plt.close() + return f"{self.data_dir}costs/{graph_title}_costs.png" + + async def graph_category(self, graph_title, cost_per_category): + bar_container = plt.barh(list(cost_per_category.keys()), cost_per_category.values()) + plt.title(f"{graph_title} Costs") + plt.xlabel("Category") + plt.ylabel("Cost") + plt.bar_label(bar_container) + plt.savefig(f"{self.data_dir}costs/{graph_title}_categories.png") + plt.close() + return f"{self.data_dir}costs/{graph_title}_categories.png" + + @commands.command( + name="costs", + brief="Get the costs for a given month and year.", + help="Get the costs for a given month and year. If no month or year is given, it will default to the current month and year.", + aliases=["cost"], + usage="costs [month] [year]", + ) + async def costs(self, ctx, month=time.strftime("%B"), year=time.strftime("%Y")): + print('working') total_cost = 0 - for category, cost in costs_dict.items(): - message += f"{category}: ${cost}\n" - total_cost += cost + cost_per_day = {} + for x in range(1,32): + daily_cost = 0 + 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 + daily_cost += cost + cost_per_day[x] = daily_cost rounded_total = round(total_cost, 2) - message += f"Total: ${rounded_total}\n" - await ctx.send(message) + if rounded_total == 0: + await ctx.send(f"No data for {month} {year}") + else: + graph_title = f"{month} {year}" + await ctx.send(f"Total cost for {month} {year}: ${rounded_total}", file=discord.File(await self.graph_cost(graph_title, cost_per_day))) + + @commands.command( + name="category_costs", + brief="Get the costs per category for a given month and year.", + help="Get the costs per category for a given month and year. If no month or year is given, it will default to the current month and year.", + aliases=["category_cost"], + usage="category_costs [month] [year] [day]", + ) + async def category_costs(self, ctx, month=time.strftime("%B"), year=time.strftime("%Y"), day=None): + total_cost = 0 + cost_per_category = {} + if day == None: + day_range = range(1,32) + else: + day_range = [int(day)] + for x in day_range: + daily_cost = 0 + 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(): + if category not in cost_per_category: + cost_per_category[category] = cost + else: + cost_per_category[category] += cost + total_cost += cost + graph_title = f"{month} {year}" + rounded_total = round(total_cost, 2) + if rounded_total == 0: + await ctx.send(f"No data for {month} {year}") + else: + await ctx.send(f"Total cost for {month} {year}: ${rounded_total}", file=discord.File(await self.graph_category(graph_title, cost_per_category))) def create_channel_config(self, filepath): config_dict = { @@ -138,6 +213,10 @@ class ChatGPT(commands.Cog): async with self.http_session.post(url, json=data, headers=self.headers) as resp: response_data = await resp.json() response = response_data['choices'][0]['message']['content'] + input_tokens = response_data['usage']['prompt_tokens'] + output_tokens = response_data['usage']['completion_tokens'] + cost = self.text_cost_calc(model, input_tokens, output_tokens) + self.add_cost(model, cost) return response except Exception as error: @@ -448,30 +527,17 @@ class ChatGPT(commands.Cog): async with ctx.channel.typing(): await asyncio.sleep(1) prompt = f"You are a {channel_vars['personality']} chat bot named Sparkytron 3000 created by @phixxy.com. Your personality should be {channel_vars['personality']}. You are currently in a {channel_vars['channel_topic']} chatroom. The message history is: {chat_history_string}" - - data = { - "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": prompt}] - } + response = await self.answer_question(prompt) + if "Sparkytron 3000:" in response[0:17]: + response = response.replace("Sparkytron 3000:", "") + max_len = 1999 + if len(response) > max_len: + messages=[response[y-max_len:y] for y in range(max_len, len(response)+max_len,max_len)] + else: + messages=[response] + for message in messages: + await ctx.channel.send(message) - url = "https://api.openai.com/v1/chat/completions" - - try: - async with self.http_session.post(url, json=data, headers=self.headers) as resp: - response_data = await resp.json() - response = response_data['choices'][0]['message']['content'] - if "Sparkytron 3000:" in response[0:17]: - response = response.replace("Sparkytron 3000:", "") - max_len = 1999 - if len(response) > max_len: - messages=[response[y-max_len:y] for y in range(max_len, len(response)+max_len,max_len)] - else: - messages=[response] - for message in messages: - await ctx.channel.send(message) - - except Exception as error: - self.logger.exception("Problem with chat_response in chatgpt") @commands.Cog.listener() async def on_reaction_add(self, reaction, user):