diff --git a/bot/chatgpt/chat_gpt_bot.py b/bot/chatgpt/chat_gpt_bot.py index d94667e..7e578c7 100644 --- a/bot/chatgpt/chat_gpt_bot.py +++ b/bot/chatgpt/chat_gpt_bot.py @@ -3,10 +3,14 @@ from bot.bot import Bot from config import conf from common.log import logger +from common.expired_dict import ExpiredDict import openai import time -user_session = dict() +if conf().get('expires_in_seconds'): + user_session = ExpiredDict(conf().get('expires_in_seconds')) +else: + user_session = dict() # OpenAI对话模型API (可用) class ChatGPTBot(Bot): diff --git a/common/expired_dict.py b/common/expired_dict.py new file mode 100644 index 0000000..ccf16c7 --- /dev/null +++ b/common/expired_dict.py @@ -0,0 +1,23 @@ +from datetime import datetime, timedelta + +class ExpiredDict(dict): + def __init__(self, expires_in_seconds): + super().__init__() + self.expires_in_seconds = expires_in_seconds + + def __getitem__(self, key): + value, expiry_time = super().__getitem__(key) + if datetime.now() > expiry_time: + del self[key] + raise KeyError("expired {}".format(key)) + self.__setitem__(key, value) + return value + + def __setitem__(self, key, value): + expiry_time = datetime.now() + timedelta(seconds=self.expires_in_seconds) + super().__setitem__(key, (value, expiry_time)) + def get(self, key, default=None): + try: + return self[key] + except KeyError: + return default \ No newline at end of file diff --git a/config-template.json b/config-template.json index 6e99c48..f83c344 100644 --- a/config-template.json +++ b/config-template.json @@ -6,5 +6,6 @@ "group_name_white_list": ["ChatGPT测试群", "ChatGPT测试群2"], "image_create_prefix": ["画", "看", "找"], "conversation_max_tokens": 1000, - "character_desc": "你是ChatGPT, 一个由OpenAI训练的大型语言模型, 你旨在回答并解决人们的任何问题,并且可以使用多种语言与人交流。" + "character_desc": "你是ChatGPT, 一个由OpenAI训练的大型语言模型, 你旨在回答并解决人们的任何问题,并且可以使用多种语言与人交流。", + "expires_in_seconds": 3600 }