✨ Feature: Add support for v1/moderations endpoint.

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
+from models import RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest
 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
@@ -191,7 +191,7 @@ app.add_middleware(
 app.add_middleware(StatsMiddleware)
 
 # 在 process_request 函数中更新成功和失败计数
-async def process_request(request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest], provider: Dict, endpoint=None, token=None):
+async def process_request(request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest], provider: Dict, endpoint=None, token=None):
     url = provider['base_url']
     parsed_url = urlparse(url)
     # print("parsed_url", parsed_url)
@@ -237,6 +237,10 @@ async def process_request(request: Union[RequestModel, ImageGenerationRequest, A
         engine = "whisper"
         request.stream = False
 
+    if endpoint == "/v1/moderations":
+        engine = "moderation"
+        request.stream = False
+
     if provider.get("engine"):
         engine = provider["engine"]
 
@@ -363,7 +367,7 @@ class ModelRequestHandler:
                 print(json.dumps(provider, indent=4, ensure_ascii=False, default=circular_list_encoder))
         return provider_list
 
-    async def request_model(self, request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest], token: str, endpoint=None):
+    async def request_model(self, request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest], token: str, endpoint=None):
         config = app.state.config
         # api_keys_db = app.state.api_keys_db
         api_list = app.state.api_list
@@ -406,7 +410,7 @@ class ModelRequestHandler:
         return await self.try_all_providers(request, matching_providers, use_round_robin, auto_retry, endpoint, token)
 
     # 在 try_all_providers 函数中处理失败的情况
-    async def try_all_providers(self, request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest], providers: List[Dict], use_round_robin: bool, auto_retry: bool, endpoint: str = None, token: str = None):
+    async def try_all_providers(self, request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest], providers: List[Dict], use_round_robin: bool, auto_retry: bool, endpoint: str = None, token: str = None):
         status_code = 500
         error_message = None
         num_providers = len(providers)
@@ -533,7 +537,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], token: str = Depends(verify_api_key)):
+async def request_model(request: Union[RequestModel, ImageGenerationRequest, AudioTranscriptionRequest, ModerationRequest], token: str = Depends(verify_api_key)):
     # logger.info(f"Request received: {request}")
     return await model_handler.request_model(request, token)
 
@@ -556,6 +560,13 @@ 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(
+    request: ModerationRequest,
+    token: str = Depends(verify_api_key)
+):
+    return await model_handler.request_model(request, token, endpoint="/v1/moderations")
+
 from fastapi import UploadFile, File, Form, HTTPException
 import io
 @app.post("/v1/audio/transcriptions", dependencies=[Depends(rate_limit_dependency)])

models.py

@@ -21,6 +21,11 @@ class AudioTranscriptionRequest(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
+class ModerationRequest(BaseModel):
+    input: str
+    model: Optional[str] = "text-moderation-latest"
+    stream: bool = False
+
 class FunctionParameter(BaseModel):
     type: str
     properties: Dict[str, Dict[str, Union[str, Dict[str, str]]]]

request.py

@@ -1043,7 +1043,7 @@ async def get_dalle_payload(request, engine, provider):
 async def get_whisper_payload(request, engine, provider):
     model = provider['model'][request.model]
     headers = {
-        "Content-Type": "application/json",
+        "Content-Type": "multipart/form-data",
     }
     if provider.get("api"):
         headers['Authorization'] = f"Bearer {provider['api'].next()}"
@@ -1066,6 +1066,22 @@ async def get_whisper_payload(request, engine, provider):
 
     return url, headers, payload
 
+async def get_moderation_payload(request, engine, provider):
+    model = provider['model'][request.model]
+    headers = {
+        "Content-Type": "application/json",
+    }
+    if provider.get("api"):
+        headers['Authorization'] = f"Bearer {provider['api'].next()}"
+    url = provider['base_url']
+    url = BaseAPI(url).moderations
+
+    payload = {
+        "input": request.input,
+    }
+
+    return url, headers, payload
+
 async def get_payload(request: RequestModel, engine, provider):
     if engine == "gemini":
         return await get_gemini_payload(request, engine, provider)
@@ -1089,5 +1105,7 @@ async def get_payload(request: RequestModel, engine, provider):
         return await get_dalle_payload(request, engine, provider)
     elif engine == "whisper":
         return await get_whisper_payload(request, engine, provider)
+    elif engine == "moderation":
+        return await get_moderation_payload(request, engine, provider)
     else:
         raise ValueError("Unknown payload")

response.py

@@ -273,7 +273,6 @@ async def fetch_response(client, url, headers, payload):
     response = None
     if payload.get("file"):
         file = payload.pop("file")
-        headers.pop("Content-Type")
         response = await client.post(url, headers=headers, data=payload, files={"file": file})
     else:
         response = await client.post(url, headers=headers, json=payload)

utils.py

@@ -308,6 +308,7 @@ class BaseAPI:
             self.chat_url: str = urlunparse(parsed_url[:2] + (before_v1 + "/v1/chat/completions",) + ("",) * 3)
         self.image_url: str = urlunparse(parsed_url[:2] + (before_v1 + "/v1/images/generations",) + ("",) * 3)
         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):
     for key in keys: