diff --git a/pyrit/executor/attack/component/adversarial_conversation_manager.py b/pyrit/executor/attack/component/adversarial_conversation_manager.py index a8f766ac16..95e1fe0ea7 100644 --- a/pyrit/executor/attack/component/adversarial_conversation_manager.py +++ b/pyrit/executor/attack/component/adversarial_conversation_manager.py @@ -13,6 +13,7 @@ from uuid import uuid4 from pyrit.exceptions import ( + BadRequestException, ComponentRole, InvalidJsonException, execution_context, @@ -755,6 +756,7 @@ async def generate_adversarial_reply_async( AdversarialReply: ``next_message`` plus the parsed ``rationale`` / ``last_response_summary``. Raises: + BadRequestException: If the adversarial chat returns a blocked or otherwise errored response. ValueError: If no response is received from the adversarial chat. InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted. """ @@ -831,6 +833,7 @@ async def _send_and_parse_async( AdversarialReply: ``next_message`` plus the parsed ``rationale`` / ``last_response_summary``. Raises: + BadRequestException: If the adversarial chat returns a blocked or otherwise errored response. ValueError: If no response is received from the adversarial chat. InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted. """ @@ -853,6 +856,10 @@ async def _send_and_parse_async( schema = self._response_json_schema def _parse(response: Message) -> AdversarialReply: + if response.is_error(): + raise BadRequestException( + message=f"Adversarial chat returned a blocked or errored response: {response.get_value()}" + ) return _parse_adversarial_reply(response.get_value(), schema=schema) with execution_context( diff --git a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py index edddd5c34e..8fb7823eb3 100644 --- a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py @@ -1,13 +1,14 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import json import logging from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest -from pyrit.exceptions import InvalidJsonException +from pyrit.exceptions import BadRequestException, InvalidJsonException from pyrit.executor.attack.component.adversarial_conversation_manager import ( _BLOCKED_FEEDBACK_TEXT, _DEFAULT_ADVERSARIAL_SCHEMA_NAME, @@ -93,6 +94,18 @@ def _response_message(value: str = "target said hi", *, data_type: str = "text", return Message(message_pieces=[piece]) +def _blocked_adversarial_response(refusal: str = "I cannot assist with that request.") -> Message: + payload = json.dumps({"status_code": 200, "message": refusal}) + piece = MessagePiece( + role="assistant", + original_value=payload, + original_value_data_type="error", + response_error="blocked", + ) + piece.mark_as_structured_refusal(refusal=refusal) + return Message(message_pieces=[piece]) + + def _seed_message(value: str = "seed prompt") -> Message: return Message(message_pieces=[MessagePiece(role="user", original_value=value, original_value_data_type="text")]) @@ -549,6 +562,36 @@ async def test_invalid_reply_raises(self): with pytest.raises(InvalidJsonException): await manager.get_next_message_async(turn_index=1, last_response=_response_message()) + async def test_blocked_reply_does_not_retry_json_parsing(self) -> None: + normalizer = _normalizer(None) + normalizer.send_prompt_async.return_value = _blocked_adversarial_response() + manager = _manager( + adversarial_system_prompt=_system_prompt(schema=SCHEMA), + prompt_normalizer=normalizer, + ) + + with pytest.raises(BadRequestException, match="blocked or errored response"): + await manager.get_next_message_async(turn_index=1, last_response=_response_message()) + + normalizer.send_prompt_async.assert_awaited_once() + + async def test_malformed_non_error_reply_retries_then_succeeds(self) -> None: + normalizer = _normalizer(None) + normalizer.send_prompt_async.side_effect = [ + Message.from_prompt(prompt="totally not json", role="assistant"), + Message.from_prompt(prompt=VALID_JSON, role="assistant"), + ] + manager = _manager( + adversarial_system_prompt=_system_prompt(schema=SCHEMA), + prompt_normalizer=normalizer, + ) + + turn = await manager.get_next_message_async(turn_index=1, last_response=_response_message()) + + assert turn.reply is not None + assert turn.reply.next_message == "hello target" + assert normalizer.send_prompt_async.call_count == 2 + async def test_non_object_reply_retries_then_succeeds(self) -> None: normalizer = _normalizer(None) normalizer.send_prompt_async.side_effect = [