add json changes to viz as I forgot it exists while creating PR

openai_cost_logger/openai_cost_logger_viz.py

@@ -1,5 +1,6 @@
 import os
-import csv
+import json
+from datetime import datetime
 from typing import Dict
 from pathlib import Path
 import matplotlib.pyplot as plt
@@ -22,10 +23,8 @@ def get_total_cost(path: str = DEFAULT_LOG_PATH) -> float:
         cost = 0
         for filename in os.listdir(path):
             with open(Path(path, filename), mode='r') as file:
-                csvreader = csv.reader(file)
-                next(csvreader)
-                for row in csvreader:
-                    cost += float(row[2])
+                data = json.load(file)
+                cost += data["total_cost"]
         return cost
 
     @staticmethod
@@ -53,12 +52,10 @@ def get_total_cost_by_model(path: str = DEFAULT_LOG_PATH) -> Dict[str, float]:
         cost_by_model = defaultdict(float)
         for filename in os.listdir(path):
             with open(Path(path, filename), mode='r') as file:
-                csvreader = csv.reader(file)
-                next(csvreader)
-                for row in csvreader:
-                    if row[1] not in cost_by_model:
-                        cost_by_model[row[1]] = 0
-                    cost_by_model[row[1]] += float(row[2])
+                data = json.load(file)
+                if data["model"] not in cost_by_model:
+                    cost_by_model[data["model"]] = 0
+                cost_by_model[data["model"]] += data["total_cost"]
         return cost_by_model
 
     def print_total_cost_by_model(path: str = DEFAULT_LOG_PATH) -> None:
@@ -84,11 +81,10 @@ def plot_cost_by_day(path: str = DEFAULT_LOG_PATH, last_n_days: int = None) -> N
         cost_by_day = defaultdict(float)
         for filename in os.listdir(path):
             with open(Path(path, filename), mode='r') as file:
-                csvreader = csv.reader(file)
-                next(csvreader)
-                for row in csvreader:
-                    day = filename.split("_")[2]
-                    cost_by_day[day] += float(row[2])
+                data = json.load(file)
+                creation_datetime = data["creation_datetime"]
+                day = creation_datetime.split(' ')[0]
+                cost_by_day[day] += data["total_cost"]
 
         cost_by_day = dict(sorted(cost_by_day.items(), key=lambda x: x[0]))
         if last_n_days:
@@ -98,4 +94,34 @@ def plot_cost_by_day(path: str = DEFAULT_LOG_PATH, last_n_days: int = None) -> N
         plt.xlabel('Day')
         plt.ylabel('Cost [$]')
         plt.title('Cost by day')
+        plt.tight_layout()
+        plt.show()
+
+    @staticmethod
+    def plot_cost_by_strftime(path: str = DEFAULT_LOG_PATH, strftime_aggregator: str = "%Y-%m-%d", last_n_days: int = None) -> None:
+        """Plot the cost by day of all the logs in the directory.
+
+        Args:
+            path (str, optional): Cost logs directory. Defaults to DEFAULT_LOG_PATH.
+                                  This method reads all the files in the specified directory.
+            last_n_days (int, optional): The number of last days to plot. Defaults to None.
+        """
+        cost_by_aggregation_key = defaultdict(float)
+        for filename in os.listdir(path):
+            with open(Path(path, filename), mode='r') as file:
+                data = json.load(file)
+                creation_datetime = datetime.strptime(data["creation_datetime"], "%Y-%m-%d %H:%M:%S")
+                aggregation_key = creation_datetime.strftime(strftime_aggregator)
+                cost_by_aggregation_key[aggregation_key] += data["total_cost"]
+
+        cost_by_aggregation_key = dict(sorted(cost_by_aggregation_key.items(), key=lambda x: x[0]))
+        if last_n_days:
+            cost_by_aggregation_key = dict(list(cost_by_aggregation_key.items())[-last_n_days:])
+
+        plt.bar(cost_by_aggregation_key.keys(), cost_by_aggregation_key.values(), width=0.5)
+        plt.xticks(rotation=30, fontsize=8)
+        plt.xlabel('Day')
+        plt.ylabel('Cost [$]')
+        plt.title('Cost by day')
+        plt.tight_layout()
         plt.show()