diff --git a/openai_cost_logger/openai_cost_logger.py b/openai_cost_logger/openai_cost_logger.py index 495dd9e..5bd64f0 100644 --- a/openai_cost_logger/openai_cost_logger.py +++ b/openai_cost_logger/openai_cost_logger.py @@ -42,6 +42,7 @@ def __init__( client_args (Dict, optional): The parameters to pass to the client. Defaults to {}. """ self.cost = 0 + self.n_responses = 0 self.model = model self.input_cost = input_cost self.log_folder = log_folder @@ -65,6 +66,7 @@ def update_cost(self, response: ChatCompletion) -> None: response: ChatCompletion object from the model. """ self.cost += self.__get_answer_cost(response) + self.n_responses += 1 self.__write_cost_to_json(response) self.__validate_cost() @@ -109,6 +111,7 @@ def __write_cost_to_json(self, response: ChatCompletion) -> None: with open(self.filepath, 'r') as file: data = json.load(file) data["total_cost"] = self.cost + data["total_responses"] = self.n_responses data["breakdown"].append(self.__build_log_breadown_entry(response)) with open(self.filepath, 'w') as file: json.dump(data, file, indent=4) @@ -126,6 +129,7 @@ def __build_log_file(self) -> None: "creation_datetime": strftime("%Y-%m-%d %H:%M:%S"), "model": self.model, "total_cost": self.cost, + "total_responses": 0, "breakdown": [] } with open(self.filepath, 'w') as file: