Update shield.py

shield.py

@@ -3,9 +3,9 @@ from flask import Flask, request, Response
 
 app = Flask(__name__)
 
-# --- CONFIGURATION ---
+# --- SETTINGS ---
 NEWAPI_INTERNAL = "http://127.0.0.1:3000"
-OPENAI_MOD_KEY = "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"
+MODERATION_KEY = "sk-proj-LYW3iVcaE..." # Your OpenAI Key for the safety check
 
 TIDB_CONFIG = {
     "host": "gateway01.eu-central-1.prod.aws.tidbcloud.com",
@@ -19,70 +19,46 @@ TIDB_CONFIG = {
     "ssl_verify_cert": True
 }
 
-def log_violation(user_id, prompt):
+def log_violation(user, prompt):
     try:
         conn = mysql.connector.connect(**TIDB_CONFIG)
         cursor = conn.cursor()
-        cursor.execute("INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)", 
-                      (str(user_id), str(prompt)[:1000]))
-        conn.commit()
-        print(f"✅ LOGGED VIOLATION: {user_id}")
+        cursor.execute("INSERT INTO safety_violations (user_id, prompt_content) VALUES (%s, %s)", (user, prompt[:1000]))
         conn.close()
-    except Exception as e:
-        print(f"❌ DATABASE ERROR: {e}")
+    except: pass
 
-# 1. SERVE UI
-@app.route('/')
-def home():
-    with open('index.html', 'r') as f:
-        return f.read()
-
-# 2. PROXY AUTH & MODEL LISTS DIRECTLY (Don't Block These)
-@app.route('/api/<path:subpath>', methods=['GET', 'POST', 'PUT', 'DELETE'])
-def proxy_api(subpath):
-    return forward_request(f"{NEWAPI_INTERNAL}/api/{subpath}")
-
-@app.route('/v1/models', methods=['GET'])
-def proxy_models():
-    return forward_request(f"{NEWAPI_INTERNAL}/v1/models")
-
-# 3. INTERCEPT & PROTECT CHAT
 @app.route('/v1/chat/completions', methods=['POST'])
 def protect_chat():
     data = request.json
-    # Extract text from latest message
-    msgs = data.get('messages', [])
-    last_msg = msgs[-1].get('content', '') if msgs else ""
+    content = " ".join([m.get('content', '') for m in data.get('messages', [])])
     
-    # Check Safety (using your OpenAI Key for moderation)
+    # Safety Check
     try:
-        mod = requests.post("https://api.openai.com/v1/moderations", 
-            headers={"Authorization": f"Bearer {OPENAI_MOD_KEY}"}, 
-            json={"input": last_msg}, timeout=3).json()
-        
-        if mod.get('results', [{}])[0].get('categories', {}).get('sexual/minors'):
-            user_token = request.headers.get('Authorization', 'Anon')
-            log_violation(user_token, last_msg)
-            return {"error": {"message": "Policy Violation: Blocked by Shield.", "type": "safety_error"}}, 403
-    except Exception as e:
-        print(f"Moderation Warning: {e}")
+        res = requests.post("https://api.openai.com/v1/moderations", 
+                           headers={"Authorization": f"Bearer {MODERATION_KEY}"}, 
+                           json={"input": content}, timeout=3).json()
+        if res.get('results', [{}])[0].get('categories', {}).get('sexual/minors'):
+            log_violation(request.headers.get('Authorization', 'Anon'), content)
+            return {"error": {"message": "Shield Block: Safety Violation"}}, 403
+    except: pass
 
-    # If Safe -> Forward to NewAPI (using the USER'S token from the UI)
-    return forward_request(f"{NEWAPI_INTERNAL}/v1/chat/completions")
+    return forward_to_newapi(request.path)
 
-def forward_request(url):
-    try:
-        resp = requests.request(
-            method=request.method,
-            url=url,
-            headers={k: v for k, v in request.headers if k.lower() != 'host'},
-            data=request.get_data(),
-            allow_redirects=False
-        )
-        return Response(resp.content, resp.status_code, resp.headers.items())
-    except Exception as e:
-        return {"error": str(e)}, 502
+@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE'])
+@app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE'])
+def catch_all(path):
+    return forward_to_newapi(path)
+
+def forward_to_newapi(path):
+    resp = requests.request(
+        method=request.method,
+        url=f"{NEWAPI_INTERNAL}/{path}",
+        headers={k: v for k, v in request.headers if k.lower() != 'host'},
+        data=request.get_data(),
+        params=request.args,
+        allow_redirects=False
+    )
+    return Response(resp.content, resp.status_code, resp.headers.items())
 
 if __name__ == '__main__':
-    # Start Shield on 7860
     app.run(host='0.0.0.0', port=7860)