Upload app.py

app.py

@@ -56,10 +56,6 @@ safety_settings = [
         "threshold": "BLOCK_NONE"
     },
 ]
-class PromptBlockedException(Exception):
-    def __init__(self, block_reason, safety_ratings=None):
-        self.block_reason = block_reason
-        self.safety_ratings = safety_ratings if safety_ratings else []
 
 class APIKeyManager:
     def __init__(self):
@@ -236,23 +232,27 @@ def handle_api_error(error, attempt):
         switch_api_key()
         return 0, None
 
-    elif isinstance(error, PromptBlockedException):
-        if error.block_reason == "SAFETY":
-            logger.error(f"{current_api_key[:11]} → 用户消息由于安全原因被阻止: {error.safety_ratings}")
-        elif error.block_reason == "OTHER":
-            logger.error(f"{current_api_key[:11]} → 用户消息由于未知原因被阻止")
-        elif error.block_reason == "BLOCKLIST":
-            logger.error(f"{current_api_key[:11]} → 用户消息由于包含黑名单词汇被阻止")
-        elif error.block_reason == "PROHIBITED_CONTENT":
-            logger.error(f"{current_api_key[:11]} → 用户消息由于包含禁止内容被阻止")
+    elif hasattr(error, 'prompt_feedback') and error.prompt_feedback.block_reason:
+        block_reason = error.prompt_feedback.block_reason
+        safety_ratings = error.prompt_feedback.safety_ratings
+        if block_reason == generation_types.BlockReason.SAFETY:
+            logger.warning(f"用户输入被阻挡，原因：SAFETY. 详细信息：{safety_ratings}")
+            return 0, None
+        elif block_reason == generation_types.BlockReason.OTHER:
+            logger.warning(f"用户输入被阻挡，原因：OTHER")
+            return 0, None
+        elif block_reason == generation_types.BlockReason.BLOCKLIST:
+            logger.warning(f"用户输入被阻挡，原因：BLOCKLIST")
+            return 0, None
+        elif block_reason == generation_types.BlockReason.PROHIBITED_CONTENT:
+            logger.warning(f"用户输入被阻挡，原因：PROHIBITED_CONTENT")
+            return 0, None
         else:
-            logger.error(f"{current_api_key[:11]} → 用户消息由于未知原因被阻止: {error.block_reason}")
+            logger.warning(f"用户输入被阻挡，未知阻挡原因：{block_reason}")
+            return 0, None
 
-        switch_api_key()
-        return 0, None
     else:
         logger.error(f"其他错误↙\n {error}")
-        switch_api_key()
         return 2, None
 
 @app.route('/hf/v1/chat/completions', methods=['POST'])