diff --git a/.env_example b/.env_example index 0d70a48cfd..1182b55e09 100644 --- a/.env_example +++ b/.env_example @@ -83,11 +83,22 @@ AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) -# Default endpoint goes here; specialized ones below +# Adversarial chat targets (used by scenario attack techniques, e.g. role-play, TAP). +# When multiple numbered endpoints are configured, adversarial_chat uses them in a round-robin. ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" +ADVERSARIAL_CHAT_UNDERLYING_MODEL="" + +ADVERSARIAL_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" +ADVERSARIAL_CHAT_KEY2="xxxxx" +ADVERSARIAL_CHAT_MODEL2="deployment-name" +ADVERSARIAL_CHAT_UNDERLYING_MODEL2="" + +ADVERSARIAL_CHAT_ENDPOINT3="https://xxxxx.openai.azure.com/openai/v1" +ADVERSARIAL_CHAT_KEY3="xxxxx" +ADVERSARIAL_CHAT_MODEL3="deployment-name" +ADVERSARIAL_CHAT_UNDERLYING_MODEL3="" ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" diff --git a/pyrit/setup/initializers/targets.py b/pyrit/setup/initializers/targets.py index 308366f734..35e491304f 100644 --- a/pyrit/setup/initializers/targets.py +++ b/pyrit/setup/initializers/targets.py @@ -198,6 +198,24 @@ class TargetConfig: underlying_model_var="ADVERSARIAL_CHAT_UNDERLYING_MODEL", temperature=1.2, ), + TargetConfig( + registry_name="adversarial_chat2", + target_class=OpenAIChatTarget, + endpoint_var="ADVERSARIAL_CHAT_ENDPOINT2", + key_var="ADVERSARIAL_CHAT_KEY2", + model_var="ADVERSARIAL_CHAT_MODEL2", + underlying_model_var="ADVERSARIAL_CHAT_UNDERLYING_MODEL2", + temperature=1.2, + ), + TargetConfig( + registry_name="adversarial_chat3", + target_class=OpenAIChatTarget, + endpoint_var="ADVERSARIAL_CHAT_ENDPOINT3", + key_var="ADVERSARIAL_CHAT_KEY3", + model_var="ADVERSARIAL_CHAT_MODEL3", + underlying_model_var="ADVERSARIAL_CHAT_UNDERLYING_MODEL3", + temperature=1.2, + ), TargetConfig( registry_name="adversarial_chat_singleturn", target_class=AzureMLChatTarget, @@ -539,6 +557,12 @@ class TargetInitializer(PyRITInitializer): await initializer.initialize_async() """ + _ADVERSARIAL_CHAT_NAMES: tuple[str, ...] = ( + "adversarial_chat", + "adversarial_chat2", + "adversarial_chat3", + ) + def __init__(self) -> None: """Initialize the TargetInitializer.""" super().__init__() @@ -605,6 +629,7 @@ async def initialize_async(self) -> None: continue self._register_target(config) + self._configure_adversarial_chat() if auto_group: self._auto_group_targets() @@ -688,6 +713,44 @@ def _register_target(self, config: TargetConfig) -> None: self._registered_names.append(config.registry_name) logger.info(f"Registered target: {config.registry_name}") + def _configure_adversarial_chat(self) -> None: + """ + Publish the configured adversarial endpoints under the canonical target name. + + Raises: + ValueError: If multiple adversarial targets have incompatible configurations. + """ + member_names = [name for name in self._ADVERSARIAL_CHAT_NAMES if name in self._registered_names] + self._registered_names = [name for name in self._registered_names if name not in self._ADVERSARIAL_CHAT_NAMES] + if not member_names: + return + + registry = TargetRegistry.get_registry_singleton() + targets = [target for name in member_names if (target := registry.instances.get(name)) is not None] + + if len(targets) == 1: + canonical_target = targets[0] + else: + try: + canonical_target = RoundRobinTarget(targets=targets) + except ValueError as ex: + raise ValueError(f"Adversarial chat round-robin targets are incompatible: {ex}") from ex + + if "adversarial_chat" in member_names: + primary = registry.instances.get("adversarial_chat") + if primary is not None: + registry.instances.register( + primary, + name="adversarial_chat_primary", + tags=[TargetInitializerTags.DEFAULT], + ) + + registry.instances.register( + canonical_target, + name="adversarial_chat", + tags=[TargetInitializerTags.DEFAULT], + ) + def _auto_group_targets(self) -> None: """ Automatically create round-robin groups from registered targets with diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 52141904e1..1a70801966 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -628,6 +628,134 @@ async def test_double_initialize_async_is_idempotent(self) -> None: assert first_default_count == second_default_count +@pytest.mark.usefixtures("patch_central_database") +class TestTargetInitializerAdversarialRoundRobin: + """Tests for explicit adversarial endpoint composition.""" + + SLOTS: tuple[tuple[str, str], ...] = ( + ("adversarial_chat", ""), + ("adversarial_chat2", "2"), + ("adversarial_chat3", "3"), + ) + + def setup_method(self) -> None: + """Reset the registry and adversarial endpoint environment.""" + TargetRegistry.reset_registry_singleton() + self._clear_env() + + def teardown_method(self) -> None: + """Reset the registry and adversarial endpoint environment.""" + TargetRegistry.reset_registry_singleton() + self._clear_env() + + def _clear_env(self) -> None: + for _, slot_suffix in self.SLOTS: + for variable in ("ENDPOINT", "KEY", "MODEL", "UNDERLYING_MODEL"): + os.environ.pop(f"ADVERSARIAL_CHAT_{variable}{slot_suffix}", None) + + def _set_slots(self, *slot_indexes: int, underlying_model: str = "grok-4.3") -> None: + for index in slot_indexes: + _, slot_suffix = self.SLOTS[index] + os.environ[f"ADVERSARIAL_CHAT_ENDPOINT{slot_suffix}"] = ( + f"https://grok-{index + 1}.openai.azure.com/openai/v1" + ) + os.environ[f"ADVERSARIAL_CHAT_KEY{slot_suffix}"] = f"key-{index + 1}" + os.environ[f"ADVERSARIAL_CHAT_MODEL{slot_suffix}"] = f"grok-deployment-{index + 1}" + os.environ[f"ADVERSARIAL_CHAT_UNDERLYING_MODEL{slot_suffix}"] = underlying_model + + @pytest.mark.parametrize("slot_index", [0, 1, 2]) + async def test_single_slot_publishes_direct_canonical_target(self, slot_index: int) -> None: + """Any single configured slot is directly available as ``adversarial_chat``.""" + self._set_slots(slot_index) + + await TargetInitializer().initialize_async() + + registry = TargetRegistry.get_registry_singleton() + member_name, _ = self.SLOTS[slot_index] + member = registry.instances.get(member_name) + assert isinstance(member, OpenAIChatTarget) + assert registry.instances.get("adversarial_chat") is member + if slot_index == 0: + assert registry.instances.get("adversarial_chat_primary") is member + + @pytest.mark.parametrize("slot_count", [2, 3]) + async def test_multiple_slots_publish_ordered_round_robin(self, slot_count: int) -> None: + """Two or three configured slots publish one ordered canonical round-robin.""" + from pyrit.prompt_target import RoundRobinTarget + + self._set_slots(*range(slot_count)) + + await TargetInitializer().initialize_async() + + registry = TargetRegistry.get_registry_singleton() + round_robin = registry.instances.get("adversarial_chat") + assert isinstance(round_robin, RoundRobinTarget) + assert len(round_robin.inner_targets) == slot_count + assert registry.instances.get("adversarial_chat_primary") is round_robin.inner_targets[0] + for index in range(1, slot_count): + member_name, _ = self.SLOTS[index] + assert registry.instances.get(member_name) is round_robin.inner_targets[index] + + async def test_noncontiguous_slots_publish_round_robin_without_inferred_duplicate(self) -> None: + """Secondary slots compose directly without producing a generic inferred group.""" + from pyrit.prompt_target import RoundRobinTarget + + self._set_slots(1, 2) + + await TargetInitializer().initialize_async() + + registry = TargetRegistry.get_registry_singleton() + round_robin = registry.instances.get("adversarial_chat") + assert isinstance(round_robin, RoundRobinTarget) + assert round_robin.inner_targets == [ + registry.instances.get("adversarial_chat2"), + registry.instances.get("adversarial_chat3"), + ] + assert registry.instances.get("adversarial_chat_primary") is None + assert registry.instances.get("OpenAIChatTarget_grok-4.3_temperature1.2_rr") is None + + async def test_explicit_round_robin_ignores_auto_group_setting(self) -> None: + """The configured adversarial pool is independent of inferred auto-grouping.""" + from pyrit.prompt_target import RoundRobinTarget + + self._set_slots(0, 1) + initializer = TargetInitializer() + initializer.params = {"tags": ["default"], "auto_group": False} + + await initializer.initialize_async() + + assert isinstance( + TargetRegistry.get_registry_singleton().instances.get("adversarial_chat"), + RoundRobinTarget, + ) + + async def test_canonical_and_member_targets_have_default_tag(self) -> None: + """The canonical pool and directly addressable members retain the default tag.""" + from pyrit.setup.initializers.targets import TargetInitializerTags + + self._set_slots(0, 1, 2) + await TargetInitializer().initialize_async() + + default_names = { + entry.name + for entry in TargetRegistry.get_registry_singleton().instances.get_by_tag(tag=TargetInitializerTags.DEFAULT) + } + assert { + "adversarial_chat", + "adversarial_chat_primary", + "adversarial_chat2", + "adversarial_chat3", + } <= default_names + + async def test_incompatible_members_fail_clearly(self) -> None: + """Different underlying models cannot form the configured adversarial pool.""" + self._set_slots(0) + self._set_slots(1, underlying_model="different-model") + + with pytest.raises(ValueError, match="Adversarial chat round-robin targets are incompatible"): + await TargetInitializer().initialize_async() + + @pytest.mark.usefixtures("patch_central_database") class TestTargetInitializerAutoGroup: """Tests for automatic round-robin grouping in TargetInitializer."""