Update src/pipeline.py

src/pipeline.py

@@ -1,17 +1,12 @@
 from dataclasses import dataclass
 from typing import Dict, Any
+
 from .detector import JailbreakDetector
 from .repair import RepairEngine
 
 
 @dataclass
 class JailbreakPipeline:
-    """
-    High-level wrapper around:
-    - jailbreak detection
-    - prompt repair
-    """
-
     consider_output: bool = False
 
     def __post_init__(self):
@@ -19,38 +14,18 @@ class JailbreakPipeline:
         self.repair_engine = RepairEngine()
 
     def detect(self, prompt: str):
-        """
-        Run jailbreak rule matching
-        """
         return self.detector.score(prompt)
 
-    def repair_prompt(self, prompt: str) -> str:
-        """
-        Rewrite unsafe content into a safe, generic assistant response.
-        """
-        return self.repair_engine.repair(prompt)
-
     def process(self, prompt: str) -> Dict[str, Any]:
-    """
-    Unified API called from Gradio UI
-    Returns:
-        {
-            risk_score: float,
-            fired_rules: list,
-            safe_output: str,
-            metadata: dict
-        }
-    """
-    result = self.detect(prompt)
-
-    # pass the risk score (or the full result) depending on your RepairEngine API
-    safe = self.repair_engine.repair(prompt, result.risk_score)
+        result = self.detect(prompt)
 
-    return {
-        "risk_score": result.risk_score,
-        "fired_rules": result.fired_rules,
-        "safe_output": safe,
-        "metadata": result.metadata,
-    }
+        # If your RepairEngine accepts only (prompt), change this to repair(prompt)
+        safe = self.repair_engine.repair(prompt, result.risk_score)
 
+        return {
+            "risk_score": result.risk_score,
+            "fired_rules": result.fired_rules,
+            "safe_output": safe,
+            "metadata": result.metadata,
+        }