feat: add gpt-4-turbo tokens calc

zhayujie · Nov 7, 2023 · db8e506 · db8e506
1 parent 099f859
commit db8e506
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/bot/chatgpt/chat_gpt_session.py b/bot/chatgpt/chat_gpt_session.py
@@ -1,5 +1,6 @@
 from bot.session_manager import Session
 from common.log import logger
+from common import const
 
 """
     e.g.  [
@@ -64,7 +65,7 @@ def num_tokens_from_messages(messages, model):
     if model in ["gpt-3.5-turbo-0301", "gpt-35-turbo"]:
         return num_tokens_from_messages(messages, model="gpt-3.5-turbo")
     elif model in ["gpt-4-0314", "gpt-4-0613", "gpt-4-32k", "gpt-4-32k-0613", "gpt-3.5-turbo-0613",
-                   "gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613", "gpt-35-turbo-16k"]:
+                   "gpt-3.5-turbo-16k", "gpt-3.5-turbo-16k-0613", "gpt-35-turbo-16k", const.GPT4_PREVIEW, const.GPT4_VISION_PREVIEW]:
         return num_tokens_from_messages(messages, model="gpt-4")
 
     try:

diff --git a/common/const.py b/common/const.py
@@ -8,3 +8,8 @@
 CLAUDEAI = "claude"
 
 MODEL_LIST = ["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4", "wenxin", "wenxin-4", "xunfei", "claude"]
+
+# model
+GPT4 = "gpt-4"
+GPT4_PREVIEW = "gpt-4-1106-preview"
+GPT4_VISION_PREVIEW = "gpt-4-vision-preview"