kaiyuan/chatgpt-on-wechat
1
0
Fork 0
mirror of https://github.com/zhayujie/chatgpt-on-wechat.git synced 2026-07-20 21:57:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: add calc_tokens method on session

Browse Source
This commit is contained in:
lanvent
2023-04-13 20:06:33 +08:00
parent 9cef559a05
commit 2989249e4b
5 changed files with 31 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
bot/chatgpt/chat_gpt_bot.py
View File

@@ -58,7 +58,7 @@ class ChatGPTBot(Bot,OpenAIImage):
# # reply in stream
# return self.reply_text_stream(query, new_query, session_id)
reply_content = self.reply_text(session, session_id, api_key, 0)
reply_content = self.reply_text(session, api_key)
logger.debug("[CHATGPT] new_query={}, session_id={}, reply_cont={}, completion_tokens={}".format(session.messages, session_id, reply_content["content"], reply_content["completion_tokens"]))
if reply_content['completion_tokens'] == 0 and len(reply_content['content']) > 0:
reply = Reply(ReplyType.ERROR, reply_content['content'])
@@ -94,7 +94,7 @@ class ChatGPTBot(Bot,OpenAIImage):
"timeout": conf().get('request_timeout', None), #重试超时时间，在这个时间内，将会自动重试
}
def reply_text(self, session:ChatGPTSession, session_id, api_key, retry_count=0) -> dict:
def reply_text(self, session:ChatGPTSession, api_key=None, retry_count=0) -> dict:
'''
call openai's ChatCompletion to get the answer
:param session: a conversation session
@@ -133,11 +133,11 @@ class ChatGPTBot(Bot,OpenAIImage):
else:
logger.warn("[CHATGPT] Exception: {}".format(e))
need_retry = False
self.sessions.clear_session(session_id)
self.sessions.clear_session(session.session_id)
if need_retry:
logger.warn("[CHATGPT] 第{}次重试".format(retry_count+1))
return self.reply_text(session, session_id, api_key, retry_count+1)
return self.reply_text(session, api_key, retry_count+1)
else:
return result

9
bot/chatgpt/chat_gpt_session.py
View File

@@ -17,7 +17,7 @@ class ChatGPTSession(Session):
def discard_exceeding(self, max_tokens, cur_tokens= None):
precise = True
try:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
except Exception as e:
precise = False
if cur_tokens is None:
@@ -29,7 +29,7 @@ class ChatGPTSession(Session):
elif len(self.messages) == 2 and self.messages[1]["role"] == "assistant":
self.messages.pop(1)
if precise:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
else:
cur_tokens = cur_tokens - max_tokens
break
@@ -40,11 +40,14 @@ class ChatGPTSession(Session):
logger.debug("max_tokens={}, total_tokens={}, len(messages)={}".format(max_tokens, cur_tokens, len(self.messages)))
break
if precise:
cur_tokens = num_tokens_from_messages(self.messages, self.model)
cur_tokens = self.calc_tokens()
else:
cur_tokens = cur_tokens - max_tokens
return cur_tokens
def calc_tokens(self):
return num_tokens_from_messages(self.messages, self.model)
# refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
def num_tokens_from_messages(messages, model):