diff --git a/pkg/audit/gatherer.py b/pkg/audit/gatherer.py index 6e0ea3ed..8e118f1a 100644 --- a/pkg/audit/gatherer.py +++ b/pkg/audit/gatherer.py @@ -43,7 +43,7 @@ class DataGatherer: def get_usage(self, key_md5): return self.usage[key_md5] if key_md5 in self.usage else {} - def report_text_model_usage(self, model, text): + def report_text_model_usage(self, model, total_tokens): key_md5 = pkg.utils.context.get_openai_manager().key_mgr.get_using_key_md5() if key_md5 not in self.usage: @@ -55,7 +55,7 @@ class DataGatherer: if model not in self.usage[key_md5]["text"]: self.usage[key_md5]["text"][model] = 0 - length = int((len(text.encode('utf-8')) - len(text)) / 2 + len(text)) + length = total_tokens self.usage[key_md5]["text"][model] += length self.dump_to_db() diff --git a/pkg/openai/manager.py b/pkg/openai/manager.py index 105b454b..a934685f 100644 --- a/pkg/openai/manager.py +++ b/pkg/openai/manager.py @@ -38,8 +38,10 @@ class OpenAIInteract: **config.completion_api_params ) + logging.debug("OpenAI response: %s", response) + self.audit_mgr.report_text_model_usage(config.completion_api_params['model'], - prompt + response['choices'][0]['text']) + response['usage']['total_tokens']) return response