Update app.py

app.py

@@ -24,10 +24,7 @@ free_keys_global = []
 unverified_keys_global = []
 valid_keys_global = []
 
-# 使用 ThreadPoolExecutor 管理并发
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=20)
-
-# 为每个模型维护一个独立的 key 索引
 model_key_indices = {}
 
 def get_credit_summary(api_key):
@@ -220,23 +217,18 @@ def select_key(request_type, model_name):
     if not available_keys:
         return None
 
-    # 获取当前模型的索引，如果没有则初始化为 0
     current_index = model_key_indices.get(model_name, 0)
 
-    # 轮询并重试
     for _ in range(len(available_keys)):
         key = available_keys[current_index % len(available_keys)]
         current_index += 1
         
-        # 检查 KEY 是否有效
         if key_is_valid(key, request_type):
-            # 更新模型索引并返回 KEY
             model_key_indices[model_name] = current_index
             return key
         else:
             logging.warning(f"KEY {key} 无效或达到限制，尝试下一个 KEY")
 
-    # 所有 KEY 都尝试过，重置索引并返回 None
     model_key_indices[model_name] = 0
     return None
 
@@ -245,7 +237,7 @@ def key_is_valid(key, request_type):
     检查 KEY 是否有效，根据不同的请求类型进行不同的检查。
     """
     if request_type == "invalid":
-        return False  # 无效 KEY 始终返回 False
+        return False
     
     credit_summary = get_credit_summary(key)
     if credit_summary is None:
@@ -254,11 +246,11 @@ def key_is_valid(key, request_type):
     total_balance = credit_summary.get("total_balance", 0)
     
     if request_type == "free":
-        return True # 免费 KEY，只要能获取到信息，就认为是有效的
+        return True
     elif request_type == "paid" or request_type == "unverified":
-        return total_balance > 0 # 付费 KEY 或未实名 KEY，需要余额大于 0 才是有效的
+        return total_balance > 0
     else:
-        return False # 未知类型，返回 False
+        return False
 
 def check_authorization(request):
     """
@@ -290,10 +282,6 @@ def index():
 
 @app.route('/check_tokens', methods=['POST'])
 def check_tokens():
-    """
-    处理前端发送的 Token 检测请求。
-    使用线程池并发处理每个 token。
-    """
     tokens = request.json.get('tokens', [])
     test_model = os.environ.get("TEST_MODEL", "Pro/google/gemma-2-9b-it")
 
@@ -322,9 +310,6 @@ def check_tokens():
 
 @app.route('/handsome/v1/chat/completions', methods=['POST'])
 def handsome_chat_completions():
-    """
-    处理 /handsome/v1/chat/completions 路由的请求，添加鉴权，并实现 KEY 的轮询和重试机制。
-    """
     if not check_authorization(request):
         return jsonify({"error": "Unauthorized"}), 401
 
@@ -345,6 +330,7 @@ def handsome_chat_completions():
     }
 
     try:
+        start_time = time.time()
         response = requests.post(
             TEST_MODEL_ENDPOINT,
             headers=headers,
@@ -357,10 +343,74 @@ def handsome_chat_completions():
             return jsonify(response.json()), 429
 
         if data.get("stream", False):
-            return Response(stream_with_context(response.iter_content(chunk_size=1024)), content_type=response.headers['Content-Type'])
+            def generate():
+                first_chunk_time = None
+                full_response_content = ""
+                for chunk in response.iter_content(chunk_size=1024):
+                    if chunk:
+                        if first_chunk_time is None:
+                            first_chunk_time = time.time()
+                        full_response_content += chunk.decode("utf-8")
+                        yield chunk
+                
+                end_time = time.time()
+                first_token_time = first_chunk_time - start_time if first_chunk_time else 0
+                total_time = end_time - start_time
+                
+                # 从完整响应中提取信息
+                try:
+                    response_json = json.loads(full_response_content.split("data: ")[-1].strip())
+                    prompt_tokens = response_json["usage"]["prompt_tokens"]
+                    completion_tokens = response_json["usage"]["completion_tokens"]
+                    response_content = response_json["choices"][0]["message"]["content"]
+                except (KeyError, ValueError, IndexError):
+                    prompt_tokens = 0
+                    completion_tokens = 0
+                    response_content = ""
+                    
+                # 提取用户输入的内容
+                user_content = ""
+                messages = data.get("messages", [])
+                for message in messages:
+                    if message["role"] == "user":
+                        user_content += message["content"] + " "
+                user_content = user_content.strip()
+
+                # 记录日志
+                logging.info(
+                    f"使用的key: {api_key}, 提示token: {prompt_tokens}, 输出token: {completion_tokens}, ���字用时: {first_token_time:.4f}秒, 总共用时: {total_time:.4f}秒, 使用的模型: {model_name}, 用户的内容: {user_content}, 输出的内容: {response_content}"
+                )
+            return Response(stream_with_context(generate()), content_type=response.headers['Content-Type'])
         else:
             response.raise_for_status()
-            return jsonify(response.json())
+            end_time = time.time()
+            response_json = response.json()
+            total_time = end_time - start_time
+
+            # 从响应中提取信息
+            try:
+                prompt_tokens = response_json["usage"]["prompt_tokens"]
+                completion_tokens = response_json["usage"]["completion_tokens"]
+                response_content = response_json["choices"][0]["message"]["content"]
+            except (KeyError, ValueError, IndexError):
+                prompt_tokens = 0
+                completion_tokens = 0
+                response_content = ""
+            
+            # 提取用户输入的内容
+            user_content = ""
+            messages = data.get("messages", [])
+            for message in messages:
+                if message["role"] == "user":
+                    user_content += message["content"] + " "
+            user_content = user_content.strip()
+
+            # 记录日志
+            logging.info(
+                f"使用的key: {api_key}, 提示token: {prompt_tokens}, 输出token: {completion_tokens}, 首字用时: 0, 总共用时: {total_time:.4f}秒, 使用的模型: {model_name}, 用户的内容: {user_content}, 输出的内容: {response_content}"
+            )
+            return jsonify(response_json)
+
     except requests.exceptions.RequestException as e:
         return jsonify({"error": str(e)}), 500
 
@@ -375,10 +425,6 @@ def list_models():
     })
 
 def get_billing_info():
-    """
-    获取所有KEY的额度信息。
-    使用线程池并发处理每个 key。
-    """
     keys = valid_keys_global + unverified_keys_global
     total_balance = 0
 
@@ -451,6 +497,7 @@ def billing_subscription():
     })
 
 if __name__ == '__main__':
+    import json
     logging.info(f"环境变量：{os.environ}")
 
     invalid_keys_global = []