Spaces:
Sleeping
Sleeping
| from flask import Flask, request, jsonify, send_from_directory | |
| from flask_cors import CORS | |
| import os | |
| import sys | |
| import time | |
| import math | |
| # Force UTF-8 for Windows | |
| sys.stdout.reconfigure(encoding="utf-8") | |
| sys.stderr.reconfigure(encoding="utf-8") | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| from pipeline import HybridClassifier | |
| app = Flask(__name__, static_folder="dashboard", static_url_path="") | |
| CORS(app) | |
| # Load pipeline once at startup | |
| print("Loading AegisPrompt Hybrid Classifier...") | |
| classifier = HybridClassifier() | |
| print("Classifier ready.") | |
| # ────────────────────────────────────────────── | |
| # Serve dashboard | |
| # ────────────────────────────────────────────── | |
| def index(): | |
| return send_from_directory("dashboard", "index.html") | |
| # ────────────────────────────────────────────── | |
| # POST /api/scan — Live Inspector | |
| # Body: { "text": "...", "w_heur": 0.4, "w_model": 0.4, "w_tfidf": 0.2, "threshold": 0.5 } | |
| # ────────────────────────────────────────────── | |
| RULE_DESCRIPTIONS = { | |
| "1_instruction_override": "Instruction Override — Matches 'ignore/forget/bypass previous instructions'", | |
| "2_system_prompt_leak": "System Prompt Leak — Detects attempts to reveal system prompt", | |
| "3_roleplay_jailbreak": "Roleplay Jailbreak — Matches DAN, sudo, developer mode persona adoption", | |
| "4_output_formatting": "Output Formatting Exploit — base64/hex encode+output patterns", | |
| "5_virtualization_shell": "Virtualization/Shell — Simulating bash/python terminal", | |
| "6_reverse_psychology": "Reverse Psychology — Social engineering via urgency or emotion", | |
| "7_payload_splitting": "Payload Splitting — Concat/spell-out obfuscation", | |
| "8_hypothetical_scenario":"Hypothetical Scenario — Fictional framing to bypass guards", | |
| } | |
| def scan(): | |
| data = request.get_json(force=True) | |
| text = data.get("text", "").strip() | |
| if not text: | |
| return jsonify({"error": "No text provided"}), 400 | |
| w_heur = float(data.get("w_heur", 0.4)) | |
| w_model = float(data.get("w_model", 0.4)) | |
| w_tfidf = float(data.get("w_tfidf", 0.2)) | |
| threshold = float(data.get("threshold", 0.5)) | |
| t_start = time.perf_counter() | |
| result = classifier.predict( | |
| text, | |
| w_heur=w_heur, w_model=w_model, w_tfidf=w_tfidf, | |
| threshold=threshold | |
| ) | |
| latency_ms = round((time.perf_counter() - t_start) * 1000, 1) | |
| # Build rules table | |
| rules_triggered = result["details"]["rules_triggered"] | |
| rules_table = [ | |
| { | |
| "id": r, | |
| "description": RULE_DESCRIPTIONS.get(r, r), | |
| "risk": "High" if r in ("1_instruction_override", "3_roleplay_jailbreak") else "Medium", | |
| } | |
| for r in rules_triggered | |
| ] | |
| # Entropy (Shannon) as a simple text complexity indicator | |
| from collections import Counter | |
| freq = Counter(text.lower()) | |
| total = sum(freq.values()) | |
| entropy = -sum((c/total) * math.log2(c/total) for c in freq.values() if c > 0) | |
| entropy = round(entropy, 2) | |
| # Complexity label | |
| words = len(text.split()) | |
| complexity = "Low" if words < 15 else ("Medium" if words < 40 else "High") | |
| return jsonify({ | |
| "prediction": result["prediction"], | |
| "decision": result["decision"], | |
| "score": round(result["score"], 4), | |
| "heuristic": round(result["details"]["heuristic_score"], 4), | |
| "model_prob": round(result["details"]["model_prob"] or 0, 4), | |
| "tfidf": round(result["details"]["tfidf_score"] or 0, 4), | |
| "rules": rules_table, | |
| "latency_ms": latency_ms, | |
| "entropy": entropy, | |
| "complexity": complexity, | |
| "flagged_substrings": len(rules_triggered), | |
| }) | |
| # ────────────────────────────────────────────── | |
| # POST /api/batch — Batch Analysis | |
| # Body: { "prompts": ["...", "..."] } | |
| # ────────────────────────────────────────────── | |
| def batch(): | |
| data = request.get_json(force=True) | |
| prompts = data.get("prompts", []) | |
| if not prompts: | |
| return jsonify({"error": "No prompts provided"}), 400 | |
| results = [] | |
| for i, text in enumerate(prompts): | |
| res = classifier.predict(text.strip()) | |
| results.append({ | |
| "id": i + 1, | |
| "text": text[:120], | |
| "prediction": res["prediction"], | |
| "decision": res["decision"], | |
| "score": round(res["score"], 4), | |
| "rules": res["details"]["rules_triggered"], | |
| "model_prob": round(res["details"]["model_prob"] or 0, 4), | |
| }) | |
| total = len(results) | |
| injections = sum(1 for r in results if r["prediction"] == 1) | |
| benign = total - injections | |
| avg_score = round(sum(r["score"] for r in results) / total, 4) if total else 0 | |
| # Score distribution buckets (0-1 in 10 bins) | |
| bins = [0] * 10 | |
| for r in results: | |
| bucket = min(int(r["score"] * 10), 9) | |
| bins[bucket] += 1 | |
| return jsonify({ | |
| "total": total, | |
| "injections": injections, | |
| "benign": benign, | |
| "avg_score": avg_score, | |
| "items": results, | |
| "histogram": bins, | |
| }) | |
| # ────────────────────────────────────────────── | |
| # GET /api/metrics — Diagnostic Metrics Center | |
| # Loads from results/metrics.json | |
| # ────────────────────────────────────────────── | |
| import json as _json | |
| def metrics(): | |
| metrics_path = os.path.join(os.path.dirname(__file__), "results", "metrics.json") | |
| if not os.path.exists(metrics_path): | |
| return jsonify({"error": "metrics.json not found. Run evaluate_pipeline.py first."}), 404 | |
| with open(metrics_path, "r") as f: | |
| m = _json.load(f) | |
| # Derive confusion matrix values from available data (67-sample test set) | |
| total_test = 67 | |
| total_benign = 40 | |
| total_injection = 27 | |
| tp = round(m["test_recall"] * total_injection) | |
| fn = total_injection - tp | |
| fp = m["test_set_fps"] | |
| tn = total_benign - fp | |
| return jsonify({ | |
| "accuracy": round(m["test_accuracy"] * 100, 2), | |
| "f1": round(m["test_f1"] * 100, 2), | |
| "precision": round(m["test_precision"] * 100, 2), | |
| "recall": round(m["test_recall"] * 100, 2), | |
| "auc": round(m["test_auc"], 4), | |
| "red_team_total": m["red_team_total"], | |
| "red_team_blocked": m["red_team_blocked"], | |
| "red_team_block_rate": round(m["red_team_block_rate"], 1), | |
| "tricky_fps": m["tricky_benign_fps"], | |
| "test_fps": m["test_set_fps"], | |
| "confusion_matrix": {"tp": tp, "fp": fp, "fn": fn, "tn": tn}, | |
| }) | |
| # ────────────────────────────────────────────── | |
| # GET /api/false_positives — FP inspector data | |
| # ────────────────────────────────────────────── | |
| def false_positives(): | |
| fp_path = os.path.join(os.path.dirname(__file__), "results", "false_positives.md") | |
| if not os.path.exists(fp_path): | |
| return jsonify({"items": []}) | |
| # Parse the false_positives.md for flagged entries | |
| items = [] | |
| with open(fp_path, "r", encoding="utf-8") as f: | |
| content = f.read() | |
| import re | |
| # Extract flagged prompt blocks | |
| blocks = re.findall( | |
| r'\[FLAGGED\].*?\n- \*\*Text\*\*: `(.+?)`\n- \*\*Combined Score\*\*: `(.+?)`\n- \*\*Heuristics Rules\*\*: `(.+?)`', | |
| content | |
| ) | |
| for text, score, rules in blocks: | |
| rules_list = eval(rules) if rules.startswith("[") else [rules] if rules != "[]" else [] | |
| items.append({ | |
| "text": text[:100], | |
| "score": float(score), | |
| "rules": rules_list, | |
| "flag_type": rules_list[0].replace("_", " ").title() if rules_list else "Semantic Shadow" | |
| }) | |
| return jsonify({"items": items}) | |
| if __name__ == "__main__": | |
| port = int(os.environ.get("PORT", 5000)) | |
| print("\n" + "="*60) | |
| print(f" AegisPrompt Dashboard running at http://localhost:{port}") | |
| print("="*60 + "\n") | |
| app.run(debug=False, port=port, host="0.0.0.0") | |