diff --git a/bot/chatgpt/chat_gpt_session.py b/bot/chatgpt/chat_gpt_session.py index 794dac3..7b8ea6e 100644 --- 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 index fa90c32..92fbfca 100644 --- a/common/const.py +++ b/common/const.py @@ -8,3 +8,8 @@ LINKAI = "linkai" 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"