Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion pyrit/message_normalizer/chat_message_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
30 changes: 28 additions & 2 deletions pyrit/message_normalizer/history_squash_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.

Expand Down Expand Up @@ -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.
Expand Down
33 changes: 31 additions & 2 deletions pyrit/message_normalizer/json_schema_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -26,7 +35,7 @@
)


class JsonSchemaNormalizer(MessageListNormalizer[Message]):
class JsonSchemaNormalizer(MessageListNormalizer[Message], Identifiable):
"""
Adapts JSON-schema metadata for targets that cannot enforce it natively.

Expand Down Expand Up @@ -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.
Expand Down
82 changes: 79 additions & 3 deletions pyrit/prompt_target/common/target_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"):
Comment thread
romanlutz marked this conversation as resolved.
Resolved
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]:
"""
Expand Down
26 changes: 26 additions & 0 deletions tests/unit/message_normalizer/test_chat_message_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
17 changes: 17 additions & 0 deletions tests/unit/message_normalizer/test_history_squash_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
7 changes: 7 additions & 0 deletions tests/unit/message_normalizer/test_json_schema_normalizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading
Loading