""" genai_app.py — GenAI Shield V2 Flask Application. Powered by Llama-Prompt-Guard-2-86M for pre-inference prompt screening. Endpoints: GET / → Chat interface GET /genai-monitoring → Real-time GenAI SIEM dashboard GET /genai-stream → SSE event stream POST /genai-chat → Send a prompt, get monitored response GET /guard-stats → Prompt Guard model statistics Configure via environment variables: GEMINI_API_KEY=... GENAI_PORT=5001 (default) """ import os import json import time import queue import threading from flask import Flask, request, jsonify, render_template, Response, stream_with_context from flask_cors import CORS from gemini_adapter import GeminiAdapter from prompt_guard_engine import PromptGuardEngine from prompt_guard_text_guard import PromptGuardTextGuard from text_monitor import TextMonitor from attachment_guard import AttachmentGuard app = Flask(__name__) CORS(app) # ── System prompt ───────────────────────────────────────────────────────────── SYSTEM_PROMPT = os.getenv( "GENAI_SYSTEM_PROMPT", "You are a helpful AI assistant. Be concise, accurate, and professional." ) # ── Initialise Prompt Guard engine ──────────────────────────────────────────── print("[GenAI Shield V2] Initialising Prompt Guard engine...") PG_ENGINE = PromptGuardEngine().load() GUARD = PromptGuardTextGuard(PG_ENGINE) print("[GenAI Shield V2] Prompt Guard ready.") # ── Initialise LLM adapter + post-inference monitor ────────────────────────── ADAPTER = GeminiAdapter(system_prompt=SYSTEM_PROMPT) MONITOR = TextMonitor(ADAPTER, system_prompt=SYSTEM_PROMPT) # ───────────────────────────────────────────────────────────────────────────── # SSE Broadcast Queue BROADCAST_QUEUES = [] def broadcast(event_type: str, data: dict): event = { "timestamp": time.strftime("%H:%M:%S"), "type": event_type, "threat_score": data.get("threat_score", 0), "flags": data.get("flags", []), "reason": data.get("reason", "CLEAN"), "source": data.get("source", "Web UI"), "prompt": data.get("prompt", "")[:120], "response": data.get("response", "")[:200], "latency_ms": data.get("latency_ms", 0), "checks": data.get("checks", {}), "model": ADAPTER.get_model_name(), "prompt_guard_score": data.get("prompt_guard_score", 0), } for q in BROADCAST_QUEUES: q.put(event) # ── Routes ──────────────────────────────────────────────────────────────────── @app.route("/") def index(): return render_template("genai.html", model=ADAPTER.get_model_name()) @app.route("/sidecar") def sidecar_ui(): """Sidecar streaming chat UI.""" return render_template("sidecar.html") @app.route("/dataflow") def dataflow_ui(): """Real-time data flow visualization dashboard.""" return render_template("dataflow.html") @app.route("/genai-monitoring") def monitoring(): return render_template("genai_monitoring.html", model=ADAPTER.get_model_name()) @app.route("/genai-stream") def stream(): def event_stream(): q = queue.Queue() BROADCAST_QUEUES.append(q) try: while True: event = q.get() yield f"data: {json.dumps(event)}\n\n" except GeneratorExit: BROADCAST_QUEUES.remove(q) return Response(stream_with_context(event_stream()), mimetype="text/event-stream") @app.route("/genai-chat", methods=["POST"]) def chat(): data = request.json prompt = data.get("prompt", "").strip() attachment = data.get("attachment") # { filename, content_b64 } source = data.get("source", "Web UI") if not prompt: return jsonify({"error": "Empty prompt"}), 400 # ── LAYER 0: Attachment Extraction & Screening ─────────────────────────── attachment_text = "" if attachment: filename = attachment.get("filename", "unknown") b64 = attachment.get("content_b64", "") extracted = AttachmentGuard.extract_text(filename, b64) if extracted["error"]: return jsonify({"error": extracted["error"]}), 400 attachment_text = extracted["text"] # Screen attachment text through Prompt Guard att_guard_result = AttachmentGuard.screen_with_guard(GUARD, filename, attachment_text) if att_guard_result["blocked"]: pg_score = att_guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0) broadcast("BLOCKED_PRE_INFERENCE", { "threat_score": att_guard_result.get("threat_score", 100), "flags": att_guard_result["flags"], "reason": att_guard_result["reason"], "prompt": prompt, "response": f"[BLOCKED — Malicious Attachment: {filename}]", "source": source, "latency_ms": 0, "checks": att_guard_result.get("checks", {}), "prompt_guard_score": pg_score, }) return jsonify({ "blocked": True, "error": "ATTACHMENT_REJECTED", "reason": att_guard_result["reason"], "threat_score": att_guard_result.get("threat_score", 100), "flags": att_guard_result["flags"] }), 403 # ── LAYER 1: Pre-Inference Guard (Prompt Guard model) ──────────────────── guard_start = time.time() guard_result = GUARD.screen(prompt) guard_lat = round((time.time() - guard_start) * 1000, 2) pg_score = guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0) if guard_result["blocked"]: broadcast("BLOCKED_PRE_INFERENCE", { "threat_score": guard_result["threat_score"], "flags": guard_result["flags"], "reason": guard_result["reason"], "prompt": prompt, "response": "[BLOCKED — LLM never called]", "source": source, "latency_ms": guard_lat, "checks": guard_result["checks"], "prompt_guard_score": pg_score, }) return jsonify({ "blocked": True, "response": None, "error": "PROMPT_REJECTED_BY_GUARD", "reason": guard_result["reason"], "threat_score": guard_result["threat_score"], "flags": guard_result["flags"], "prompt_guard_score": pg_score, "latency_breakdown": { "guard_ms": guard_lat, "model_ms": 0, "monitor_ms": 0 } }), 403 # ── LAYER 2: LLM Inference ──────────────────────────────────────────────── try: model_start = time.time() full_prompt = prompt if attachment_text: full_prompt = ( f"Context from attachment '{filename}':\n---\n{attachment_text}\n" f"---\nUser prompt: {prompt}" ) response = ADAPTER.chat(full_prompt, system_prompt=SYSTEM_PROMPT) model_lat = round((time.time() - model_start) * 1000, 2) except Exception as e: return jsonify({"error": f"LLM error: {str(e)}"}), 500 # ── LAYER 3: Post-Inference Monitor ─────────────────────────────────────── monitor_start = time.time() monitor_result = MONITOR.analyze(prompt, response, source=source) monitor_lat = round((time.time() - monitor_start) * 1000, 2) total_lat = round(guard_lat + model_lat + monitor_lat, 2) # Determine final threat level threat_score = max(guard_result["threat_score"], monitor_result["threat_score"]) all_flags = guard_result["flags"] + monitor_result["flags"] # Broadcast to dashboard event_type = "SUSPICIOUS" if threat_score >= 30 else "INFERENCE" broadcast(event_type, { "threat_score": threat_score, "flags": all_flags, "reason": monitor_result["reason"], "prompt": prompt, "response": response, "source": source, "latency_ms": total_lat, "prompt_guard_score": pg_score, "checks": { "guard": guard_result["checks"], "monitor": monitor_result["checks"], "breakdown": { "guard_ms": guard_lat, "model_ms": model_lat, "monitor_ms": monitor_lat } }, }) return jsonify({ "blocked": False, "response": response, "threat_score": threat_score, "flags": all_flags, "latency_ms": total_lat, "prompt_guard_score": pg_score, "latency_breakdown": { "guard_ms": guard_lat, "model_ms": model_lat, "monitor_ms": monitor_lat }, "model": ADAPTER.get_model_name(), }) @app.route("/guard-stats") def guard_stats(): """Return Prompt Guard engine statistics.""" return jsonify(PG_ENGINE.stats()) def _start_sidecar_subprocess(): """ Optionally launch the sidecar as a sub-process so both UIs run together. Controlled via LAUNCH_SIDECAR=true env var. """ import subprocess, sys sidecar_port = int(os.getenv("SIDECAR_PORT", 5050)) print(f"[GenAI Shield] Launching sidecar on :{sidecar_port}...") proc = subprocess.Popen( [sys.executable, "-m", "uvicorn", "sidecar.app:app", "--host", "0.0.0.0", "--port", str(sidecar_port), "--log-level", "warning"], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) def _pipe(): for line in proc.stdout: print("[sidecar]", line.decode(errors='replace').rstrip()) threading.Thread(target=_pipe, daemon=True).start() return proc if __name__ == "__main__": port = int(os.getenv("GENAI_PORT", 5001)) print(f"GenAI Shield V2 starting on port {port}") print(f"LLM Model: {ADAPTER.get_model_name()}") print(f"Guard: Llama-Prompt-Guard-2-86M") print(f"Sidecar UI: http://localhost:{port}/sidecar") if os.getenv("LAUNCH_SIDECAR", "").lower() in ("1", "true", "yes"): _start_sidecar_subprocess() app.run(host="0.0.0.0", port=port, debug=False)