Update shield.py

shield.py

@@ -3,55 +3,72 @@ from flask import Flask, request, Response
 
 app = Flask(__name__)
 
-# YOUR OPENAI PROJECT KEY
+# --- CONFIGURATION ---
 OPENAI_KEY = "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"
 
-# TIDB CONNECTION DATA
 TIDB_CONFIG = {
     "host": "gateway01.eu-central-1.prod.aws.tidbcloud.com",
     "port": 4000,
     "user": "uiSKPXCQ9Gzb4co.root",
     "password": "Bxg8rpU27gyH60E0",
     "database": "test",
-    "ssl_verify_cert": True,
-    "ssl_ca": "/etc/ssl/certs/ca-certificates.crt" 
+    "autocommit": True,
+    "ssl_verify_cert": False,  # Bypasses the certificate path issue in HF
+    "use_pure": True           # More stable in Docker environments
 }
 
-def log_to_tidb(user, prompt):
+def log_to_tidb(user_id, prompt):
+    """Force logs the violation into TiDB."""
     try:
+        print(f"DEBUG: Attempting to log violation for user {user_id}...")
         conn = mysql.connector.connect(**TIDB_CONFIG)
         cursor = conn.cursor()
-        cursor.execute("INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)", (user, prompt[:1000]))
+        
+        # SQL matches the columns you created in your DESCRIBE screenshot
+        query = "INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)"
+        values = (str(user_id), str(prompt)[:1000])
+        
+        cursor.execute(query, values)
         conn.commit()
+        
+        print("✅ SUCCESS: Violation logged to TiDB.")
         cursor.close()
         conn.close()
     except Exception as e:
-        print(f"TiDB Error: {e}")
+        # If it still fails, this will show up in your 'Logs' tab on Hugging Face
+        print(f"❌ DATABASE ERROR: {str(e)}")
 
 @app.route('/v1/chat/completions', methods=['POST'])
 def protect_and_proxy():
     data = request.json
     messages = data.get('messages', [])
-    text_to_check = " ".join([m.get('content', '') for m in messages])
+    full_text = " ".join([m.get('content', '') for m in messages])
 
-    # 1. OpenAI Moderation Check (FREE)
+    # 1. Moderation Check (FREE)
     res = requests.post(
         "https://api.openai.com/v1/moderations",
         headers={"Authorization": f"Bearer {OPENAI_KEY}"},
-        json={"input": text_to_check}
+        json={"input": full_text}
     ).json()
     
-    # 2. Block only for 'sexual/minors' (CSAM)
-    results = res['results'][0]
-    if results['categories'].get('sexual/minors'):
+    # 2. Block only for 'sexual/minors'
+    results = res.json().get('results', [{}])[0]
+    if results.get('categories', {}).get('sexual/minors'):
         user_auth = request.headers.get('Authorization', 'Anonymous')
-        log_to_tidb(user_auth, text_to_check)
+        print(f"!!! CSAM DETECTED: {user_auth} !!!")
+        
+        # Trigger the log function
+        log_to_tidb(user_auth, full_text)
+        
         return {"error": {"message": "Policy Violation: CSAM is strictly prohibited.", "type": "safety_error"}}, 403
 
-    # 3. Safe? Forward to NewAPI on internal port 3000
-    resp = requests.post("http://127.0.0.1:3000/v1/chat/completions", 
-                         json=data, headers=dict(request.headers))
-    return Response(resp.content, resp.status_code, resp.headers.items())
+    # 3. If safe, pass to the real NewAPI on internal port 3000
+    try:
+        resp = requests.post("http://127.0.0.1:3000/v1/chat/completions", 
+                             json=data, headers=dict(request.headers), timeout=60)
+        return Response(resp.content, resp.status_code, resp.headers.items())
+    except Exception as e:
+        return {"error": {"message": f"Proxy error: {str(e)}", "type": "internal_error"}}, 500
 
 # Proxy all other dashboard/admin traffic
 @app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])