🐛 Bug: Fix the bug where the request to generate images cannot be parsed.

.github/workflows/main.yml

@@ -68,7 +68,7 @@ jobs:
         git config --global user.name 'github-actions[bot]'
         git config --global user.email 'github-actions[bot]@users.noreply.github.com'
         git add VERSION
-        git commit -m "Bump version to ${{ steps.bump_version.outputs.new_version }}"
+        git commit -m "📖 Bump version to ${{ steps.bump_version.outputs.new_version }}"
         git push
 
     - name: Build and push Docker image

main.py

@@ -12,7 +12,7 @@ from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.exceptions import RequestValidationError
 
-from models import RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest
+from models import RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest, UnifiedRequest
 from request import get_payload
 from response import fetch_response, fetch_response_stream
 from utils import error_handling_wrapper, post_all_models, load_config, safe_get, circular_list_encoder
@@ -164,30 +164,27 @@ class StatsMiddleware(BaseHTTPMiddleware):
         super().__init__(app)
 
     async def dispatch(self, request: Request, call_next):
-        if request.headers.get("x-api-key"):
-            token = request.headers.get("x-api-key")
-        elif request.headers.get("Authorization"):
-            token = request.headers.get("Authorization").split(" ")[1]
-        else:
-            token = None
-
         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 = ""
+        enable_moderation = False  # 默认不开启道德审查
 
         config = app.state.config
-        api_list = app.state.api_list
-
         # 根据token决定是否启用道德审查
+        if request.headers.get("x-api-key"):
+            token = request.headers.get("x-api-key")
+        elif request.headers.get("Authorization"):
+            token = request.headers.get("Authorization").split(" ")[1]
+        else:
+            token = None
         if token:
             try:
+                api_list = app.state.api_list
                 api_index = api_list.index(token)
                 enable_moderation = safe_get(config, 'api_keys', api_index, "preferences", "ENABLE_MODERATION", default=False)
             except ValueError:
@@ -197,11 +194,16 @@ class StatsMiddleware(BaseHTTPMiddleware):
             # 如果token为None，检查全局设置
             enable_moderation = config.get('ENABLE_MODERATION', False)
 
-        if request.state.parsed_body:
+        parsed_body = await parse_request_body(request)
+        if parsed_body:
             try:
-                request_model = RequestModel(**request.state.parsed_body)
+                request_model = UnifiedRequest.model_validate(parsed_body).data
                 model = request_model.model
-                moderated_content = request_model.get_last_text_message()
+
+                if request_model.request_type == "chat":
+                    moderated_content = request_model.get_last_text_message()
+                elif request_model.request_type == "image":
+                    moderated_content = request_model.prompt
 
                 if enable_moderation and moderated_content:
                     moderation_response = await self.moderate_content(moderated_content, token)
@@ -636,7 +638,7 @@ def verify_admin_api_key(credentials: HTTPAuthorizationCredentials = Depends(sec
     return token
 
 @app.post("/v1/chat/completions", dependencies=[Depends(rate_limit_dependency)])
-async def request_model(request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest], token: str = Depends(verify_api_key)):
+async def request_model(request: RequestModel, token: str = Depends(verify_api_key)):
     # logger.info(f"Request received: {request}")
     return await model_handler.request_model(request, token)
 

models.py

@@ -1,30 +1,6 @@
 from io import IOBase
-from pydantic import BaseModel, Field
-from typing import List, Dict, Optional, Union, Tuple
-
-class ImageGenerationRequest(BaseModel):
-    model: str
-    prompt: str
-    n: int
-    size: str
-    stream: bool = False
-
-class AudioTranscriptionRequest(BaseModel):
-    file: Tuple[str, IOBase, str]
-    model: str
-    language: Optional[str] = None
-    prompt: Optional[str] = None
-    response_format: Optional[str] = None
-    temperature: Optional[float] = None
-    stream: bool = False
-
-    class Config:
-        arbitrary_types_allowed = True
-
-class ModerationRequest(BaseModel):
-    input: str
-    model: Optional[str] = "text-moderation-latest"
-    stream: bool = False
+from pydantic import BaseModel, Field, model_validator
+from typing import List, Dict, Optional, Union, Tuple, Literal
 
 class FunctionParameter(BaseModel):
     type: str
@@ -82,6 +58,7 @@ class ToolChoice(BaseModel):
     function: Optional[FunctionChoice] = None
 
 class RequestModel(BaseModel):
+    request_type: Literal["chat"] = "chat"
     model: str
     messages: List[Message]
     logprobs: Optional[bool] = None
@@ -107,4 +84,54 @@ class RequestModel(BaseModel):
                     for item in reversed(message.content):
                         if item.type == "text" and item.text:
                             return item.text
-        return ""
+        return ""
+
+class ImageGenerationRequest(BaseModel):
+    request_type: Literal["image"] = "image"
+    prompt: str
+    model: Optional[str] = "dall-e-3"
+    n:  Optional[int] = 1
+    size: Optional[str] = "1024x1024"
+    stream: bool = False
+
+class AudioTranscriptionRequest(BaseModel):
+    request_type: Literal["audio"] = "audio"
+    file: Tuple[str, IOBase, str]
+    model: str
+    language: Optional[str] = None
+    prompt: Optional[str] = None
+    response_format: Optional[str] = None
+    temperature: Optional[float] = None
+    stream: bool = False
+
+    class Config:
+        arbitrary_types_allowed = True
+
+class ModerationRequest(BaseModel):
+    request_type: Literal["moderation"] = "moderation"
+    input: str
+    model: Optional[str] = "text-moderation-latest"
+    stream: bool = False
+
+class UnifiedRequest(BaseModel):
+    data: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest] = Field(..., discriminator="request_type")
+
+    @model_validator(mode='before')
+    @classmethod
+    def set_request_type(cls, values):
+        if isinstance(values, dict):
+            if "messages" in values:
+                values["request_type"] = "chat"
+                values["data"] = RequestModel(**values)
+            elif "prompt" in values:
+                values["request_type"] = "image"
+                values["data"] = ImageGenerationRequest(**values)
+            elif "file" in values:
+                values["request_type"] = "audio"
+                values["data"] = AudioTranscriptionRequest(**values)
+            elif "input" in values:
+                values["request_type"] = "moderation"
+                values["data"] = ModerationRequest(**values)
+            else:
+                raise ValueError("无法确定请求类型")
+        return values