Преглед на файлове

添加会话过期时间

lanvent преди 3 години
родител
ревизия
2c4d79e952
променени са 3 файла, в които са добавени 30 реда и са изтрити 2 реда
  1. 5 1
      bot/chatgpt/chat_gpt_bot.py
  2. 23 0
      common/expired_dict.py
  3. 2 1
      config-template.json

+ 5 - 1
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):

+ 23 - 0
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

+ 2 - 1
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
 }