From b7ae1fa516a7eb2caa1835e53fabf2572b98aeff Mon Sep 17 00:00:00 2001
From: Rock Chin <1010553892@qq.com>
Date: Sun, 5 Feb 2023 15:12:54 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E5=A4=8D=E9=94=99=E8=AF=AF?=
 =?UTF-8?q?=E7=9A=84=E4=BD=BF=E7=94=A8=E9=87=8F=E8=AE=A1=E7=AE=97=E6=96=B9?=
 =?UTF-8?q?=E5=BC=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/audit/gatherer.py | 4 ++--
 pkg/openai/manager.py | 4 +++-
 2 files changed, 5 insertions(+), 3 deletions(-)

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