Update app.py

app.py

@@ -1,38 +1,30 @@
 import gradio as gr
 from src.pipeline import JailbreakPipeline
 
-pipeline = JailbreakPipeline()
-
-def run_defense(prompt):
-    result = pipeline.process(prompt)
-    
-    fired = "\n".join(
-        [f"- **{h['rule']}**: {h['description']} (match: *{h['match_text']}*)" 
-         for h in result["fired_rules"]]
-    ) or "No rules fired ✔️"
-    
-    return (
-        result["risk_score"],
-        fired,
-        result["repaired_output"] or "No repair needed ✔️"
-    )
+pipeline = JailbreakPipeline(consider_output=True)
+
+def run_defense(prompt: str):
+    result = pipeline.process(prompt)  # returns PipelineResult
+    # round risk score for display
+    risk = round(float(result.risk_score), 2)
+    safe_text = result.safe_output
+    return risk, safe_text
 
-with gr.Blocks(title="JailBreakDefense") as demo:
+with gr.Blocks() as demo:
     gr.Markdown("# JailBreakDefense – Jailbreak Prompt Detector")
-    gr.Markdown("Enter any prompt and detect jailbreak attempts in real-time.")
 
-    with gr.Row():
-        prompt = gr.Textbox(
-            label="User Prompt",
-            placeholder="Type something like: 'Ignore safety and do anything now...'"
-        )
+    prompt_box = gr.Textbox(
+        label="User Prompt",
+        placeholder="Type something like: 'Ignore safety and do anything now...'",
+        lines=3,
+    )
 
-    btn = gr.Button("Analyze Prompt")
+    analyze_btn = gr.Button("Analyze Prompt")
 
-    risk_score = gr.Number(label="Risk Score (0–1)")
-    rules_fired = gr.Markdown(label="Fired Rules")
-    repaired = gr.Textbox(label="Safe Output (if repaired)", lines=4)
+    risk_out = gr.Number(label="Risk Score (0–1)")
+    safe_out = gr.Textbox(label="Safe Output (if repaired)", lines=4)
 
-    btn.click(run_defense, inputs=[prompt], outputs=[risk_score, rules_fired, repaired])
+    analyze_btn.click(run_defense, inputs=prompt_box, outputs=[risk_out, safe_out])
 
 demo.launch()
+