count token used.

yihong0618 · Mar 8, 2023 · f74ae37 · f74ae37
1 parent f7b3daa
commit f74ae37
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/book_maker/translator/chatgptapi_translator.py b/book_maker/translator/chatgptapi_translator.py
@@ -28,6 +28,7 @@ def translate(self, text):
             }
         ]
         count_tokens = num_tokens_from_messages(message_log)
+        consumed_tokens = 0
         t_text = ""
         if count_tokens > 4000:
             print("too long!")
@@ -60,6 +61,7 @@ def translate(self, text):
                     .decode()
                 )
                 print(t_sub_text)
+                consumed_tokens += completion["usage"]["prompt_tokens"]
 
                 t_text = t_text + t_sub_text
 
@@ -82,6 +84,8 @@ def translate(self, text):
                         .encode("utf8")
                         .decode()
                     )
+                    consumed_tokens += completion["usage"]["prompt_tokens"]
+
                 except Exception as e:
                     # TIME LIMIT for open api please pay
                     key_len = self.key.count(",") + 1
@@ -105,5 +109,8 @@ def translate(self, text):
                         .encode("utf8")
                         .decode()
                     )
+                    consumed_tokens += completion["usage"]["prompt_tokens"]
+
         print(t_text)
+        print(f"{consumed_tokens} prompt tokens used.")
         return t_text