Update main.py

main.py

@@ -4,7 +4,7 @@ import os
 import secrets
 import string
 import time
-from typing import List, Optional, Union # <-- Added Union
+from typing import List, Optional, Union
 
 import httpx
 from dotenv import load_dotenv
@@ -28,7 +28,6 @@ AVAILABLE_MODELS = [
     {"id": "gpt-4o", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "gpt-3.5-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
     {"id": "dall-e-3", "object": "model", "created": int(time.time()), "owned_by": "system"},
-    # Added moderation model to the list
     {"id": "text-moderation-stable", "object": "model", "created": int(time.time()), "owned_by": "system"},
 ]
 
@@ -226,7 +225,7 @@ async def generate_images(request: ImageGenerationRequest):
     return {"created": int(time.time()), "data": results}
 
 
-# === NEW: Moderation Endpoint ===
+# === Moderation Endpoint ===
 
 class ModerationRequest(BaseModel):
     input: Union[str, List[str]]
@@ -236,6 +235,7 @@ class ModerationRequest(BaseModel):
 async def create_moderation(request: ModerationRequest):
     """
     Handles moderation requests, conforming to the OpenAI API specification.
+    Includes a custom 'reason' field in the result if provided by the upstream API.
     """
     input_texts = [request.input] if isinstance(request.input, str) else request.input
     if not input_texts:
@@ -246,9 +246,6 @@ async def create_moderation(request: ModerationRequest):
         'Content-Type': 'application/json',
         'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/138.0.0.0 Safari/537.36',
         'Referer': 'https://www.chatwithmono.xyz/',
-        'sec-ch-ua-platform': '"Windows"',
-        'sec-ch-ua': '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
-        'sec-ch-ua-mobile': '?0',
     }
 
     results = []
@@ -261,28 +258,15 @@ async def create_moderation(request: ModerationRequest):
                 upstream_data = resp.json()
 
                 # --- Transform upstream response to OpenAI format ---
-                # Based on your example, we assume upstream gives: {"overall_sentiment": "...", "categories": {"hate": ...}}
                 upstream_categories = upstream_data.get("categories", {})
-
-                # OpenAI has more specific categories; we'll map the ones we can.
                 openai_categories = {
-                    "hate": upstream_categories.get("hate", False),
-                    "hate/threatening": False, # No data from upstream
-                    "harassment": False, # No data from upstream
-                    "harassment/threatening": False, # No data from upstream
-                    "self-harm": upstream_categories.get("self-harm", False),
-                    "self-harm/intent": False, # No data from upstream
-                    "self-harm/instructions": False, # No data from upstream
-                    "sexual": upstream_categories.get("sexual", False),
-                    "sexual/minors": False, # No data from upstream
-                    "violence": upstream_categories.get("violence", False),
-                    "violence/graphic": False, # No data from upstream
+                    "hate": upstream_categories.get("hate", False), "hate/threatening": False,
+                    "harassment": False, "harassment/threatening": False,
+                    "self-harm": upstream_categories.get("self-harm", False), "self-harm/intent": False, "self-harm/instructions": False,
+                    "sexual": upstream_categories.get("sexual", False), "sexual/minors": False,
+                    "violence": upstream_categories.get("violence", False), "violence/graphic": False,
                 }
-                
-                # Generate scores (1.0 for true, 0.0 for false) as upstream doesn't provide them
                 category_scores = {k: 1.0 if v else 0.0 for k, v in openai_categories.items()}
-                
-                # Determine overall 'flagged' status
                 flagged = upstream_data.get("overall_sentiment") == "flagged"
 
                 result_item = {
@@ -290,6 +274,13 @@ async def create_moderation(request: ModerationRequest):
                     "categories": openai_categories,
                     "category_scores": category_scores,
                 }
+                
+                # --- NEW: Conditionally add the 'reason' field ---
+                # This is a custom extension to the OpenAI spec to provide more detail.
+                reason = upstream_data.get("reason")
+                if reason:
+                    result_item["reason"] = reason
+                
                 results.append(result_item)
 
     except httpx.HTTPStatusError as e:
@@ -313,7 +304,4 @@ async def create_moderation(request: ModerationRequest):
 
 if __name__ == "__main__":
     import uvicorn
-    # To run this file:
-    # 1. Make sure you have a .env file with your SNAP key.
-    # 2. Run in your terminal: uvicorn your_script_name:app --reload --port 8000
     uvicorn.run(app, host="0.0.0.0", port=8000)