From a98d5cf785aa20376b08394efbd17262290d6671 Mon Sep 17 00:00:00 2001 From: behnam Date: Wed, 29 Apr 2026 14:56:26 -0700 Subject: [PATCH 1/2] Add deprecation warnings when Message Piece is constructed with labels --- pyrit/backend/mappers/attack_mappers.py | 23 +++++++++++++++---- pyrit/backend/services/attack_service.py | 20 ++++++++-------- .../attack/component/conversation_manager.py | 23 +++++++++++++++---- pyrit/executor/attack/multi_turn/crescendo.py | 2 +- .../executor/attack/multi_turn/red_teaming.py | 2 +- .../attack/multi_turn/tree_of_attacks.py | 2 +- pyrit/executor/promptgen/anecdoctor.py | 4 ++-- pyrit/models/message.py | 2 +- pyrit/models/message_piece.py | 12 ++++++++-- pyrit/prompt_normalizer/prompt_normalizer.py | 10 +++++++- .../common/prompt_chat_target.py | 11 +++++++-- .../openai/openai_response_target.py | 4 ++-- .../playwright_copilot_target.py | 2 +- pyrit/prompt_target/text_target.py | 2 +- pyrit/score/conversation_scorer.py | 2 +- 15 files changed, 87 insertions(+), 34 deletions(-) diff --git a/pyrit/backend/mappers/attack_mappers.py b/pyrit/backend/mappers/attack_mappers.py index 0245e2af12..67b9e7931b 100644 --- a/pyrit/backend/mappers/attack_mappers.py +++ b/pyrit/backend/mappers/attack_mappers.py @@ -33,6 +33,7 @@ Score, TargetInfo, ) +from pyrit.common.deprecation import print_deprecation_message from pyrit.models import AttackResult, ChatMessageRole, PromptDataType from pyrit.models import Message as PyritMessage from pyrit.models import MessagePiece as PyritMessagePiece @@ -409,7 +410,7 @@ def request_piece_to_pyrit_message_piece( role: ChatMessageRole, conversation_id: str, sequence: int, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> PyritMessagePiece: """ Convert a single request piece DTO to a PyRIT MessagePiece domain object. @@ -420,10 +421,17 @@ def request_piece_to_pyrit_message_piece( conversation_id: The conversation/attack ID. sequence: The message sequence number. labels: Optional labels to attach to the piece. + Deprecated: This parameter will be removed in a release 0.16.0. Returns: PyritMessagePiece domain object. """ + if labels is not None: + print_deprecation_message( + old_item="request_piece_to_pyrit_message_piece(..., labels=...)", + new_item="request_piece_to_pyrit_message_piece(...)", + removed_in="0.16.0", + ) metadata: Optional[dict[str, str | int]] = None if piece.prompt_metadata: metadata = dict(piece.prompt_metadata) @@ -439,7 +447,7 @@ def request_piece_to_pyrit_message_piece( conversation_id=conversation_id, sequence=sequence, prompt_metadata=metadata, - labels=labels or {}, + labels=labels or {}, # deprecated original_prompt_id=original_prompt_id, ) @@ -449,7 +457,7 @@ def request_to_pyrit_message( request: AddMessageRequest, conversation_id: str, sequence: int, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> PyritMessage: """ Build a PyRIT Message from an AddMessageRequest DTO. @@ -459,17 +467,24 @@ def request_to_pyrit_message( conversation_id: The conversation/attack ID. sequence: The message sequence number. labels: Optional labels to attach to each piece. + Deprecated: This parameter will be removed in a release 0.16.0. Returns: PyritMessage ready to send to the target. """ + if labels is not None: + print_deprecation_message( + old_item="request_to_pyrit_message(..., labels=...)", + new_item="request_to_pyrit_message(...)", + removed_in="0.16.0", + ) pieces = [ request_piece_to_pyrit_message_piece( piece=p, role=request.role, conversation_id=conversation_id, sequence=sequence, - labels=labels, + labels=labels, # deprecated ) for p in request.pieces ] diff --git a/pyrit/backend/services/attack_service.py b/pyrit/backend/services/attack_service.py index 41b98739aa..d7b3c5b023 100644 --- a/pyrit/backend/services/attack_service.py +++ b/pyrit/backend/services/attack_service.py @@ -340,7 +340,7 @@ async def create_attack_async(self, *, request: CreateAttackRequest) -> CreateAt await self._store_prepended_messages( conversation_id=conversation_id, prepended=request.prepended_conversation, - labels=labels, + labels=labels, # deprecated ) return CreateAttackResponse( @@ -614,14 +614,14 @@ async def add_message_async(self, *, attack_result_id: str, request: AddMessageR target_registry_name=target_registry_name, request=request, sequence=sequence, - labels=attack_labels, + labels=attack_labels, # deprecated ) else: await self._store_message_only_async( conversation_id=msg_conversation_id, request=request, sequence=sequence, - labels=attack_labels, + labels=attack_labels, # deprecated ) await self._update_attack_after_message_async(attack_result_id=attack_result_id, ar=ar, request=request) @@ -852,7 +852,7 @@ def _duplicate_conversation_up_to( # Apply optional overrides to the fresh pieces before persisting for piece in all_pieces: if labels_override is not None: - piece.labels = dict(labels_override) + piece.labels = dict(labels_override) # deprecated if remap_assistant_to_simulated and piece.api_role == "assistant": piece._role = "simulated_assistant" @@ -943,7 +943,7 @@ async def _store_prepended_messages( self, conversation_id: str, prepended: list[Any], - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> None: """Store prepended conversation messages in memory.""" for seq, msg in enumerate(prepended): @@ -953,7 +953,7 @@ async def _store_prepended_messages( role=msg.role, conversation_id=conversation_id, sequence=seq, - labels=labels, + labels=labels, # deprecated ) self._memory.add_message_pieces_to_memory(message_pieces=[piece]) @@ -964,7 +964,7 @@ async def _send_and_store_message_async( target_registry_name: str, request: AddMessageRequest, sequence: int, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> None: """Send message to target via normalizer and store response.""" target_obj = get_target_service().get_target_object(target_registry_name=target_registry_name) @@ -979,7 +979,7 @@ async def _send_and_store_message_async( request=request, conversation_id=conversation_id, sequence=sequence, - labels=labels, + labels=labels, # deprecated ) converter_configs = self._get_converter_configs(request) @@ -1000,7 +1000,7 @@ async def _store_message_only_async( conversation_id: str, request: AddMessageRequest, sequence: int, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> None: """Store message without sending (send=False).""" await self._persist_base64_pieces_async(request) @@ -1010,7 +1010,7 @@ async def _store_message_only_async( role=request.role, conversation_id=conversation_id, sequence=sequence, - labels=labels, + labels=labels, # deprecated ) self._memory.add_message_pieces_to_memory(message_pieces=[piece]) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 7a27cb5666..1104676830 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Optional +from pyrit.common.deprecation import print_deprecation_message from pyrit.common.utils import combine_dict from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, @@ -57,7 +58,7 @@ def get_adversarial_chat_messages( adversarial_chat_conversation_id: str, attack_identifier: ComponentIdentifier, adversarial_chat_target_identifier: ComponentIdentifier, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> list[Message]: """ Transform prepended conversation messages for adversarial chat with swapped roles. @@ -76,10 +77,17 @@ def get_adversarial_chat_messages( attack_identifier (ComponentIdentifier): Attack identifier to associate with messages. adversarial_chat_target_identifier (ComponentIdentifier): Target identifier for the adversarial chat. labels: Optional labels to associate with the messages. + Deprecated: This parameter will be removed in a release 0.16.0. Returns: List of transformed messages with swapped roles and new IDs. """ + if labels is not None: + print_deprecation_message( + old_item="get_adversarial_chat_messages(..., labels=...)", + new_item="get_adversarial_chat_messages(...)", + removed_in="0.16.0", + ) if not prepended_conversation: return [] @@ -110,7 +118,7 @@ def get_adversarial_chat_messages( conversation_id=adversarial_chat_conversation_id, attack_identifier=attack_identifier, prompt_target_identifier=adversarial_chat_target_identifier, - labels=labels, + labels=labels, # deprecated ) result.append(adversarial_piece.to_message()) @@ -245,7 +253,7 @@ def set_system_prompt( target: PromptChatTarget, conversation_id: str, system_prompt: str, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> None: """ Set or update the system prompt for a conversation. @@ -255,12 +263,19 @@ def set_system_prompt( conversation_id: Unique identifier for the conversation. system_prompt: The system prompt text. labels: Optional labels to associate with the system prompt. + Deprecated: This parameter will be removed in a release 0.16.0. """ + if labels is not None: + print_deprecation_message( + old_item="set_system_prompt(..., labels=...)", + new_item="set_system_prompt(...)", + removed_in="0.16.0", + ) target.set_system_prompt( system_prompt=system_prompt, conversation_id=conversation_id, attack_identifier=self._attack_identifier, - labels=labels, + labels=labels, # deprecated ) async def initialize_context_async( diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 4a180d5df3..9625a7c455 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -311,7 +311,7 @@ async def _setup_async(self, *, context: CrescendoAttackContext) -> None: system_prompt=system_prompt, conversation_id=context.session.adversarial_chat_conversation_id, attack_identifier=self.get_identifier(), - labels=context.memory_labels, + labels=context.memory_labels, # deprecated ) # Initialize backtrack count in context diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index a8778f664a..0dd784a0f0 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -256,7 +256,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: system_prompt=adversarial_system_prompt, conversation_id=context.session.adversarial_chat_conversation_id, attack_identifier=self.get_identifier(), - labels=context.memory_labels, + labels=context.memory_labels, # deprecated ) async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> AttackResult: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 7ea7f927b7..b0dfcc3ca1 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -982,7 +982,7 @@ async def _generate_first_turn_prompt_async(self, objective: str) -> str: system_prompt=system_prompt, conversation_id=self.adversarial_chat_conversation_id, attack_identifier=self._attack_id, - labels=self._memory_labels, + labels=self._memory_labels, # deprecated ) logger.debug(f"Node {self.node_id}: Using initial seed prompt for first turn") diff --git a/pyrit/executor/promptgen/anecdoctor.py b/pyrit/executor/promptgen/anecdoctor.py index f30dcc5c43..1d7450f5ce 100644 --- a/pyrit/executor/promptgen/anecdoctor.py +++ b/pyrit/executor/promptgen/anecdoctor.py @@ -213,7 +213,7 @@ async def _setup_async(self, *, context: AnecdoctorContext) -> None: system_prompt=system_prompt, conversation_id=context.conversation_id, attack_identifier=self.get_identifier(), - labels=context.memory_labels, + labels=context.memory_labels, # deprecated ) async def _perform_async(self, *, context: AnecdoctorContext) -> AnecdoctorResult: @@ -376,7 +376,7 @@ async def _extract_knowledge_graph_async(self, *, context: AnecdoctorContext) -> system_prompt=kg_system_prompt, conversation_id=kg_conversation_id, attack_identifier=self.get_identifier(), - labels=self._memory_labels, + labels=self._memory_labels, # deprecated ) # Format examples for knowledge graph extraction using few-shot format diff --git a/pyrit/models/message.py b/pyrit/models/message.py index 21e785fdc2..16a77efaab 100644 --- a/pyrit/models/message.py +++ b/pyrit/models/message.py @@ -555,7 +555,7 @@ def construct_response_from_request( role="assistant", original_value=resp_text, conversation_id=request.conversation_id, - labels=request.labels, + labels=request.labels, # deprecated prompt_target_identifier=request.prompt_target_identifier, attack_identifier=request.attack_identifier, original_value_data_type=response_type, diff --git a/pyrit/models/message_piece.py b/pyrit/models/message_piece.py index f4109786e9..52abdeb7d6 100644 --- a/pyrit/models/message_piece.py +++ b/pyrit/models/message_piece.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union, get_args from uuid import uuid4 +from pyrit.common.deprecation import print_deprecation_message from pyrit.identifiers.component_identifier import ComponentIdentifier from pyrit.models.literals import ChatMessageRole, PromptDataType, PromptResponseError @@ -69,6 +70,7 @@ def __init__( Defaults to None. sequence: The order of the conversation within a conversation_id. Defaults to -1. labels: The labels associated with the memory entry. Several can be standardized. Defaults to None. + Deprecated: This parameter will be removed in a release 0.16.0. prompt_metadata: The metadata associated with the prompt. This can be specific to any scenarios. Because memory is how components talk with each other, this can be component specific. e.g. the URI from a file uploaded to a blob store, or a document type you want to upload. @@ -117,6 +119,12 @@ def __init__( self.timestamp = timestamp.replace(tzinfo=timezone.utc) else: self.timestamp = timestamp + if labels is not None: + print_deprecation_message( + old_item="MessagePiece(..., labels=...)", + new_item="MessagePiece(...)", + removed_in="0.16.0", + ) self.labels = labels or {} self.prompt_metadata = prompt_metadata or {} @@ -212,7 +220,7 @@ def copy_lineage_from(self, source: MessagePiece) -> None: source: The piece whose lineage metadata is authoritative. """ self.conversation_id = source.conversation_id - self.labels = dict(source.labels) + self.labels = dict(source.labels) # deprecated self.attack_identifier = source.attack_identifier self.prompt_target_identifier = source.prompt_target_identifier self.prompt_metadata = dict(source.prompt_metadata) @@ -327,7 +335,7 @@ def to_dict(self) -> dict[str, object]: "conversation_id": self.conversation_id, "sequence": self.sequence, "timestamp": self.timestamp.isoformat() if self.timestamp else None, - "labels": self.labels, + "labels": self.labels, # deprecated "targeted_harm_categories": self.targeted_harm_categories if self.targeted_harm_categories else None, "prompt_metadata": self.prompt_metadata, "converter_identifiers": [conv.to_dict() for conv in self.converter_identifiers], diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index e8c96de607..528782dee6 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -8,6 +8,7 @@ from typing import Any, Optional from uuid import uuid4 +from pyrit.common.deprecation import print_deprecation_message from pyrit.exceptions import ( ComponentRole, EmptyResponseException, @@ -80,6 +81,7 @@ async def send_prompt_async( response_converter_configurations (list[PromptConverterConfiguration], optional): Configurations for converting the response. Defaults to an empty list. labels (Optional[dict[str, str]], optional): Labels associated with the request. Defaults to None. + Deprecated: This parameter will be removed in a release 0.16.0. attack_identifier (Optional[ComponentIdentifier], optional): Identifier for the attack. Defaults to None. @@ -90,6 +92,12 @@ async def send_prompt_async( Exception: If an error occurs during the request processing. ValueError: If the message pieces are not part of the same sequence. """ + if labels is not None: + print_deprecation_message( + old_item="send_prompt_async(..., labels=...)", + new_item="send_prompt_async(...)", + removed_in="0.16.0", + ) # Validates that the MessagePieces in the Message are part of the same sequence request_converter_configurations = request_converter_configurations or [] response_converter_configurations = response_converter_configurations or [] @@ -103,7 +111,7 @@ async def send_prompt_async( for piece in request.message_pieces: piece.conversation_id = conversation_id if labels: - piece.labels = labels + piece.labels = labels # deprecated piece.prompt_target_identifier = target.get_identifier() if attack_identifier: piece.attack_identifier = attack_identifier diff --git a/pyrit/prompt_target/common/prompt_chat_target.py b/pyrit/prompt_target/common/prompt_chat_target.py index ce1f254678..c58b94be40 100644 --- a/pyrit/prompt_target/common/prompt_chat_target.py +++ b/pyrit/prompt_target/common/prompt_chat_target.py @@ -3,6 +3,7 @@ from typing import Optional +from pyrit.common.deprecation import print_deprecation_message from pyrit.identifiers import ComponentIdentifier from pyrit.models import MessagePiece from pyrit.models.json_response_config import _JsonResponseConfig @@ -70,7 +71,7 @@ def set_system_prompt( system_prompt: str, conversation_id: str, attack_identifier: Optional[ComponentIdentifier] = None, - labels: Optional[dict[str, str]] = None, + labels: Optional[dict[str, str]] = None, # deprecated ) -> None: """ Set the system prompt for the prompt target. May be overridden by subclasses. @@ -78,6 +79,12 @@ def set_system_prompt( Raises: RuntimeError: If the conversation already exists. """ + if labels is not None: + print_deprecation_message( + old_item="set_system_prompt(..., labels=...)", + new_item="set_system_prompt(...)", + removed_in="0.16.0", + ) messages = self._memory.get_conversation(conversation_id=conversation_id) if messages: @@ -91,7 +98,7 @@ def set_system_prompt( converted_value=system_prompt, prompt_target_identifier=self.get_identifier(), attack_identifier=attack_identifier, - labels=labels, + labels=labels, # deprecated ).to_message() ) diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 60367107c5..1666ccc2c9 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -678,7 +678,7 @@ def _parse_response_output_section( role="assistant", original_value=piece_value, conversation_id=message_piece.conversation_id, - labels=message_piece.labels, + labels=message_piece.labels, # deprecated prompt_target_identifier=message_piece.prompt_target_identifier, attack_identifier=message_piece.attack_identifier, original_value_data_type=piece_type, @@ -786,7 +786,7 @@ def _make_tool_piece(self, output: dict[str, Any], call_id: str, *, reference_pi ), original_value_data_type="function_call_output", conversation_id=reference_piece.conversation_id, - labels={"call_id": call_id}, + labels={"call_id": call_id}, # deprecated prompt_target_identifier=reference_piece.prompt_target_identifier, attack_identifier=reference_piece.attack_identifier, ) diff --git a/pyrit/prompt_target/playwright_copilot_target.py b/pyrit/prompt_target/playwright_copilot_target.py index 58d30f088b..21439db4d3 100644 --- a/pyrit/prompt_target/playwright_copilot_target.py +++ b/pyrit/prompt_target/playwright_copilot_target.py @@ -243,7 +243,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me role="assistant", original_value=piece_data, conversation_id=request_piece.conversation_id, - labels=request_piece.labels, + labels=request_piece.labels, # deprecated prompt_target_identifier=request_piece.prompt_target_identifier, attack_identifier=request_piece.attack_identifier, original_value_data_type=piece_type, diff --git a/pyrit/prompt_target/text_target.py b/pyrit/prompt_target/text_target.py index 65d97d85d8..d6a4cb7835 100644 --- a/pyrit/prompt_target/text_target.py +++ b/pyrit/prompt_target/text_target.py @@ -87,7 +87,7 @@ def import_scores_from_csv(self, csv_file_path: Path) -> list[MessagePiece]: original_value_data_type=row.get("data_type", None), # type: ignore[arg-type] conversation_id=row.get("conversation_id", None), sequence=int(sequence_str) if sequence_str else 0, - labels=labels, + labels=labels, # deprecated response_error=row.get("response_error", None), # type: ignore[arg-type] prompt_target_identifier=self.get_identifier(), ) diff --git a/pyrit/score/conversation_scorer.py b/pyrit/score/conversation_scorer.py index e2b7a5ce13..7e5f03e2ab 100644 --- a/pyrit/score/conversation_scorer.py +++ b/pyrit/score/conversation_scorer.py @@ -84,7 +84,7 @@ async def _score_async(self, message: Message, *, objective: Optional[str] = Non converted_value=conversation_text, id=original_piece.id, conversation_id=original_piece.conversation_id, - labels=original_piece.labels, + labels=original_piece.labels, # deprecated prompt_target_identifier=original_piece.prompt_target_identifier, attack_identifier=original_piece.attack_identifier, original_value_data_type=original_piece.original_value_data_type, From ae76fa92edf3301bfc74763063daccc3341ec977 Mon Sep 17 00:00:00 2001 From: behnam Date: Wed, 29 Apr 2026 15:36:33 -0700 Subject: [PATCH 2/2] cc tests Co-authored-by: Copilot --- tests/unit/backend/test_mappers.py | 44 +++++++++++++++++++ .../component/test_conversation_manager.py | 38 +++++++++++++++- tests/unit/models/test_message_piece.py | 7 +++ .../test_prompt_normalizer.py | 17 +++++++ .../prompt_target/test_prompt_chat_target.py | 17 ++++++- 5 files changed, 121 insertions(+), 2 deletions(-) diff --git a/tests/unit/backend/test_mappers.py b/tests/unit/backend/test_mappers.py index 0f483b3f10..78a1937b44 100644 --- a/tests/unit/backend/test_mappers.py +++ b/tests/unit/backend/test_mappers.py @@ -691,6 +691,29 @@ def test_converts_request_to_domain(self) -> None: assert result.message_pieces[0].conversation_id == "conv-1" assert result.message_pieces[0].sequence == 0 + def test_labels_emit_deprecation_warning(self) -> None: + """Test that passing labels emits deprecation warning through mapper helper.""" + request = MagicMock() + request.role = "user" + piece = MagicMock() + piece.data_type = "text" + piece.original_value = "hello" + piece.converted_value = None + piece.prompt_metadata = None + piece.mime_type = None + piece.original_prompt_id = None + request.pieces = [piece] + + with patch("pyrit.backend.mappers.attack_mappers.print_deprecation_message") as mock_deprecation: + request_to_pyrit_message( + request=request, + conversation_id="conv-1", + sequence=0, + labels={"env": "prod"}, + ) + + assert mock_deprecation.call_count == 2 + class TestRequestPieceToPyritMessagePiece: """Tests for request_piece_to_pyrit_message_piece function.""" @@ -811,6 +834,27 @@ def test_labels_are_stamped_on_piece(self) -> None: assert result.labels == {"env": "prod"} + def test_labels_emit_deprecation_warning(self) -> None: + """Test that passing labels emits deprecation warning.""" + piece = MagicMock() + piece.data_type = "text" + piece.original_value = "hello" + piece.converted_value = None + piece.mime_type = None + piece.prompt_metadata = None + piece.original_prompt_id = None + + with patch("pyrit.backend.mappers.attack_mappers.print_deprecation_message") as mock_deprecation: + request_piece_to_pyrit_message_piece( + piece=piece, + role="user", + conversation_id="conv-1", + sequence=0, + labels={"env": "prod"}, + ) + + mock_deprecation.assert_called_once() + def test_labels_default_to_empty_dict(self) -> None: """Test that labels default to empty dict when not provided.""" piece = MagicMock() diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index c86e741e9c..49b301d533 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -19,7 +19,7 @@ import uuid from typing import Optional -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from unit.mocks import get_mock_scorer_identifier @@ -370,6 +370,24 @@ def test_applies_labels(self) -> None: assert result[0].get_piece().labels == labels + def test_labels_emit_deprecation_warning(self) -> None: + """Test that passing labels emits deprecation warning.""" + piece = MessagePiece(role="user", original_value="Message", conversation_id="original") + messages = [Message(message_pieces=[piece])] + + with patch( + "pyrit.executor.attack.component.conversation_manager.print_deprecation_message" + ) as mock_deprecation: + get_adversarial_chat_messages( + messages, + adversarial_chat_conversation_id="adversarial_conv", + attack_identifier=ComponentIdentifier(class_name="TestAttack", class_module="test_module"), + adversarial_chat_target_identifier=_mock_target_id("adversarial_target"), + labels={"env": "prod"}, + ) + + mock_deprecation.assert_called_once() + class TestBuildConversationContextStringAsync: """Tests for the build_conversation_context_string_async helper function.""" @@ -642,6 +660,24 @@ def test_set_system_prompt_without_labels( call_args = mock_chat_target.set_system_prompt.call_args assert call_args.kwargs["labels"] is None + def test_set_system_prompt_labels_emit_deprecation_warning( + self, attack_identifier: ComponentIdentifier, mock_chat_target: MagicMock + ) -> None: + """Test that passing labels emits deprecation warning.""" + manager = ConversationManager(attack_identifier=attack_identifier) + + with patch( + "pyrit.executor.attack.component.conversation_manager.print_deprecation_message" + ) as mock_deprecation: + manager.set_system_prompt( + target=mock_chat_target, + conversation_id=str(uuid.uuid4()), + system_prompt="You are a helpful assistant", + labels={"type": "system"}, + ) + + mock_deprecation.assert_called_once() + # ============================================================================= # Test Class: Initialize Context diff --git a/tests/unit/models/test_message_piece.py b/tests/unit/models/test_message_piece.py index 752347c71d..75c5799d7c 100644 --- a/tests/unit/models/test_message_piece.py +++ b/tests/unit/models/test_message_piece.py @@ -1121,6 +1121,13 @@ def test_targeted_harm_categories_none_no_warning(self): deprecation_msgs = [x for x in w if issubclass(x.category, DeprecationWarning)] assert not any("targeted_harm_categories" in str(m.message) for m in deprecation_msgs) + def test_labels_emits_deprecation_warning(self): + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + MessagePiece(role="user", original_value="Hello", labels={"env": "prod"}) + deprecation_msgs = [x for x in w if issubclass(x.category, DeprecationWarning)] + assert any("labels" in str(m.message) for m in deprecation_msgs) + class TestCopyLineageFrom: """Tests for MessagePiece.copy_lineage_from.""" diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index bc10d3323e..7fc30e8fa8 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -136,6 +136,23 @@ async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, s assert_message_piece_hashes_set(response) +@pytest.mark.asyncio +async def test_send_prompt_async_labels_emit_deprecation_warning(mock_memory_instance, seed_group): + prompt_target = MagicMock() + prompt_target.send_prompt_async = AsyncMock( + return_value=[MessagePiece(role="assistant", original_value="ok", conversation_id="conv-1").to_message()] + ) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + + normalizer = PromptNormalizer() + message = Message.from_prompt(prompt=seed_group.prompts[0].value, role="user") + + with patch("pyrit.prompt_normalizer.prompt_normalizer.print_deprecation_message") as mock_deprecation: + await normalizer.send_prompt_async(message=message, target=prompt_target, labels={"env": "prod"}) + + mock_deprecation.assert_called_once() + + @pytest.mark.asyncio async def test_send_prompt_async_empty_response_exception_handled(mock_memory_instance, seed_group): # Use MagicMock with send_prompt_async as AsyncMock to avoid coroutine warnings on other methods diff --git a/tests/unit/prompt_target/test_prompt_chat_target.py b/tests/unit/prompt_target/test_prompt_chat_target.py index 3ffcb6341a..cd2f5103f5 100644 --- a/tests/unit/prompt_target/test_prompt_chat_target.py +++ b/tests/unit/prompt_target/test_prompt_chat_target.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. import warnings -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from unit.mocks import MockPromptTarget, get_mock_attack_identifier @@ -52,6 +52,21 @@ def test_set_system_prompt_raises_if_conversation_exists(): # target that uses the real implementation. +@pytest.mark.usefixtures("patch_central_database") +def test_base_set_system_prompt_labels_emit_deprecation_warning(): + target = MockPromptTarget() + + with patch("pyrit.prompt_target.common.prompt_chat_target.print_deprecation_message") as mock_deprecation: + PromptChatTarget.set_system_prompt( + target, + system_prompt="You are a helpful assistant.", + conversation_id="conv-base-1", + labels={"key": "value"}, + ) + + mock_deprecation.assert_called_once() + + @pytest.mark.usefixtures("patch_central_database") def test_is_response_format_json_false_when_no_metadata(): target = MockPromptTarget()