'fix:修改获取token逻辑'

novel_api/apps/chatbot/models.py

@@ -1,4 +1,5 @@
 import mongoengine
+import datetime
 
 
 class Memory(mongoengine.Document):
@@ -8,7 +9,16 @@ class Memory(mongoengine.Document):
 
 class Access_token_pool(mongoengine.Document):
     access_token = mongoengine.StringField()
-    now_time = mongoengine.StringField()
+    now_time = mongoengine.DateTimeField()
+
+    @classmethod
+    def get_oldest_token(cls):
+        current_time = datetime.datetime.now()
+        oldest_token = cls.objects.order_by('now_time').first()  # 未取到不会报错
+        if oldest_token:
+            oldest_token.now_time = current_time
+            oldest_token.save()
+            return oldest_token  # 如果不存在会为None
 
 
 class User(mongoengine.Document):
@@ -22,12 +32,8 @@ class Paragraph(mongoengine.Document):
 class Choice(mongoengine.Document):
     text = mongoengine.ListField()
 
-
-
-
-
 # "novel_id": 1,
 # "novel_title": "我对总裁大人有偏见",
 # "novel_image": "https://cn.bing.com/images/search?view=detailV2&ccid=Bq5jD730&id=E91F2971887D91E927CE16AB8F45BEB8C1185543&thid=OIP.Bq5jD730RoSsMF3c1yWIWwHaJ4&mediaurl=https%3A%2F%2Fstatic.zongheng.com%2Fupload%2Fcover%2Fshucheng%2F16%2F15416195.jpg&exph=3200&expw=2400&q=%e5%b0%8f%e8%af%b4%e5%9b%be%e7%89%87&simid=608000767812436123&form=IRPRST&ck=ADD7A0A91334D779ECBD217BF7F86F67&selectedindex=2&ajaxhist=0&ajaxserp=0&vt=0&sim=11",
 # "novel_tag": ["言情"],
-# "novel_visit": '10.0万',
+# "novel_visit": '10.0万',

novel_api/apps/chatbot/tests.py

@@ -0,0 +1,54 @@
+import datetime
+import threading
+from rest_framework.exceptions import APIException
+from mongoengine import Document, StringField, DateTimeField, connect
+
+# 连接 MongoDB 数据库
+connect('your_database_name')
+
+class Access_token_pool(Document):
+    access_token = StringField()
+    now_time = DateTimeField()
+
+    @classmethod
+    def create_tokens(cls, num_tokens):
+        current_time = datetime.datetime.now()
+        for i in range(num_tokens):
+            token = cls(access_token=str(i), now_time=current_time)
+            token.save()
+
+    @classmethod
+    def get_oldest_token(cls):
+        current_time = datetime.datetime.now()
+        oldest_token = cls.objects.order_by('now_time').first()  # 未取到不会报错
+        print(oldest_token)
+        if oldest_token:
+            oldest_token.now_time = current_time
+            oldest_token.save()
+        return oldest_token  # 如果不存在会为None
+
+# token_lock = threading.Lock()  # 创建互斥锁
+
+def concurrent_get_token():
+    try:
+        # with token_lock:  # 使用互斥锁
+            oldest_token = Access_token_pool.get_oldest_token()
+            print(oldest_token.access_token)
+            return oldest_token
+    except Exception as e:
+        APIException("获取 token 出错:", str(e))
+
+
+# 创建测试数据
+# Access_token_pool.create_tokens(50)
+
+# 模拟并发获取 token
+threads = []
+for _ in range(100):
+    thread = threading.Thread(target=concurrent_get_token)
+    threads.append(thread)
+    thread.start()
+
+# 等待所有线程完成
+for thread in threads:
+    thread.join()

novel_api/apps/chatbot/views.py

@@ -1,6 +1,3 @@
-import uuid
-import time
-
 from rest_framework.exceptions import APIException
 from rest_framework.viewsets import ViewSet
 from rest_framework.decorators import action
@@ -9,43 +6,33 @@
 from utils.common_response import APIResponse
 from .models import Access_token_pool, Paragraph, Choice
 
+# token_lock = threading.Lock()  # 注释掉这两行代码，就变成了不加锁的版本
 
-class User(object):
-
-    def __init__(self):
-        self.user_id = str(uuid.uuid4())
+def concurrent_get_token():
+    try:
+        # with token_lock:  # 注释掉这两行代码，就变成了不加锁的版本
+            oldest_token = Access_token_pool.get_oldest_token()
+            return oldest_token.access_token
+    except Exception as e:
+        raise APIException(f"获取 token 出错:{e}")
 
 
 def get_response_streaming(prompt):
-    """
-    Args:
-        prompt:提示词
-        access_token
-
-    Returns:流式
-    """
-    acp_obj = Access_token_pool.objects[0]
-    access_token = acp_obj.access_token
+    access_token = concurrent_get_token()
     chatbot = Chatbot(config={
         "access_token": access_token,
         "collect_analytics": True,
-        # 服务器挂代理
         "proxy": "socks5h://127.0.0.1:1090"
     })
     try:
         result = chatbot.ask(prompt)
     except Exception as e:
-        acp_obj.delete()
-        raise APIException('chatgpt报错')
-    acp_obj.delete()
-    Access_token_pool(access_token=access_token, now_time=str(time.time())).save()
-    # 存回来
+        raise APIException("chatgpt报错:", str(e))
     return result
 
 
 def get_response(prompt):  # 这个是仅仅只有总结接口使用 不会返回流式输出
-    acp_obj = Access_token_pool.objects[0]
-    access_token = acp_obj.access_token
+    access_token = concurrent_get_token()
     chatbot = Chatbot(config={
         "access_token": access_token,
         "collect_analytics": True,
@@ -57,10 +44,7 @@ def get_response(prompt):  # 这个是仅仅只有总结接口使用 不会返
             message = data["message"][len(prev_text):]
             prev_text = data["message"]
     except Exception as e:
-        acp_obj.delete()
-        raise APIException('chatgpt报错')
-    acp_obj.delete()
-    Access_token_pool(access_token=access_token, now_time=str(time.time())).save()
+        raise APIException("chatgpt报错:", str(e))
     return prev_text
 
 
@@ -79,9 +63,7 @@ def get_content(self, request, *args, **kwargs):
             'long_memory': '',
             "index": request.data.get('index', '')
         }
-
         del pre_data["index"]
-
         prompt = lambda background, relationship, character, summary, content, question, choice, long_memory: f"""
         现在你是一个经验丰富的写对话小说的网文作家，你需要续写这本小说，小说的大部分内容都是对话，注意你续写的只是小说的开头部分，发展要缓慢，续写应该停在突然的地方，比如话说到一半，人物动作做到一半
         续写指的是接着小说的末尾创作出新的内容，创作出的新内容与小说之前的内容不矛盾，输出时不用将小说之前内容输出！！！

novel_api/settings/dev.py

@@ -89,10 +89,14 @@
 }
 
 # mongodb配置
-from mongoengine import connect
-from urllib.parse import quote_plus
+# from mongoengine import connect
+# from urllib.parse import quote_plus
+#
+# connect('novel_api_h5', host="mongodb://%s:%s@%s" % (quote_plus("aiwaves"), quote_plus("bxzn2023"), "47.96.122.196"))
 
-connect('novel_api_h5', host="mongodb://%s:%s@%s" % (quote_plus("aiwaves"), quote_plus("bxzn2023"), "47.96.122.196"))
+from mongoengine import connect
+# 连接 MongoDB 数据库
+connect('your_database_name')
 
 # Password validation
 # https://docs.djangoproject.com/en/4.2/ref/settings/#auth-password-validators

script/add_lock.py

@@ -0,0 +1,63 @@
+import datetime
+import threading
+from rest_framework.exceptions import APIException
+from mongoengine import Document, StringField, DateTimeField, connect
+
+# 连接本地 MongoDB 数据库
+connect('your_database_name')
+
+# # 连接远程 MongoDB 数据库
+# from mongoengine import connect
+# from urllib.parse import quote_plus
+#
+# connect('novel_api_h5', host="mongodb://%s:%s@%s" % (quote_plus("aiwaves"), quote_plus("bxzn2023"), "47.96.122.196"))
+
+
+class Access_token_pool(Document):
+    access_token = StringField()
+    now_time = DateTimeField()
+
+    # @classmethod
+    # def create_tokens(cls, num_tokens):
+    #     current_time = datetime.datetime.now()
+    #     for i in range(num_tokens):
+    #         token = cls(access_token=str(i), now_time=current_time)
+    #         token.save()
+
+    @classmethod
+    def get_oldest_token(cls):
+        current_time = datetime.datetime.now()
+        oldest_token = cls.objects.order_by('now_time').first()  # 未取到不会报错
+        print(oldest_token.access_token)
+        if oldest_token:
+            oldest_token.now_time = current_time
+            oldest_token.save()
+        return oldest_token  # 如果不存在会为None
+
+
+
+# token_lock = threading.Lock()  # 注释掉这两行代码，就变成了不加锁的版本
+
+def concurrent_get_token():
+    try:
+        # with token_lock:  # 注释掉这两行代码，就变成了不加锁的版本
+            oldest_token = Access_token_pool.get_oldest_token()
+            print(oldest_token.access_token)
+            return oldest_token
+    except Exception as e:
+        APIException("获取 token 出错:", str(e))
+
+
+# 创建测试数据
+# Access_token_pool.create_tokens(50)
+
+# 模拟并发获取 token
+threads = []
+for _ in range(100):
+    thread = threading.Thread(target=concurrent_get_token)
+    threads.append(thread)
+    thread.start()
+
+# 等待所有线程完成
+for thread in threads:
+    thread.join()

script/add_token.py

@@ -0,0 +1,28 @@
+from mongoengine import Document, StringField, DateTimeField, connect
+import datetime
+
+# 当前时间
+current_time = datetime.datetime.now()
+
+# 连接本地 MongoDB 数据库
+connect('your_database_name')
+
+
+# # 连接远程 MongoDB 数据库
+# from mongoengine import connect
+# from urllib.parse import quote_plus
+#
+# connect('novel_api_h5', host="mongodb://%s:%s@%s" % (quote_plus("aiwaves"), quote_plus("bxzn2023"), "47.96.122.196"))
+
+class Access_token_pool(Document):
+    access_token = StringField()
+    now_time = DateTimeField()
+
+
+access_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6Ik1UaEVOVUpHTkVNMVFURTRNMEZCTWpkQ05UZzVNRFUxUlRVd1FVSkRNRU13UmtGRVFrRXpSZyJ9.eyJodHRwczovL2FwaS5vcGVuYWkuY29tL3Byb2ZpbGUiOnsiZW1haWwiOiJqaWF0a2IwNjU2QHdlcC5lbWFpbCIsImVtYWlsX3ZlcmlmaWVkIjp0cnVlfSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7InVzZXJfaWQiOiJ1c2VyLWdGV3RRS2QxYkJuWmNsbWgzeXdrNllzSCJ9LCJpc3MiOiJodHRwczovL2F1dGgwLm9wZW5haS5jb20vIiwic3ViIjoiYXV0aDB8NjQ0YTFiMmUxZGVjOGQyZWRhOGQxZmU0IiwiYXVkIjpbImh0dHBzOi8vYXBpLm9wZW5haS5jb20vdjEiLCJodHRwczovL29wZW5haS5vcGVuYWkuYXV0aDBhcHAuY29tL3VzZXJpbmZvIl0sImlhdCI6MTY4NjExMDMzNywiZXhwIjoxNjg3MzE5OTM3LCJhenAiOiJUZEpJY2JlMTZXb1RIdE45NW55eXdoNUU0eU9vNkl0RyIsInNjb3BlIjoib3BlbmlkIHByb2ZpbGUgZW1haWwgbW9kZWwucmVhZCBtb2RlbC5yZXF1ZXN0IG9yZ2FuaXphdGlvbi5yZWFkIG9yZ2FuaXphdGlvbi53cml0ZSJ9.0fI-3RYbF1eg_kRzE3Q3QYNGa6kgOgsaPvt4m9HREjRkYFKWGVP3lBWxGQkG3553imEr2vZjETn4hyHFGj-LLYg1XpLpa-_zvhjITyzzT4flUZijFp6tzZg185WIySeXrmoKmZTxL7ij1LF1Un6Y2IRPSIbvOvugr5wWvrPZtZD0xHcgmzmj-A1xJytjM7zhdlKtaJiK_VyeR6wlTv5fLdjPWxU5nutmxSnaV6i9drR9g8HCMk2_mtckyS_zAUJIW2KQVZ81d0OFijL2aiaQiaSOtNiV2jKRr0jrm9lysUkvKDwszQMGtOZXhPR2KMWEr_nLHZpryBJdYPf3yF8E7A"
+# Access_token_pool(access_token=access_token,now_time=current_time).save()
+
+
+token_obj = Access_token_pool.objects.first()
+token_obj.access_token = access_token
+token_obj.save()