Update main.py

main.py

@@ -1,10 +1,10 @@
 import base64
 import json
 import os
-import secrets  # <-- Import secrets
-import string   # <-- Import string
+import secrets
+import string
 import time
-from typing import List, Optional
+from typing import List, Optional, Union # <-- Added Union
 
 import httpx
 from dotenv import load_dotenv
@@ -28,6 +28,8 @@ 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"},
 ]
 
 MODEL_ALIASES = {}
@@ -45,9 +47,6 @@ app = FastAPI(
 def generate_random_id(prefix: str, length: int = 29) -> str:
     """
     Generates a cryptographically secure, random alphanumeric ID.
-    The default length of 29 characters is common for OpenAI IDs.
-    The example 'bwvaLjbI0KEKMadGmFbSsjYNLgaI' is 30 characters.
-    You can adjust the length as needed.
     """
     population = string.ascii_letters + string.digits
     random_part = "".join(secrets.choice(population) for _ in range(length))
@@ -77,13 +76,9 @@ class ChatRequest(BaseModel):
 async def chat_completion(request: ChatRequest):
     """
     Handles chat completion requests, supporting both streaming and non-streaming responses.
-    This endpoint now uses a long, random ID for completions.
     """
     model_id = MODEL_ALIASES.get(request.model, request.model)
-    
-    # Generate the ID once for the entire request
     chat_id = generate_random_id("chatcmpl-")
-
     headers = {
         'accept': 'text/event-stream',
         'content-type': 'application/json',
@@ -91,7 +86,6 @@ async def chat_completion(request: ChatRequest):
         'referer': 'https://www.chatwithmono.xyz/',
         'user-agent': 'Mozilla/5.0',
     }
-
     payload = {
         "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
         "model": model_id
@@ -100,7 +94,6 @@ async def chat_completion(request: ChatRequest):
     if request.stream:
         async def event_stream():
             created = int(time.time())
-            
             is_first_chunk = True
             usage_info = None
 
@@ -108,58 +101,38 @@ async def chat_completion(request: ChatRequest):
                 async with httpx.AsyncClient(timeout=120) as client:
                     async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
                         response.raise_for_status()
-                        
                         async for line in response.aiter_lines():
-                            if not line:
-                                continue
-
+                            if not line: continue
                             if line.startswith("0:"):
                                 try:
                                     content_piece = json.loads(line[2:])
-                                    delta = {
-                                        "content": content_piece,
-                                        "function_call": None,
-                                        "tool_calls": None
-                                    }
+                                    delta = {"content": content_piece, "function_call": None, "tool_calls": None}
                                     if is_first_chunk:
                                         delta["role"] = "assistant"
                                         is_first_chunk = False
-
                                     chunk_data = {
-                                        "id": chat_id, # Use the pre-generated ID
-                                        "object": "chat.completion.chunk",
-                                        "created": created,
+                                        "id": chat_id, "object": "chat.completion.chunk", "created": created,
                                         "model": model_id,
                                         "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
                                         "usage": None
                                     }
                                     yield f"data: {json.dumps(chunk_data)}\n\n"
-                                except json.JSONDecodeError:
-                                    continue
-                            
+                                except json.JSONDecodeError: continue
                             elif line.startswith(("e:", "d:")):
                                 try:
-                                    end_data = json.loads(line[2:])
-                                    usage_info = end_data.get("usage")
-                                except (json.JSONDecodeError, AttributeError):
-                                    pass
+                                    usage_info = json.loads(line[2:]).get("usage")
+                                except (json.JSONDecodeError, AttributeError): pass
                                 break
-                
                 final_usage = None
                 if usage_info:
                     prompt_tokens = usage_info.get("promptTokens", 0)
                     completion_tokens = usage_info.get("completionTokens", 0)
                     final_usage = {
-                        "prompt_tokens": prompt_tokens,
-                        "completion_tokens": completion_tokens,
+                        "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
                         "total_tokens": prompt_tokens + completion_tokens,
                     }
-
                 done_chunk = {
-                    "id": chat_id, # Use the pre-generated ID
-                    "object": "chat.completion.chunk",
-                    "created": created,
-                    "model": model_id,
+                    "id": chat_id, "object": "chat.completion.chunk", "created": created, "model": model_id,
                     "choices": [{
                         "index": 0,
                         "delta": {"role": "assistant", "content": None, "function_call": None, "tool_calls": None},
@@ -168,54 +141,33 @@ async def chat_completion(request: ChatRequest):
                     "usage": final_usage
                 }
                 yield f"data: {json.dumps(done_chunk)}\n\n"
-
             except httpx.HTTPStatusError as e:
                 error_content = {
                     "error": {
                         "message": f"Upstream API error: {e.response.status_code}. Details: {e.response.text}",
-                        "type": "upstream_error",
-                        "code": str(e.response.status_code)
+                        "type": "upstream_error", "code": str(e.response.status_code)
                     }
                 }
                 yield f"data: {json.dumps(error_content)}\n\n"
-            
             finally:
                 yield "data: [DONE]\n\n"
-
         return StreamingResponse(event_stream(), media_type="text/event-stream")
-    else:
-        # Non-streaming logic
-        assistant_response = ""
-        usage_info = {}
-
+    else: # Non-streaming
+        assistant_response, usage_info = "", {}
         try:
             async with httpx.AsyncClient(timeout=120) as client:
                 async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
                     response.raise_for_status()
                     async for chunk in response.aiter_lines():
                         if chunk.startswith("0:"):
-                            try:
-                                piece = json.loads(chunk[2:])
-                                assistant_response += piece
-                            except:
-                                continue
+                            try: assistant_response += json.loads(chunk[2:])
+                            except: continue
                         elif chunk.startswith(("e:", "d:")):
-                            try:
-                                data = json.loads(chunk[2:])
-                                usage_info = data.get("usage", {})
-                            except:
-                                continue
-            
+                            try: usage_info = json.loads(chunk[2:]).get("usage", {})
+                            except: continue
             return JSONResponse(content={
-                "id": chat_id, # Use the pre-generated ID
-                "object": "chat.completion",
-                "created": int(time.time()),
-                "model": model_id,
-                "choices": [{
-                    "index": 0,
-                    "message": {"role": "assistant", "content": assistant_response},
-                    "finish_reason": "stop"
-                }],
+                "id": chat_id, "object": "chat.completion", "created": int(time.time()), "model": model_id,
+                "choices": [{"index": 0, "message": {"role": "assistant", "content": assistant_response}, "finish_reason": "stop"}],
                 "usage": {
                     "prompt_tokens": usage_info.get("promptTokens", 0),
                     "completion_tokens": usage_info.get("completionTokens", 0),
@@ -223,15 +175,7 @@ async def chat_completion(request: ChatRequest):
                 }
             })
         except httpx.HTTPStatusError as e:
-            return JSONResponse(
-                status_code=e.response.status_code,
-                content={
-                    "error": {
-                        "message": f"Upstream API error. Details: {e.response.text}",
-                        "type": "upstream_error"
-                    }
-                }
-            )
+            return JSONResponse(status_code=e.response.status_code, content={"error": {"message": f"Upstream API error. Details: {e.response.text}", "type": "upstream_error"}})
 
 
 # === Image Generation ===
@@ -247,26 +191,18 @@ class ImageGenerationRequest(BaseModel):
 async def generate_images(request: ImageGenerationRequest):
     """Handles image generation requests."""
     results = []
-
     try:
         async with httpx.AsyncClient(timeout=120) as client:
             for _ in range(request.n):
                 model = request.model or "default"
-
                 if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
-                    headers = {
-                        'Content-Type': 'application/json',
-                        'User-Agent': 'Mozilla/5.0',
-                        'Referer': 'https://www.chatwithmono.xyz/',
-                    }
+                    headers = {'Content-Type': 'application/json', 'User-Agent': 'Mozilla/5.0', 'Referer': 'https://www.chatwithmono.xyz/'}
                     payload = {"prompt": request.prompt, "model": model}
                     resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
                     resp.raise_for_status()
                     data = resp.json()
                     b64_image = data.get("image")
-                    if not b64_image:
-                        return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
-
+                    if not b64_image: return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
                     if SNAPZION_API_KEY:
                         upload_headers = {"Authorization": SNAPZION_API_KEY}
                         upload_files = {'file': ('image.png', base64.b64decode(b64_image), 'image/png')}
@@ -276,7 +212,6 @@ async def generate_images(request: ImageGenerationRequest):
                         image_url = upload_data.get("url")
                     else:
                         image_url = f"data:image/png;base64,{b64_image}"
-
                     results.append({"url": image_url, "b64_json": b64_image, "revised_prompt": data.get("revised_prompt")})
                 else:
                     params = {"prompt": request.prompt, "aspect_ratio": request.aspect_ratio, "link": "typegpt.net"}
@@ -284,17 +219,97 @@ async def generate_images(request: ImageGenerationRequest):
                     resp.raise_for_status()
                     data = resp.json()
                     results.append({"url": data.get("image_link"), "b64_json": data.get("base64_output")})
+    except httpx.HTTPStatusError as e:
+        return JSONResponse(status_code=502, content={"error": f"Image generation failed. Upstream error: {e.response.status_code}", "details": e.response.text})
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
+    return {"created": int(time.time()), "data": results}
+
+
+# === NEW: Moderation Endpoint ===
+
+class ModerationRequest(BaseModel):
+    input: Union[str, List[str]]
+    model: Optional[str] = "text-moderation-stable"
+
+@app.post("/v1/moderations")
+async def create_moderation(request: ModerationRequest):
+    """
+    Handles moderation requests, conforming to the OpenAI API specification.
+    """
+    input_texts = [request.input] if isinstance(request.input, str) else request.input
+    if not input_texts:
+         return JSONResponse(status_code=400, content={"error": {"message": "Request must have at least one input string.", "type": "invalid_request_error"}})
+
+    moderation_url = "https://www.chatwithmono.xyz/api/moderation"
+    headers = {
+        '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 = []
+    try:
+        async with httpx.AsyncClient(timeout=30) as client:
+            for text_input in input_texts:
+                payload = {"text": text_input}
+                resp = await client.post(moderation_url, headers=headers, json=payload)
+                resp.raise_for_status()
+                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
+                }
+                
+                # 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 = {
+                    "flagged": flagged,
+                    "categories": openai_categories,
+                    "category_scores": category_scores,
+                }
+                results.append(result_item)
+
     except httpx.HTTPStatusError as e:
         return JSONResponse(
-            status_code=502,
-            content={"error": f"Image generation failed. Upstream error: {e.response.status_code}", "details": e.response.text}
+            status_code=502, # Bad Gateway
+            content={"error": {"message": f"Moderation failed. Upstream error: {e.response.status_code}", "type": "upstream_error", "details": e.response.text}}
         )
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
+        return JSONResponse(status_code=500, content={"error": {"message": "An internal error occurred during moderation.", "type": "internal_error", "details": str(e)}})
 
+    # Build the final OpenAI-compatible response
+    final_response = {
+        "id": generate_random_id("modr-"),
+        "model": request.model,
+        "results": results,
+    }
+    return JSONResponse(content=final_response)
 
-    return {"created": int(time.time()), "data": results}
 
+# --- Main Execution ---
 
 if __name__ == "__main__":
     import uvicorn