Enhance project configuration and structure: Added build-system requirements in pyproject.toml, updated README for installation instructions, and refactored attack classes to utilize a new create_test_case method. Introduced a plugins.yaml for component registration and removed the legacy registry initialization file.

2026-07-11 07:23:45 +02:00 · 2025-12-18 23:32:04 +08:00
parent becb9c8624
commit 84a4d1708e
49 changed files with 1474 additions and 727 deletions
@@ -27,6 +27,10 @@ class BaseComponent(ABC):
        """
        self.config = config or {}
        self.name = self.__class__.__name__
+        # Registry name injected by UnifiedRegistry (canonical component id, e.g. "llama_guard_3")
+        self.registry_name = (
+            config.get("_registry_name") if isinstance(config, dict) else None
+        )

        # Initialize logger (using the module name where component is located)
        self.logger = logging.getLogger(self.__class__.__module__)
@@ -77,6 +81,13 @@ class BaseComponent(ABC):
 class BaseAttack(BaseComponent, ABC):
    """Attack method base class (enhanced version)"""

+    # Standard Metadata Keys
+    META_KEY_ATTACK_METHOD = "attack_method"
+    META_KEY_ORIGINAL_PROMPT = "original_prompt"
+    META_KEY_JAILBREAK_PROMPT = "jailbreak_prompt"
+    META_KEY_JAILBREAK_IMAGE = "jailbreak_image_path"
+    META_KEY_TARGET_MODEL = "target_model"
+
    def __init__(
        self, config: Dict[str, Any] = None, output_image_dir: Optional[str] = None
    ):
@@ -135,6 +146,61 @@ class BaseAttack(BaseComponent, ABC):
        """
        pass

+    def create_test_case(
+        self,
+        case_id: str,
+        jailbreak_prompt: str,
+        jailbreak_image_path: str,
+        original_prompt: str,
+        original_image_path: str = None,
+        metadata: Dict[str, Any] = None,
+        **kwargs,
+    ) -> TestCase:
+        """
+        Create a standardized `TestCase` with populated metadata.
+
+        Args:
+            case_id: Unique test case ID
+            jailbreak_prompt: The generated adversarial prompt
+            jailbreak_image_path: Path to the generated adversarial image
+            original_prompt: The original harmful prompt
+            original_image_path: Path to the original image (optional)
+            metadata: Additional metadata
+            **kwargs: Other parameters to store in metadata
+
+        Returns:
+            TestCase object
+        """
+        attack_method = (
+            getattr(self, "registry_name", None)
+            or getattr(self, "name", None)
+            or self.__class__.__name__.replace("Attack", "").lower()
+        )
+
+        base_metadata = {
+            self.META_KEY_ATTACK_METHOD: attack_method,
+            self.META_KEY_ORIGINAL_PROMPT: original_prompt,
+            self.META_KEY_JAILBREAK_PROMPT: jailbreak_prompt,
+            self.META_KEY_JAILBREAK_IMAGE: jailbreak_image_path,
+        }
+
+        if hasattr(self, "cfg") and hasattr(self.cfg, "target_model_name"):
+            if self.cfg.target_model_name:
+                base_metadata[self.META_KEY_TARGET_MODEL] = self.cfg.target_model_name
+
+        if original_image_path:
+            base_metadata["original_image_path"] = original_image_path
+
+        # Merge additional metadata
+        final_metadata = {**base_metadata, **(metadata or {}), **kwargs}
+
+        return TestCase(
+            test_case_id=str(case_id),
+            prompt=jailbreak_prompt,
+            image_path=jailbreak_image_path,
+            metadata=final_metadata,
+        )
+
    def __str__(self):
        return f"{self.name}(config={self.config})"

@@ -173,6 +239,13 @@ class BaseModel(BaseComponent, ABC):
 class BaseDefense(BaseComponent, ABC):
    """Defense method base class (enhanced version)"""

+    # Metadata keys
+    META_KEY_SHOULD_BLOCK = "should_return_default"
+    META_KEY_BLOCKED = "defense_blocked"
+    META_KEY_DEFAULT_RESP = "default_response"
+    META_KEY_GENERATED_RESP = "defense_generated_response"
+    META_KEY_GENERATED_RESP_USED = "defense_generated_response_used"
+
    def __init__(self, config: Dict[str, Any] = None):
        """
        Initialize defense method
@@ -219,6 +292,90 @@ class BaseDefense(BaseComponent, ABC):
        """
        pass

+    def create_defended_case(
+        self,
+        test_case: TestCase,
+        defended_prompt: str = None,
+        defended_image_path: str = None,
+        metadata: Dict[str, Any] = None,
+    ) -> TestCase:
+        """
+        Create a new `TestCase` after applying defense (common helper).
+
+        This keeps downstream pipeline logic and metadata conventions consistent.
+        """
+        defense_method = (
+            getattr(self, "registry_name", None)
+            or getattr(self, "name", None)
+            or self.__class__.__name__.replace("Defense", "").lower()
+        )
+
+        return TestCase(
+            test_case_id=test_case.test_case_id,
+            prompt=defended_prompt,
+            image_path=defended_image_path,
+            metadata={
+                **(test_case.metadata or {}),
+                "defense_method": defense_method,
+                "defended_prompt": defended_prompt,
+                "defended_image_path": defended_image_path,
+                **(metadata or {}),
+            },
+        )
+
+    def block_input(
+        self,
+        test_case: TestCase,
+        default_response: str = "I'm sorry, I can't assist with that.",
+        defended_prompt: str = None,
+        defended_image_path: str = None,
+        metadata: Dict[str, Any] = None,
+    ) -> TestCase:
+        """
+        Mark the test case as 'blocked' by the defense.
+        Pipeline will intercept this and return `default_response` immediately.
+        """
+        defended_case = self.create_defended_case(
+            test_case=test_case,
+            defended_prompt=defended_prompt if defended_prompt is not None else test_case.prompt,
+            defended_image_path=defended_image_path
+            if defended_image_path is not None
+            else test_case.image_path,
+            metadata={
+                **(metadata or {}),
+                self.META_KEY_SHOULD_BLOCK: True,
+                self.META_KEY_BLOCKED: True,
+                self.META_KEY_DEFAULT_RESP: default_response,
+            },
+        )
+        return defended_case
+
+    def reply_directly(
+        self,
+        test_case: TestCase,
+        response_text: str,
+        defended_prompt: str = None,
+        defended_image_path: str = None,
+        metadata: Dict[str, Any] = None,
+    ) -> TestCase:
+        """
+        Defense generates the response directly (bypassing the target model).
+        Pipeline will use `response_text` as the final output.
+        """
+        defended_case = self.create_defended_case(
+            test_case=test_case,
+            defended_prompt=defended_prompt if defended_prompt is not None else test_case.prompt,
+            defended_image_path=defended_image_path
+            if defended_image_path is not None
+            else test_case.image_path,
+            metadata={
+                **(metadata or {}),
+                self.META_KEY_GENERATED_RESP: response_text,
+                self.META_KEY_GENERATED_RESP_USED: True,
+            },
+        )
+        return defended_case
+
    def __str__(self):
        return f"{self.name}(config={self.config})"

@@ -237,7 +394,9 @@ class BaseEvaluator(BaseComponent, ABC):
        super().__init__(config)

    @abstractmethod
-    def evaluate(self, model_response: ModelResponse, **kwargs) -> EvaluationResult:
+    def evaluate_response(
+        self, model_response: ModelResponse, **kwargs
+    ) -> EvaluationResult:
        """
        Evaluate model response

@@ -1,69 +0,0 @@
-def initialize_registry_lazy_imports():
-    """Initialize lazy loading mappings to unified registry"""
-
-    # Attack method lazy loading mappings
-    attack_lazy_imports = {
-        "figstep": ("attacks.figstep.attack", "FigStepAttack"),
-        "cs_dj": ("attacks.cs_dj.attack", "CSDJAttack"),
-        "hades": ("attacks.hades.attack", "HadesAttack"),
-        "qr": ("attacks.qr.attack", "QRAttack"),
-        "mml": ("attacks.mml.attack", "MMLAttack"),
-        "si": ("attacks.si.attack", "SIAttack"),
-        "jood": ("attacks.jood.attack", "JOODAttack"),
-        "himrd": ("attacks.himrd.attack", "HIMRDAttack"),
-        "bap": ("attacks.bap.attack", "BAPAttack"),
-        "visual_adv": ("attacks.visual_adv.attack", "VisualAdvAttack"),
-        "viscra": ("attacks.viscra.attack", "VisCRAAttack"),
-        "umk": ("attacks.umk.attack", "UMKAttack"),
-        "pba": ("attacks.pba.attack", "PBAAttack"),
-        "imgjp": ("attacks.imgjp.attack", "ImgJPAttack"),
-    }
-
-    # Model lazy loading mappings
-    model_lazy_imports = {
-        "openai": ("models.openai_model", "OpenAIModel"),
-        "google": ("models.google_model", "GoogleModel"),
-        "anthropic": ("models.anthropic_model", "AnthropicModel"),
-        "qwen": ("models.qwen_model", "QwenModel"),
-        "doubao": ("models.doubao_model", "DoubaoModel"),
-        "vllm": ("models.vllm_model", "VLLMModel"),
-        "mistral": ("models.mistral_model", "MistralModel"),
-        "any": (
-            "models.openai_model",
-            "OpenAIModel",
-        ),  # "any" uses OpenAI-compatible API
-    }
-
-    # Defense method lazy loading mappings
-    defense_lazy_imports = {
-        "adashield": ("defenses.adashield", "AdaShieldDefense"),
-        "dps": ("defenses.dps", "DPSDefense"),
-        "ecso": ("defenses.ecso", "ECSODefense"),
-        "jailguard": ("defenses.jailguard", "JailGuardDefense"),
-        "llavaguard": ("defenses.llavaguard", "LlavaGuardDefense"),
-        "qguard": ("defenses.qguard", "QGuardDefense"),
-        "shieldlm": ("defenses.shieldlm", "ShieldLMDefense"),
-        "uniguard": ("defenses.uniguard", "UniguardDefense"),
-        "cider": ("defenses.cider", "CIDERDefense"),
-        "mllm_protector": ("defenses.mllm_protector", "MLLMProtectorDefense"),
-        "llama_guard_3": ("defenses.llama_guard_3", "LlamaGuard3Defense"),
-        "hiddendetect": ("defenses.hiddendetect", "HiddenDetectDefense"),
-        "guardreasoner_vl": ("defenses.guardreasoner_vl", "GuardReasonerVLDefense"),
-        "llama_guard_4": ("defenses.llama_guard_4", "LlamaGuard4Defense"),
-        "vlguard": ("defenses.vlguard", "VLGuardDefense"),
-        "spa_vl": ("defenses.spa_vl", "SPAVLDefense"),
-        "saferlhf_v": ("defenses.saferlhf_v", "SafeRLHFVDefense"),
-        "coca": ("defenses.coca", "CoCADefense"),
-    }
-
-    # Evaluator lazy loading mappings
-    evaluator_lazy_imports = {
-        "default_judge": ("evaluators.default_judge", "DefaultJudge"),
-    }
-
-    return {
-        "attacks": attack_lazy_imports,
-        "models": model_lazy_imports,
-        "defenses": defense_lazy_imports,
-        "evaluators": evaluator_lazy_imports,
-    }
@@ -2,12 +2,19 @@
 Unified registry system - merged version, removed redundant code
 """

-from typing import Dict, Any, Optional, Type, List
+from typing import Dict, Any, Optional, Type, List, TYPE_CHECKING
 import logging
 import importlib
 from typing import Set
+from pathlib import Path
+import yaml

-from .base_classes import BaseAttack, BaseModel, BaseDefense, BaseEvaluator
+# NOTE:
+# - `core.base_classes` is the canonical place for component contracts (BaseAttack/BaseModel/BaseDefense/BaseEvaluator).
+# - This registry module should NOT re-export Base* at runtime; keep imports here type-checking only.
+
+if TYPE_CHECKING:
+    from .base_classes import BaseAttack, BaseModel, BaseDefense, BaseEvaluator


 class UnifiedRegistry:
@@ -22,6 +29,49 @@ class UnifiedRegistry:
        self.logger = logging.getLogger(__name__)
        # Imported module cache to avoid duplicate imports
        self._imported_modules: Set[str] = set()
+        # Lazy mapping cache (names only). This should not trigger heavy imports.
+        self._lazy_mappings_cache: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def _get_lazy_mappings(self) -> Dict[str, Dict[str, Any]]:
+        """Get lazy import mappings from config/plugins.yaml (cached)."""
+        if self._lazy_mappings_cache is not None:
+            return self._lazy_mappings_cache
+        self._lazy_mappings_cache = self._load_plugins_yaml()
+        return self._lazy_mappings_cache
+
+    def _load_plugins_yaml(self) -> Dict[str, Dict[str, Any]]:
+        """Load plugin mappings from `config/plugins.yaml`."""
+        # Repo root = parent of `core/`
+        repo_root = Path(__file__).resolve().parent.parent
+        plugins_file = repo_root / "config" / "plugins.yaml"
+        if not plugins_file.exists():
+            raise FileNotFoundError(
+                f"config/plugins.yaml not found at {plugins_file}. This project requires config/plugins.yaml."
+            )
+
+        with open(plugins_file, "r", encoding="utf-8") as f:
+            data = yaml.safe_load(f) or {}
+
+        plugins = data.get("plugins") or {}
+        mappings: Dict[str, Dict[str, Any]] = {}
+        for section in ("attacks", "models", "defenses", "evaluators"):
+            section_map = plugins.get(section) or {}
+            converted: Dict[str, Any] = {}
+            for name, spec in section_map.items():
+                if not isinstance(spec, (list, tuple)) or len(spec) != 2:
+                    raise ValueError(
+                        f"Invalid plugin spec for {section}.{name}: expected [module, class], got {spec!r}"
+                    )
+                module_path, class_name = spec
+                converted[name] = (module_path, class_name)
+            mappings[section] = converted
+
+        # Ensure all sections exist
+        mappings.setdefault("attacks", {})
+        mappings.setdefault("models", {})
+        mappings.setdefault("defenses", {})
+        mappings.setdefault("evaluators", {})
+        return mappings

    def register_attack(self, name: str, attack_class: Type["BaseAttack"]) -> None:
        """Register attack method"""
@@ -64,11 +114,9 @@ class UnifiedRegistry:
        if name in self.attack_registry:
            return self.attack_registry[name]

-        # Get mapping information from registry_init
+        # Get mapping information from config/plugins.yaml
        try:
-            from .registry_init import initialize_registry_lazy_imports
-
-            mappings = initialize_registry_lazy_imports()
+            mappings = self._get_lazy_mappings()
            if name in mappings["attacks"]:
                module_path, class_name = mappings["attacks"][name]
                module = importlib.import_module(module_path)
@@ -98,11 +146,9 @@ class UnifiedRegistry:
        if name in self.model_registry:
            return self.model_registry[name]

-        # Get mapping information from registry_init
+        # Get mapping information from config/plugins.yaml
        try:
-            from .registry_init import initialize_registry_lazy_imports
-
-            mappings = initialize_registry_lazy_imports()
+            mappings = self._get_lazy_mappings()
            if name in mappings["models"]:
                module_path, class_name = mappings["models"][name]
                module = importlib.import_module(module_path)
@@ -127,11 +173,9 @@ class UnifiedRegistry:
        if name in self.defense_registry:
            return self.defense_registry[name]

-        # Get mapping information from registry_init
+        # Get mapping information from config/plugins.yaml
        try:
-            from .registry_init import initialize_registry_lazy_imports
-
-            mappings = initialize_registry_lazy_imports()
+            mappings = self._get_lazy_mappings()
            if name in mappings["defenses"]:
                module_path, class_name = mappings["defenses"][name]
                module = importlib.import_module(module_path)
@@ -160,11 +204,9 @@ class UnifiedRegistry:
        if name in self.evaluator_registry:
            return self.evaluator_registry[name]

-        # Get mapping information from registry_init
+        # Get mapping information from config/plugins.yaml
        try:
-            from .registry_init import initialize_registry_lazy_imports
-
-            mappings = initialize_registry_lazy_imports()
+            mappings = self._get_lazy_mappings()
            if name in mappings["evaluators"]:
                module_path, class_name = mappings["evaluators"][name]
                module = importlib.import_module(module_path)
@@ -199,7 +241,9 @@ class UnifiedRegistry:
        attack_class = self.get_attack(name)
        if attack_class:
            try:
-                return attack_class(config=config, output_image_dir=output_image_dir)
+                cfg = dict(config or {})
+                cfg.setdefault("_registry_name", name)
+                return attack_class(config=cfg, output_image_dir=output_image_dir)
            except Exception as e:
                self.logger.error(f"Failed to create attack method '{name}': {e}")
        return None
@@ -217,22 +261,29 @@ class UnifiedRegistry:

        if model_class:
            try:
-                # Extract parameters from configuration
-                model_name = config.get("model_name", name)
-                api_key = config.get("api_key", "")
-                base_url = config.get("base_url", None)
+                # Prefer factory method if provided by the model class.
+                # This removes provider-specific branching from the registry.
+                if hasattr(model_class, "from_config") and callable(
+                    getattr(model_class, "from_config")
+                ):
+                    # NOTE: models are instantiated via `models.*` which has its own config parsing.
+                    # We keep `_registry_name` injection for symmetry, but it is optional for models.
+                    cfg = dict(config or {})
+                    cfg.setdefault("_registry_name", name)
+                    return model_class.from_config(name=name, config=cfg)

-                # Pass appropriate parameters based on provider type
-                if provider in ["openai", "qwen", "vllm", "any"]:
-                    # OpenAI-compatible models need model_name, api_key, base_url
+                # ---- Backward-compatible fallback ----
+                cfg = dict(config or {})
+                cfg.setdefault("_registry_name", name)
+                model_name = cfg.get("model_name", name)
+                api_key = cfg.get("api_key", "")
+                base_url = cfg.get("base_url", None)
+                try:
                    return model_class(
                        model_name=model_name, api_key=api_key, base_url=base_url
                    )
-                elif provider in ["google", "anthropic", "doubao", "mistral"]:
-                    # Other models only need model_name and api_key
-                    return model_class(model_name=model_name, api_key=api_key)
-                else:
-                    # Default case
+                except TypeError:
+                    # Some providers don't take base_url
                    return model_class(model_name=model_name, api_key=api_key)
            except Exception as e:
                self.logger.error(
@@ -249,7 +300,9 @@ class UnifiedRegistry:
        defense_class = self.get_defense(name)
        if defense_class:
            try:
-                return defense_class(config=config)
+                cfg = dict(config or {})
+                cfg.setdefault("_registry_name", name)
+                return defense_class(config=cfg)
            except Exception as e:
                self.logger.error(f"Failed to create defense method '{name}': {e}")
        return None
@@ -261,42 +314,60 @@ class UnifiedRegistry:
        evaluator_class = self.get_evaluator(name)
        if evaluator_class:
            try:
-                return evaluator_class(config=config)
+                cfg = dict(config or {})
+                cfg.setdefault("_registry_name", name)
+                return evaluator_class(config=cfg)
            except Exception as e:
                self.logger.error(f"Failed to create evaluator '{name}': {e}")
        return None

    def list_attacks(self) -> List[str]:
        """List all registered attack methods"""
-        return sorted(self.attack_registry.keys())
+        names = set(self.attack_registry.keys())
+        names.update(self._get_lazy_mappings().get("attacks", {}).keys())
+        return sorted(names)

    def list_models(self) -> List[str]:
        """List all registered models"""
-        return sorted(self.model_registry.keys())
+        names = set(self.model_registry.keys())
+        names.update(self._get_lazy_mappings().get("models", {}).keys())
+        return sorted(names)

    def list_defenses(self) -> List[str]:
        """List all registered defense methods"""
-        return sorted(self.defense_registry.keys())
+        names = set(self.defense_registry.keys())
+        names.update(self._get_lazy_mappings().get("defenses", {}).keys())
+        return sorted(names)

    def list_evaluators(self) -> List[str]:
        """List all registered evaluators"""
-        return sorted(self.evaluator_registry.keys())
+        names = set(self.evaluator_registry.keys())
+        names.update(self._get_lazy_mappings().get("evaluators", {}).keys())
+        return sorted(names)

    def validate_attack(self, name: str) -> bool:
        """Validate if attack method exists"""
-        return name in self.attack_registry
+        if name in self.attack_registry:
+            return True
+        return name in self._get_lazy_mappings().get("attacks", {})

    def validate_model(self, name: str) -> bool:
        """Validate if model exists"""
-        return name in self.model_registry
+        if name in self.model_registry:
+            return True
+        return name in self._get_lazy_mappings().get("models", {})

    def validate_defense(self, name: str) -> bool:
        """Validate if defense method exists"""
-        return name in self.defense_registry
+        if name in self.defense_registry:
+            return True
+        return name in self._get_lazy_mappings().get("defenses", {})

    def validate_evaluator(self, name: str) -> bool:
        """Validate if evaluator exists"""
-        return name in self.evaluator_registry
+        if name in self.evaluator_registry:
+            return True
+        return name in self._get_lazy_mappings().get("evaluators", {})

    def get_component_summary(self) -> Dict[str, List[str]]:
        """Get component summary"""