添加优先消耗免费次数

chatgpt-next-web/service.py

@@ -265,8 +265,11 @@ def proxy_chat_completions():
         # 获取 OpenAI API 的响应数据
         result = response.iter_content(chunk_size=8192)
 
-        user_package_key = f'user:{user_id}:package'
-        redis.hincrby(user_package_key, 'basic_chat_limit', -1)
+        if get_free_count(user_id) > 0 and model == 'gpt-3.5-turbo':
+            redis.hincrby(f'user:{user_id}:free', 'basic_chat_count', -1)
+        else:
+            user_package_key = f'user:{user_id}:package'
+            redis.hincrby(user_package_key, 'basic_chat_limit', -1)
 
         # 返回 OpenAI API 的响应给客户端
         return Response(result, content_type=response.headers['content-type'])
@@ -278,8 +281,11 @@ def proxy_chat_completions():
         # 获取 OpenAI API 的响应数据
         result = response.json()
 
-        user_package_key = f'user:{user_id}:package'
-        redis.hincrby(user_package_key, 'basic_chat_limit', -1)
+        if get_free_count(user_id) > 0 and model == 'gpt-3.5-turbo':
+            redis.hincrby(f'user:{user_id}:free', 'basic_chat_count', -1)
+        else:
+            user_package_key = f'user:{user_id}:package'
+            redis.hincrby(user_package_key, 'basic_chat_limit', -1)
 
         # 返回 OpenAI API 的响应给客户端
         return result, response.status_code
@@ -452,17 +458,16 @@ def get_package_expiration(user_id):
 # 检查用户聊天次数是否超过限制
 def exceeded_chat_limit(user_id, package, model):
     if model == 'gpt-3.5-turbo':
-        user_basic_chat_key = f'user:{user_id}:basic_chat'
-        basic_chat_limit = int(package.get(
-            b'basic_chat_limit', 0).decode('utf-8'))
-        if basic_chat_limit >= 0 and int(redis.get(user_basic_chat_key) or 0) >= basic_chat_limit:
+        basic_chat_limit = int(package.get(b'basic_chat_limit', 0))
+        print('basic_chat_limit:', basic_chat_limit)
+        if get_free_count(user_id) > 0:
+            return False
+        if basic_chat_limit <= 0:
             return True
 
     if model == 'gpt-4':
-        user_advanced_chat_key = f'user:{user_id}:advanced_chat'
-        advanced_chat_limit = int(package.get(
-            b'advanced_chat_limit', 0).decode('utf-8'))
-        if advanced_chat_limit >= 0 and int(redis.get(user_advanced_chat_key) or 0) >= advanced_chat_limit:
+        advanced_chat_limit = int(package.get(b'advanced_chat_limit', 0))
+        if advanced_chat_limit <= 0:
             return True
 
     return False
@@ -546,5 +551,11 @@ def pick_up_free_chat_count(user_id):
     return True
 
 
+def get_free_count(user_id):
+    free_data = get_user_free_data(user_id)
+    basic_chat_count = free_data.get(b'basic_chat_count', 0)
+    return int(basic_chat_count)
+
+
 if __name__ == '__main__':
     app.run(debug=True)