Update app/main.py

app/main.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
 from glin_profanity import Filter
+import re
 
 toxicity_pipe = pipeline("text-classification", model="textdetox/twitter-xlmr-toxicity-classifier")
 detox_pipe = pipeline("text2text-generation", model="s-nlp/mt0-xl-detox-orpo")
@@ -18,25 +19,38 @@ class ModerationResponse(BaseModel):
     is_flagged: bool
     rewritten_text: list[str] | None
 
+# 1. Define explicit threat keywords (You can expand this list)
+THREAT_KEYWORDS = [
+    r"\bguns?\b", r"\bweapons?\b", r"\bbombs?\b", r"\bexplosives?\b", 
+    r"\bkill\b", r"\bmurder\b", r"\bshoot\b", r"\bstab\b", 
+    r"\bdie\b", r"\bdeath\b", r"\bhurt\b"
+]
+# Compile the regex pattern for fast matching
+threat_pattern = re.compile("|".join(THREAT_KEYWORDS), re.IGNORECASE)
+
 @app.post("/moderate", response_model=ModerationResponse)
 async def moderate_text(request: TextRequest):
     text = request.text
     
+    # Check 1: Profanity filter
     profanity_result = profanity_filter.check_profanity(text)
     
-    # 1. Get the full result from the toxicity model
+    # Check 2: AI Toxicity model
     tox_result = toxicity_pipe(text)[0]
-    
-    # 2. Check if the LABEL is actually toxic AND the confidence is high
-    # (The model usually outputs 'toxic' or 'neutral' / 'LABEL_1' or 'LABEL_0')
     is_toxic_by_model = (tox_result['label'].lower() in ['toxic', 'label_1']) and (tox_result['score'] > 0.5)
     
-    # 3. Final Flag Decision
-    is_flagged = profanity_result["contains_profanity"] or is_toxic_by_model
+    # Check 3: Threat Keyword Regex
+    is_threat_by_keyword = bool(threat_pattern.search(text))
+    
+    # Final Decision: Flag if ANY check is true
+    is_flagged = (
+        profanity_result["contains_profanity"] or 
+        is_toxic_by_model or 
+        is_threat_by_keyword
+    )
     
     rewritten = None
     if is_flagged:
-        # Added temperature=0.8 to make the 5 sentences more diverse
         outputs = detox_pipe(
             text, 
             max_length=100, 
@@ -50,9 +64,9 @@ async def moderate_text(request: TextRequest):
     return ModerationResponse(
         original_text=text,
         detected_categories={
-            "profanity": profanity_result["contains_profanity"], 
-            "model_label": tox_result['label'],
-            "model_score": tox_result['score']
+            "profanity": profanity_result["contains_profanity"],
+            "ai_toxicity": is_toxic_by_model,
+            "threat_keyword": is_threat_by_keyword
         },
         is_flagged=is_flagged,
         rewritten_text=rewritten