🐛 Bug: 1. Fix the bug that causes an error when Claude uploads a PNG image.

2. Fix the bug where fields are not automatically added when the database does not have specific fields.

✨ Feature: Add user message ethics review support.

📖 Docs: Update documentation

README_CN.md

@@ -105,10 +105,12 @@ api_keys:
     model:
       - anthropic/claude-3-5-sonnet # 可以使用的模型名称，仅可以使用名为 anthropic 提供商提供的 claude-3-5-sonnet 模型。其他提供商的 claude-3-5-sonnet 模型不可以使用。这种写法不会匹配到other-provider提供的名为anthropic/claude-3-5-sonnet的模型。
       - <anthropic/claude-3-5-sonnet> # 通过在模型名两侧加上尖括号，这样就不会去名为anthropic的渠道下去寻找claude-3-5-sonnet模型，而是将整个 anthropic/claude-3-5-sonnet 作为模型名称。这种写法可以匹配到other-provider提供的名为 anthropic/claude-3-5-sonnet 的模型。但不会匹配到anthropic下面的claude-3-5-sonnet模型。
+      - openai-test/text-moderation-latest # 当开启消息道德审查后，可以使用名为 openai-test 渠道下的 text-moderation-latest 模型进行道德审查。
     preferences:
       USE_ROUND_ROBIN: true # 是否使用轮询负载均衡，true 为使用，false 为不使用，默认为 true。开启轮训后每次请求模型按照 model 配置的顺序依次请求。与 providers 里面原始的渠道顺序无关。因此你可以设置每个 API key 请求顺序不一样。
       AUTO_RETRY: true # 是否自动重试，自动重试下一个提供商，true 为自动重试，false 为不自动重试，默认为 true
       RATE_LIMIT: 2/min # 支持限流，每分钟最多请求次数，可以设置为整数，如 2/min，2 次每分钟、5/hour，5 次每小时、10/day，10 次每天，10/month，10 次每月，10/year，10 次每年。默认60/min，选填
+      ENABLE_MODERATION: true # 是否开启消息道德审查，true 为开启，false 为不开启，默认为 false，当开启后，会对用户的消息进行道德审查，如果发现不当的消息，会返回错误信息。
 
   # 渠道级加权负载均衡配置示例
   - api: sk-KjjI60Yf0JFWtxxxxxxxxxxxxxxwmRWpWpQRo

main.py

@@ -24,10 +24,50 @@ from urllib.parse import urlparse
 import os
 is_debug = bool(os.getenv("DEBUG", False))
 
+from sqlalchemy import inspect, text
+from sqlalchemy.sql import sqltypes
+
 async def create_tables():
     async with engine.begin() as conn:
         await conn.run_sync(Base.metadata.create_all)
 
+        # 检查并添加缺失的列
+        def check_and_add_columns(connection):
+            inspector = inspect(connection)
+            for table in [RequestStat, ChannelStat]:
+                table_name = table.__tablename__
+                existing_columns = {col['name']: col['type'] for col in inspector.get_columns(table_name)}
+
+                for column_name, column in table.__table__.columns.items():
+                    if column_name not in existing_columns:
+                        col_type = _map_sa_type_to_sql_type(column.type)
+                        default = _get_default_sql(column.default)
+                        connection.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {col_type}{default}"))
+
+        await conn.run_sync(check_and_add_columns)
+
+def _map_sa_type_to_sql_type(sa_type):
+    type_map = {
+        sqltypes.Integer: "INTEGER",
+        sqltypes.String: "TEXT",
+        sqltypes.Float: "REAL",
+        sqltypes.Boolean: "BOOLEAN",
+        sqltypes.DateTime: "DATETIME",
+        sqltypes.Text: "TEXT"
+    }
+    return type_map.get(type(sa_type), "TEXT")
+
+def _get_default_sql(default):
+    if default is None:
+        return ""
+    if isinstance(default.arg, bool):
+        return f" DEFAULT {str(default.arg).upper()}"
+    if isinstance(default.arg, (int, float)):
+        return f" DEFAULT {default.arg}"
+    if isinstance(default.arg, str):
+        return f" DEFAULT '{default.arg}'"
+    return ""
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # 启动时的代码
@@ -79,7 +119,7 @@ async def parse_request_body(request: Request):
 
 from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
 from sqlalchemy.orm import declarative_base, sessionmaker
-from sqlalchemy import Column, Integer, String, Float, DateTime, select, Boolean
+from sqlalchemy import Column, Integer, String, Float, DateTime, select, Boolean, Text
 from sqlalchemy.sql import func
 
 # 定义数据库模型
@@ -93,6 +133,8 @@ class RequestStat(Base):
     token = Column(String)
     total_time = Column(Float)
     model = Column(String)
+    is_flagged = Column(Boolean, default=False)
+    moderated_content = Column(Text)
     timestamp = Column(DateTime(timezone=True), server_default=func.now())
 
 class ChannelStat(Base):
@@ -113,6 +155,7 @@ data_dir = os.path.dirname(db_path)
 os.makedirs(data_dir, exist_ok=True)
 
 # 创建异步引擎和会话
+# engine = create_async_engine('sqlite+aiosqlite:///' + db_path, echo=False)
 engine = create_async_engine('sqlite+aiosqlite:///' + db_path, echo=is_debug)
 async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
 
@@ -132,37 +175,76 @@ class StatsMiddleware(BaseHTTPMiddleware):
         start_time = time()
 
         request.state.parsed_body = await parse_request_body(request)
+        endpoint = f"{request.method} {request.url.path}"
+        client_ip = request.client.host
 
         model = "unknown"
+        enable_moderation = False  # 默认不开启道德审查
+        is_flagged = False
+        moderated_content = ""
+
+        config = app.state.config
+        api_list = app.state.api_list
+
+        # 根据token决定是否启用道德审查
+        if token:
+            try:
+                api_index = api_list.index(token)
+                enable_moderation = safe_get(config, 'api_keys', api_index, "preferences", "ENABLE_MODERATION", default=False)
+            except ValueError:
+                # token不在api_list中，使用默认值（不开启）
+                pass
+        else:
+            # 如果token为None，检查全局设置
+            enable_moderation = config.get('ENABLE_MODERATION', False)
+
         if request.state.parsed_body:
             try:
                 request_model = RequestModel(**request.state.parsed_body)
                 model = request_model.model
+                moderated_content = request_model.get_last_text_message()
+
+                if enable_moderation and moderated_content:
+                    moderation_response = await self.moderate_content(moderated_content, token)
+                    moderation_result = moderation_response.body
+                    moderation_data = json.loads(moderation_result)
+                    is_flagged = moderation_data.get('results', [{}])[0].get('flagged', False)
+
+                    if is_flagged:
+                        logger.error(f"Content did not pass the moral check: %s", moderated_content)
+                        process_time = time() - start_time
+                        await self.update_stats(endpoint, process_time, client_ip, model, token, is_flagged, moderated_content)
+                        return JSONResponse(
+                            status_code=400,
+                            content={"error": "Content did not pass the moral check, please modify and try again."}
+                        )
             except RequestValidationError:
                 pass
             except Exception as e:
-                logger.error(f"Error processing request: {str(e)}")
+                if is_debug:
+                    import traceback
+                    traceback.print_exc()
+
+                logger.error(f"处理请求或进行道德检查时出错: {str(e)}")
 
         response = await call_next(request)
         process_time = time() - start_time
 
-        endpoint = f"{request.method} {request.url.path}"
-        client_ip = request.client.host
-
         # 异步更新数据库
-        await self.update_stats(endpoint, process_time, client_ip, model, token)
+        await self.update_stats(endpoint, process_time, client_ip, model, token, is_flagged, moderated_content)
 
         return response
 
-    async def update_stats(self, endpoint, process_time, client_ip, model, token):
+    async def update_stats(self, endpoint, process_time, client_ip, model, token, is_flagged, moderated_content):
         async with self.db as session:
-            # 为每个请求创建一条新的记录
             new_request_stat = RequestStat(
                 endpoint=endpoint,
                 ip=client_ip,
                 token=token,
                 total_time=process_time,
-                model=model
+                model=model,
+                is_flagged=is_flagged,
+                moderated_content=moderated_content
             )
             session.add(new_request_stat)
             await session.commit()
@@ -179,6 +261,14 @@ class StatsMiddleware(BaseHTTPMiddleware):
             session.add(channel_stat)
             await session.commit()
 
+    async def moderate_content(self, content, token):
+        moderation_request = ModerationRequest(input=content)
+
+        # 直接调用 moderations 函数
+        response = await moderations(moderation_request, token)
+
+        return response
+
 # 配置 CORS 中间件
 app.add_middleware(
     CORSMiddleware,
@@ -561,7 +651,7 @@ async def images_generations(
     return await model_handler.request_model(request, token, endpoint="/v1/images/generations")
 
 @app.post("/v1/moderations", dependencies=[Depends(rate_limit_dependency)])
-async def images_generations(
+async def moderations(
     request: ModerationRequest,
     token: str = Depends(verify_api_key)
 ):
@@ -601,9 +691,6 @@ def generate_api_key():
     return JSONResponse(content={"api_key": api_key})
 
 # 在 /stats 路由中返回成功和失败百分比
-from collections import defaultdict
-from sqlalchemy import func
-
 from collections import defaultdict
 from sqlalchemy import func, desc, case
 

models.py

@@ -96,4 +96,15 @@ class RequestModel(BaseModel):
     n: Optional[int] = 1
     user: Optional[str] = None
     tool_choice: Optional[Union[str, ToolChoice]] = None
-    tools: Optional[List[Tool]] = None
+    tools: Optional[List[Tool]] = None
+
+    def get_last_text_message(self) -> Optional[str]:
+        for message in reversed(self.messages):
+            if message.content:
+                if isinstance(message.content, str):
+                    return message.content
+                elif isinstance(message.content, list):
+                    for item in reversed(message.content):
+                        if item.type == "text" and item.text:
+                            return item.text
+        return ""

request.py

@@ -8,9 +8,20 @@ import urllib.parse
 from models import RequestModel
 from utils import c35s, c3s, c3o, c3h, gem, BaseAPI
 
+import imghdr
+
 def encode_image(image_path):
-  with open(image_path, "rb") as image_file:
-    return base64.b64encode(image_file.read()).decode('utf-8')
+    with open(image_path, "rb") as image_file:
+        file_content = image_file.read()
+        file_type = imghdr.what(None, file_content)
+        base64_encoded = base64.b64encode(file_content).decode('utf-8')
+
+        if file_type == 'png':
+            return f"data:image/png;base64,{base64_encoded}"
+        elif file_type in ['jpeg', 'jpg']:
+            return f"data:image/jpeg;base64,{base64_encoded}"
+        else:
+            raise ValueError(f"不支持的图片格式: {file_type}")
 
 async def get_doc_from_url(url):
     filename = urllib.parse.unquote(url.split("/")[-1])
@@ -37,12 +48,28 @@ async def get_encode_image(image_url):
     filename = await get_doc_from_url(image_url)
     image_path = os.getcwd() + "/" + filename
     base64_image = encode_image(image_path)
-    if filename.endswith(".png"):
-        prompt = f"data:image/png;base64,{base64_image}"
-    else:
-        prompt = f"data:image/jpeg;base64,{base64_image}"
     os.remove(image_path)
-    return prompt
+    return base64_image
+
+from PIL import Image
+import io
+def validate_image(image_data, image_type):
+    try:
+        decoded_image = base64.b64decode(image_data)
+        image = Image.open(io.BytesIO(decoded_image))
+
+        # 检查图片格式是否与声明的类型匹配
+        # print("image.format", image.format)
+        if image_type == "image/png" and image.format != "PNG":
+            raise ValueError("Image is not a valid PNG")
+        elif image_type == "image/jpeg" and image.format not in ["JPEG", "JPG"]:
+            raise ValueError("Image is not a valid JPEG")
+
+        # 如果没有异常,则图片有效
+        return True
+    except Exception as e:
+        print(f"Image validation failed: {str(e)}")
+        return False
 
 async def get_image_message(base64_image, engine = None):
     if base64_image.startswith("http"):
@@ -59,6 +86,8 @@ async def get_image_message(base64_image, engine = None):
             }
         }
     if "claude" == engine or "vertex-claude" == engine:
+        # if not validate_image(base64_image.split(",")[1], image_type):
+        #     raise ValueError(f"Invalid image format. Expected {image_type}")
         return {
             "type": "image",
             "source": {

utils.py

@@ -310,10 +310,10 @@ class BaseAPI:
         self.audio_transcriptions: str = urlunparse(parsed_url[:2] + (before_v1 + "/v1/audio/transcriptions",) + ("",) * 3)
         self.moderations: str = urlunparse(parsed_url[:2] + (before_v1 + "/v1/moderations",) + ("",) * 3)
 
-def safe_get(data, *keys):
+def safe_get(data, *keys, default=None):
     for key in keys:
         try:
             data = data[key] if isinstance(data, (dict, list)) else data.get(key)
         except (KeyError, IndexError, AttributeError, TypeError):
-            return None
+            return default
     return data