diff --git a/pyrit/message_normalizer/chat_message_normalizer.py b/pyrit/message_normalizer/chat_message_normalizer.py index 254fa913c5..c064435c7e 100644 --- a/pyrit/message_normalizer/chat_message_normalizer.py +++ b/pyrit/message_normalizer/chat_message_normalizer.py @@ -16,7 +16,7 @@ SystemMessageBehavior, apply_system_message_behavior_async, ) -from pyrit.models import ChatMessage, Message, MessagePiece +from pyrit.models import ChatMessage, ComponentIdentifier, Message, MessagePiece if TYPE_CHECKING: from pyrit.models.literals import ChatMessageRole @@ -60,6 +60,27 @@ def __init__( self.use_developer_role = use_developer_role self.system_message_behavior = system_message_behavior + def as_identifier_params(self) -> dict[str, Any]: + """ + Return the configuration parameters that affect ChatMessage normalization. + + Returns: + dict[str, Any]: Parameters including use_developer_role and system_message_behavior. + """ + return { + "use_developer_role": self.use_developer_role, + "system_message_behavior": self.system_message_behavior, + } + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the behavioral identity for this normalizer. + + Returns: + ComponentIdentifier: The frozen identity snapshot. + """ + return ComponentIdentifier.of(self, params=self.as_identifier_params()) + async def normalize_async(self, messages: list[Message]) -> list[ChatMessage]: """ Convert a list of Messages to a list of ChatMessages. diff --git a/pyrit/message_normalizer/history_squash_normalizer.py b/pyrit/message_normalizer/history_squash_normalizer.py index fedfda98d1..f0d25c4ced 100644 --- a/pyrit/message_normalizer/history_squash_normalizer.py +++ b/pyrit/message_normalizer/history_squash_normalizer.py @@ -4,6 +4,7 @@ import copy import logging import uuid +from typing import Any from pyrit.message_normalizer._helpers import ( format_message_piece_for_context, @@ -12,12 +13,12 @@ from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer -from pyrit.models import Message, MessagePiece +from pyrit.models import ComponentIdentifier, Identifiable, Message, MessagePiece logger = logging.getLogger(__name__) -class HistorySquashNormalizer(MessageListNormalizer[Message]): +class HistorySquashNormalizer(MessageListNormalizer[Message], Identifiable): """ Combine conversation history and the current request into one message. @@ -58,6 +59,31 @@ def __init__( self._message_normalizer = message_normalizer self._expected_history_message_count = expected_history_message_count + def as_identifier_params(self) -> dict[str, Any]: + """ + Return the configuration parameters that affect history squashing. + + Returns: + dict[str, Any]: Parameters including expected count and child formatter identifier. + """ + params: dict[str, Any] = {} + if self._message_normalizer is not None: + params["message_normalizer"] = self._message_normalizer.get_identifier().model_dump( + mode="json", exclude={"pyrit_version", "hash", "eval_hash"} + ) + if self._expected_history_message_count is not None: + params["expected_history_message_count"] = self._expected_history_message_count + return params + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the behavioral identity for this normalizer. + + Returns: + ComponentIdentifier: The frozen identity snapshot. + """ + return ComponentIdentifier.of(self, params=self.as_identifier_params()) + async def normalize_async(self, messages: list[Message]) -> list[Message]: """ Combine history and the current request into one target-facing message. diff --git a/pyrit/message_normalizer/json_schema_normalizer.py b/pyrit/message_normalizer/json_schema_normalizer.py index 4937696041..ffdc5c2e07 100644 --- a/pyrit/message_normalizer/json_schema_normalizer.py +++ b/pyrit/message_normalizer/json_schema_normalizer.py @@ -2,9 +2,18 @@ # Licensed under the MIT license. import json +from typing import Any from pyrit.message_normalizer.message_normalizer import MessageListNormalizer -from pyrit.models import JSON_SCHEMA_METADATA_KEY, JsonResponseConfig, JsonSchemaDefinition, Message, MessagePiece +from pyrit.models import ( + JSON_SCHEMA_METADATA_KEY, + ComponentIdentifier, + Identifiable, + JsonResponseConfig, + JsonSchemaDefinition, + Message, + MessagePiece, +) # Default text appended to a message piece when the target cannot natively # enforce a JSON schema. The exact phrasing mirrors the legacy in-prompt @@ -26,7 +35,7 @@ ) -class JsonSchemaNormalizer(MessageListNormalizer[Message]): +class JsonSchemaNormalizer(MessageListNormalizer[Message], Identifiable): """ Adapts JSON-schema metadata for targets that cannot enforce it natively. @@ -80,6 +89,26 @@ def __init__( raise ValueError("schema_instructions_template must contain a '{schema_json}' placeholder.") self._schema_instructions_template = schema_instructions_template + def as_identifier_params(self) -> dict[str, Any]: + """ + Return the configuration parameters that affect JSON schema normalization. + + Returns: + dict[str, Any]: Parameters including schema_instructions_template. + """ + return { + "schema_instructions_template": self._schema_instructions_template, + } + + def _build_identifier(self) -> ComponentIdentifier: + """ + Build the behavioral identity for this normalizer. + + Returns: + ComponentIdentifier: The frozen identity snapshot. + """ + return ComponentIdentifier.of(self, params=self.as_identifier_params()) + async def normalize_async(self, messages: list[Message]) -> list[Message]: """ Return messages adapted for a target that does not support JSON schemas. diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index c5b19fa259..3d3d8bded2 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -7,7 +7,7 @@ from typing import Any from pyrit.message_normalizer import MessageListNormalizer -from pyrit.models import Message +from pyrit.models import ComponentIdentifier, Identifiable, Message from pyrit.prompt_target.common.conversation_normalization_pipeline import ConversationNormalizationPipeline from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -194,11 +194,87 @@ def as_identifier_params(self) -> dict[str, Any]: # ``self._pipeline``. Attack-owned overrides are represented by the # attack identifier instead. "normalization_pipeline": [ - f"{type(normalizer).__module__}.{type(normalizer).__qualname__}" - for normalizer in self._pipeline.normalizers + self._serialize_normalizer(normalizer) for normalizer in self._pipeline.normalizers ], } + @classmethod + def _serialize_normalizer(cls, normalizer: Any) -> dict[str, Any]: + """ + Serialize a normalizer into a deterministic dictionary containing its type, parameters, and children. + + Args: + normalizer (Any): The message normalizer instance to serialize. + + Returns: + dict[str, Any]: Dictionary with 'type', 'params', and optionally 'children' keys. + """ + params: dict[str, Any] = {} + children: dict[str, Any] = {} + if isinstance(normalizer, Identifiable): + ident = normalizer.get_identifier() + if ident.params: + params = ident.params + if ident.children: + children = cls._serialize_identifier_children(ident.children) + elif hasattr(normalizer, "as_identifier_params"): + val = normalizer.as_identifier_params() + if isinstance(val, dict): + params = val + + result: dict[str, Any] = { + "type": f"{type(normalizer).__module__}.{type(normalizer).__qualname__}", + "params": params, + } + if children: + result["children"] = children + return result + + @classmethod + def _serialize_identifier(cls, ident: ComponentIdentifier) -> dict[str, Any]: + """ + Serialize a ComponentIdentifier into a deterministic dictionary. + + Args: + ident (ComponentIdentifier): The component identifier to serialize. + + Returns: + dict[str, Any]: Dictionary containing 'type', 'params', and optionally 'children'. + """ + result: dict[str, Any] = { + "type": f"{ident.class_module}.{ident.class_name}", + "params": ident.params, + } + if ident.children: + result["children"] = cls._serialize_identifier_children(ident.children) + return result + + @classmethod + def _serialize_identifier_children( + cls, children: Mapping[str, ComponentIdentifier | list[ComponentIdentifier] | Any] + ) -> dict[str, Any]: + """ + Recursively serialize a mapping of child identifiers. + + Args: + children (Mapping[str, ComponentIdentifier | list[ComponentIdentifier] | Any]): The mapping of children. + + Returns: + dict[str, Any]: Recursively serialized dictionary of child configurations. + """ + serialized: dict[str, Any] = {} + for key in sorted(children.keys()): + val = children[key] + if isinstance(val, list): + serialized[key] = [ + cls._serialize_identifier(c) if isinstance(c, ComponentIdentifier) else c for c in val + ] + elif isinstance(val, ComponentIdentifier): + serialized[key] = cls._serialize_identifier(val) + else: + serialized[key] = val + return serialized + @staticmethod def _capabilities_to_identifier_params(capabilities: TargetCapabilities) -> dict[str, Any]: """ diff --git a/tests/unit/message_normalizer/test_chat_message_normalizer.py b/tests/unit/message_normalizer/test_chat_message_normalizer.py index 6233abd36c..9ec4f8b79f 100644 --- a/tests/unit/message_normalizer/test_chat_message_normalizer.py +++ b/tests/unit/message_normalizer/test_chat_message_normalizer.py @@ -319,3 +319,29 @@ async def test_returns_list_of_dicts(self): assert isinstance(result[0], dict) assert result[0]["role"] == "user" assert result[0]["content"] == "Hello" + + +class TestChatMessageNormalizerIdentifier: + """Tests for ChatMessageNormalizer identifier and configuration parameters.""" + + def test_as_identifier_params_default(self): + normalizer = ChatMessageNormalizer() + params = normalizer.as_identifier_params() + assert params == { + "use_developer_role": False, + "system_message_behavior": "keep", + } + + def test_as_identifier_params_custom(self): + normalizer = ChatMessageNormalizer(use_developer_role=True, system_message_behavior="squash") + params = normalizer.as_identifier_params() + assert params == { + "use_developer_role": True, + "system_message_behavior": "squash", + } + + def test_get_identifier(self): + normalizer = ChatMessageNormalizer(use_developer_role=True, system_message_behavior="ignore") + ident = normalizer.get_identifier() + assert ident.params["use_developer_role"] is True + assert ident.params["system_message_behavior"] == "ignore" diff --git a/tests/unit/message_normalizer/test_history_squash_normalizer.py b/tests/unit/message_normalizer/test_history_squash_normalizer.py index ba90695b5a..5ece906d06 100644 --- a/tests/unit/message_normalizer/test_history_squash_normalizer.py +++ b/tests/unit/message_normalizer/test_history_squash_normalizer.py @@ -268,3 +268,20 @@ async def test_history_squash_propagates_last_message_metadata(): assert len(result) == 1 squashed_piece = result[0].message_pieces[0] assert squashed_piece.prompt_metadata == {"json_schema": schema, "scenario": "regression"} + + +def test_history_squash_as_identifier_params_empty(): + normalizer = HistorySquashNormalizer() + assert normalizer.as_identifier_params() == {} + + +def test_history_squash_as_identifier_params_with_config(): + from pyrit.message_normalizer import ConversationContextNormalizer + + child = ConversationContextNormalizer() + normalizer = HistorySquashNormalizer(message_normalizer=child, expected_history_message_count=3) + params = normalizer.as_identifier_params() + assert params["expected_history_message_count"] == 3 + assert "message_normalizer" in params + assert params["message_normalizer"]["class_name"] == "ConversationContextNormalizer" + assert normalizer.get_identifier().params == params diff --git a/tests/unit/message_normalizer/test_json_schema_normalizer.py b/tests/unit/message_normalizer/test_json_schema_normalizer.py index 9f9a83994a..d9e000355a 100644 --- a/tests/unit/message_normalizer/test_json_schema_normalizer.py +++ b/tests/unit/message_normalizer/test_json_schema_normalizer.py @@ -213,3 +213,10 @@ def test_default_template_is_exposed_as_class_attribute(self) -> None: """Callers can extend the default phrasing instead of redefining it.""" assert "{schema_json}" in JsonSchemaNormalizer.DEFAULT_SCHEMA_INSTRUCTIONS_TEMPLATE assert "### Response format" in JsonSchemaNormalizer.DEFAULT_SCHEMA_INSTRUCTIONS_TEMPLATE + + def test_get_identifier(self) -> None: + """Builds identifier incorporating custom template.""" + custom_template = "Format: {schema_json}" + normalizer = JsonSchemaNormalizer(schema_instructions_template=custom_template) + ident = normalizer.get_identifier() + assert ident.params["schema_instructions_template"] == custom_template diff --git a/tests/unit/prompt_target/target/test_target_configuration.py b/tests/unit/prompt_target/target/test_target_configuration.py index 9ed1d6cafe..3fddc55f20 100644 --- a/tests/unit/prompt_target/target/test_target_configuration.py +++ b/tests/unit/prompt_target/target/test_target_configuration.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest @@ -11,7 +12,7 @@ JsonSchemaNormalizer, MessageListNormalizer, ) -from pyrit.models import Message, MessagePiece +from pyrit.models import ComponentIdentifier, Identifiable, Message, MessagePiece from pyrit.models.literals import ChatMessageRole from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -453,3 +454,178 @@ def test_capabilities_to_identifier_params_differing_caps_produce_differing_para assert TargetConfiguration._capabilities_to_identifier_params( caps_a ) != TargetConfiguration._capabilities_to_identifier_params(caps_b) + + +def test_as_identifier_params_pipeline_distinguishes_normalizer_templates(): + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + config_a = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template="[TEMPLATE_A]\n{schema_json}") + }, + ) + config_b = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template="[TEMPLATE_B]\n{schema_json}") + }, + ) + + params_a = config_a.as_identifier_params() + params_b = config_b.as_identifier_params() + + assert params_a["normalization_pipeline"] != params_b["normalization_pipeline"] + assert params_a != params_b + + +def test_as_identifier_params_pipeline_matches_when_normalizer_templates_equal(): + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + template = "[CUSTOM_TEMPLATE]\n{schema_json}" + config_a = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template=template)}, + ) + config_b = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template=template)}, + ) + + params_a = config_a.as_identifier_params() + params_b = config_b.as_identifier_params() + + assert params_a["normalization_pipeline"] == params_b["normalization_pipeline"] + assert params_a == params_b + + +def test_serialize_normalizer_honors_identifiable_message_list_normalizer(): + class CustomIdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"custom_key": "custom_val"}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + norm = CustomIdentifiableListNormalizer() + serialized = TargetConfiguration._serialize_normalizer(norm) + assert serialized["params"] == {"custom_key": "custom_val"} + assert "CustomIdentifiableListNormalizer" in serialized["type"] + + +def test_serialize_normalizer_falls_back_to_as_identifier_params(): + class DuckTypedNormalizer: + def as_identifier_params(self) -> dict[str, Any]: + return {"duck_key": "duck_val"} + + norm = DuckTypedNormalizer() + serialized = TargetConfiguration._serialize_normalizer(norm) + assert serialized["params"] == {"duck_key": "duck_val"} + assert "DuckTypedNormalizer" in serialized["type"] + + +def test_serialize_normalizer_defaults_to_empty_params(): + class PlainListNormalizer(MessageListNormalizer[Message]): + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + norm = PlainListNormalizer() + serialized = TargetConfiguration._serialize_normalizer(norm) + assert serialized["params"] == {} + assert "PlainListNormalizer" in serialized["type"] + + +def test_serialize_normalizer_retains_child_identifier_configuration(): + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + child = CustomPrefixNormalizer(prefix="[CHILD]") + wrapper = CustomWrapperNormalizer(child) + serialized = TargetConfiguration._serialize_normalizer(wrapper) + + assert serialized["params"] == {} + assert "children" in serialized + assert "normalizer" in serialized["children"] + child_serialized = serialized["children"]["normalizer"] + assert child_serialized["params"] == {"prefix": "[CHILD]"} + assert "CustomPrefixNormalizer" in child_serialized["type"] + + +def test_serialize_normalizer_recursively_serializes_nested_children(): + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + leaf = CustomPrefixNormalizer(prefix="[LEAF]") + inner_wrapper = CustomWrapperNormalizer(leaf) + outer_wrapper = CustomWrapperNormalizer(inner_wrapper) + + serialized = TargetConfiguration._serialize_normalizer(outer_wrapper) + assert serialized["params"] == {} + assert "children" in serialized + + level1 = serialized["children"]["normalizer"] + assert level1["params"] == {} + assert "CustomWrapperNormalizer" in level1["type"] + assert "children" in level1 + + level2 = level1["children"]["normalizer"] + assert level2["params"] == {"prefix": "[LEAF]"} + assert "CustomPrefixNormalizer" in level2["type"] diff --git a/tests/unit/prompt_target/test_round_robin_target.py b/tests/unit/prompt_target/test_round_robin_target.py index ade87f676c..4df1f02586 100644 --- a/tests/unit/prompt_target/test_round_robin_target.py +++ b/tests/unit/prompt_target/test_round_robin_target.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. + import pytest from unit.mocks import MockPromptTarget @@ -497,3 +498,418 @@ def test_init_uses_model_name_fallback_for_underlying_model(): # Both resolve to "gpt-4o" — should not raise _validate_behavioral_consistency([t1, t2]) + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_rejects_mismatched_normalizer_templates(): + from pyrit.message_normalizer import JsonSchemaNormalizer + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template="[WRAPPER_A]\n{schema_json}") + }, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template="[WRAPPER_B]\n{schema_json}") + }, + ) + + with pytest.raises(ValueError, match="identical configurations"): + RoundRobinTarget(targets=[t1, t2]) + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_accepts_identical_custom_normalizer_templates(): + from pyrit.message_normalizer import JsonSchemaNormalizer + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + template = "[CUSTOM_SCHEMA]\n{schema_json}" + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template=template)}, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: JsonSchemaNormalizer(schema_instructions_template=template)}, + ) + + rr = RoundRobinTarget(targets=[t1, t2]) + assert rr.configuration.as_identifier_params() == t1.configuration.as_identifier_params() + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_rejects_custom_identifiable_normalizers_with_mismatched_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class CustomPrefixNormalizer(Identifiable, MessageListNormalizer[Message]): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: CustomPrefixNormalizer(prefix="[A]")}, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: CustomPrefixNormalizer(prefix="[B]")}, + ) + + with pytest.raises(ValueError, match="identical configurations"): + RoundRobinTarget(targets=[t1, t2]) + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_accepts_custom_identifiable_normalizers_with_matching_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class CustomPrefixNormalizer(Identifiable, MessageListNormalizer[Message]): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: CustomPrefixNormalizer(prefix="[MATCH]")}, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={CapabilityName.JSON_SCHEMA: CustomPrefixNormalizer(prefix="[MATCH]")}, + ) + + rr = RoundRobinTarget(targets=[t1, t2]) + assert rr.configuration.as_identifier_params() == t1.configuration.as_identifier_params() + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_rejects_delegating_normalizers_with_mismatched_child_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[A]")) + }, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[B]")) + }, + ) + + with pytest.raises(ValueError, match="identical configurations"): + RoundRobinTarget(targets=[t1, t2]) + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_accepts_delegating_normalizers_with_matching_child_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[MATCH]")) + }, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[MATCH]")) + }, + ) + + rr = RoundRobinTarget(targets=[t1, t2]) + assert rr.configuration.as_identifier_params() == t1.configuration.as_identifier_params() + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_rejects_multi_level_delegating_normalizers_with_mismatched_leaf_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer( + CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[A]")) + ) + }, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer( + CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[B]")) + ) + }, + ) + + with pytest.raises(ValueError, match="identical configurations"): + RoundRobinTarget(targets=[t1, t2]) + + +@pytest.mark.usefixtures("patch_central_database") +def test_init_accepts_multi_level_delegating_normalizers_with_matching_leaf_params(): + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import ComponentIdentifier, Identifiable + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + UnsupportedCapabilityBehavior, + ) + + class _IdentifiableListNormalizer(Identifiable, MessageListNormalizer[Message]): + pass + + class CustomPrefixNormalizer(_IdentifiableListNormalizer): + def __init__(self, prefix: str): + self._prefix = prefix + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of(self, params={"prefix": self._prefix}) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return messages + + class CustomWrapperNormalizer(_IdentifiableListNormalizer): + def __init__(self, normalizer: _IdentifiableListNormalizer): + self.normalizer = normalizer + + def _build_identifier(self) -> ComponentIdentifier: + return ComponentIdentifier.of( + self, + children={"normalizer": self.normalizer.get_identifier()}, + ) + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + return await self.normalizer.normalize_async(messages) + + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_schema=False, + ) + policy = CapabilityHandlingPolicy(behaviors={CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT}) + t1 = MockPromptTarget() + t2 = MockPromptTarget() + t1._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer( + CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[MATCH]")) + ) + }, + ) + t2._configuration = TargetConfiguration( + capabilities=caps, + policy=policy, + normalizer_overrides={ + CapabilityName.JSON_SCHEMA: CustomWrapperNormalizer( + CustomWrapperNormalizer(CustomPrefixNormalizer(prefix="[MATCH]")) + ) + }, + ) + + rr = RoundRobinTarget(targets=[t1, t2]) + assert rr.configuration.as_identifier_params() == t1.configuration.as_identifier_params()