added cost calculation and cost graphs

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 graph_cost(self, graph_title, costs_per_day):
-    async def costs(self, ctx):
+        plt.plot(costs_per_day.values())
-        day = time.strftime("%d")
+        plt.title(f"{graph_title} Costs")
-        month = time.strftime("%B")
+        plt.xlabel("Day")
-        year = time.strftime("%Y")
+        plt.ylabel("Cost")
-        filepath = f"{self.data_dir}costs/{month}_{day}_{year}.json"
+        plt.savefig(f"{self.data_dir}costs/{graph_title}_costs.png")
-        if not os.path.exists(filepath):
+        plt.close()
-            with open(filepath, "w") as f:
+        return f"{self.data_dir}costs/{graph_title}_costs.png"
-                json.dump({},f)
+    
-        with open(filepath, "r") as f:
+    async def graph_category(self, graph_title, cost_per_category):
-            costs_dict = json.loads(f.readline())
+        bar_container = plt.barh(list(cost_per_category.keys()), cost_per_category.values())
-        message = ""
+        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():
+        cost_per_day = {}
-            message += f"{category}: ${cost}\n"
+        for x in range(1,32):
-            total_cost += cost
+            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"
+        if rounded_total == 0:
-        await ctx.send(message)
+            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}"
+            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)
 
-            data = { 
-                "model": "gpt-3.5-turbo", 
-                "messages": [{"role": "user", "content": prompt}]
-            }
-
-            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):