src/pipeline.py

from dataclasses import dataclass
from typing import Dict, Any

from .detector import JailbreakDetector
from .repair import RepairEngine


@dataclass
class JailbreakPipeline:
    consider_output: bool = False

    def __post_init__(self):
        self.detector = JailbreakDetector(consider_output=self.consider_output)
        self.repair_engine = RepairEngine()

    def detect(self, prompt: str):
        return self.detector.score(prompt)

    def process(self, prompt: str) -> Dict[str, Any]:
        result = self.detect(prompt)

        # 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,
        }