From ed1e3ff43709e8c20a0f7a2617fb08711111b38c Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Fri, 7 Feb 2025 13:49:46 -0500 Subject: [PATCH 01/26] Adding AWS Bedrock Anthropic Claude target class --- .../aws_bedrock_claude_target.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 pyrit/prompt_target/aws_bedrock_claude_target.py diff --git a/pyrit/prompt_target/aws_bedrock_claude_target.py b/pyrit/prompt_target/aws_bedrock_claude_target.py new file mode 100644 index 000000000..ae7f8c8b4 --- /dev/null +++ b/pyrit/prompt_target/aws_bedrock_claude_target.py @@ -0,0 +1,105 @@ +import logging +import json +import boto3 +from typing import Optional +import asyncio + +from botocore.exceptions import ClientError + +from pyrit.models import PromptRequestResponse, construct_response_from_request +from pyrit.prompt_target import PromptTarget, limit_requests_per_minute + +logger = logging.getLogger(__name__) + +class AWSBedrockClaudeTarget(PromptTarget): + """ + This class initializes an AWS Bedrock target for any of the Anthropic Claude models. + Local AWS credentials (typically stored in ~/.aws) are used for authentication. + See the following for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html + + Parameters: + model_id (str): The model ID for target claude model + max_tokens (int): maximum number of tokens to generate + temperature (float, optional): The amount of randomness injected into the response. + top_p (float, optional): Use nucleus sampling + top_k (int, optional): Only sample from the top K options for each subsequent token + verify (bool, optional): whether or not to perform SSL certificate verification + """ + + def __init__( + self, + *, + model_id: str, + max_tokens: int, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + verify: bool = True, + max_requests_per_minute: Optional[int] = None, + ) -> None: + super().__init__(max_requests_per_minute=max_requests_per_minute) + + self._model_id = model_id + self._max_tokens = max_tokens + self._temperature = temperature + self._top_p = top_p + self._top_k = top_k + self._verify = verify + + @limit_requests_per_minute + async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: + + self._validate_request(prompt_request=prompt_request) + request = prompt_request.request_pieces[0] + + logger.info(f"Sending the following prompt to the prompt target: {request}") + + response = await self._complete_text_async(request.converted_value) + + response_entry = construct_response_from_request(request=request, response_text_pieces=[response]) + + return response_entry + + def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: + if len(prompt_request.request_pieces) != 1: + raise ValueError("This target only supports a single prompt request piece.") + + if prompt_request.request_pieces[0].converted_value_data_type != "text": + raise ValueError("This target only supports text prompt input.") + + async def _complete_text_async(self, text: str) -> str: + brt = boto3.client(service_name="bedrock-runtime", verify=self._verify) + + native_request = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": self._max_tokens, + "messages": [ + { + "role": "user", + "content": text + } + ] + } + + if self._temperature: + native_request['temperature'] = self._temperature + if self._top_p: + native_request['top_p'] = self._top_p + if self._top_k: + native_request['top_k'] = self._top_k + + request = json.dumps(native_request) + + try: + #response = brt.invoke_model(modelId=self._model_id, body=request) + response = await asyncio.to_thread(brt.invoke_model, modelId=self._model_id, body=request) + except (ClientError, Exception) as e: + print(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") + exit() + + model_response = json.loads(response["body"].read()) + + answer = model_response["content"][0]["text"] + + logger.info(f'Received the following response from the prompt target "{answer}"') + return answer From 63a9b2ef6ec25135a0386cee1132d16c2c07bef5 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Fri, 7 Feb 2025 13:50:46 -0500 Subject: [PATCH 02/26] Adding unit tests for AWSBedrockClaudeTarget class --- tests/unit/test_aws_bedrock_claude_target.py | 73 ++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 tests/unit/test_aws_bedrock_claude_target.py diff --git a/tests/unit/test_aws_bedrock_claude_target.py b/tests/unit/test_aws_bedrock_claude_target.py new file mode 100644 index 000000000..6baf855f4 --- /dev/null +++ b/tests/unit/test_aws_bedrock_claude_target.py @@ -0,0 +1,73 @@ +import pytest +import json +from unittest.mock import AsyncMock, patch, MagicMock + +from pyrit.models import PromptRequestResponse, PromptRequestPiece +from pyrit.prompt_target.aws_bedrock_claude_target import AWSBedrockClaudeTarget + +@pytest.fixture +def aws_target() -> AWSBedrockClaudeTarget: + return AWSBedrockClaudeTarget( + model_id="anthropic.claude-v2", + max_tokens=100, + temperature=0.7, + top_p=0.9, + top_k=50, + verify=True, + ) + +@pytest.fixture +def mock_prompt_request(): + request_piece = PromptRequestPiece( + role="user", + original_value="Hello, Claude!", + converted_value="Hello, how are you?" + ) + return PromptRequestResponse(request_pieces=[request_piece]) + +@pytest.mark.asyncio +async def test_send_prompt_async(aws_target, mock_prompt_request): + with patch("boto3.client", new_callable=MagicMock) as mock_boto: + mock_client = mock_boto.return_value + mock_client.invoke_model.return_value = { + "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "I'm good, thanks!"}]}))) + } + + response = await aws_target.send_prompt_async(prompt_request=mock_prompt_request) + + assert response.request_pieces[0].converted_value == "I'm good, thanks!" + +@pytest.mark.asyncio +async def test_validate_request_valid(aws_target, mock_prompt_request): + aws_target._validate_request(prompt_request=mock_prompt_request) + +@pytest.mark.asyncio +async def test_validate_request_invalid_multiple_pieces(aws_target): + request_pieces = [ + PromptRequestPiece(role="user", original_value="test", converted_value="Text 1", converted_value_data_type="text"), + PromptRequestPiece(role="user", original_value="test", converted_value="Text 2", converted_value_data_type="text") + ] + invalid_request = PromptRequestResponse(request_pieces=request_pieces) + + with pytest.raises(ValueError, match="This target only supports a single prompt request piece."): + aws_target._validate_request(prompt_request=invalid_request) + +@pytest.mark.asyncio +async def test_validate_request_invalid_data_type(aws_target): + request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="image_path")] + invalid_request = PromptRequestResponse(request_pieces=request_pieces) + + with pytest.raises(ValueError, match="This target only supports text prompt input."): + aws_target._validate_request(prompt_request=invalid_request) + +@pytest.mark.asyncio +async def test_complete_text_async(aws_target): + with patch("boto3.client", new_callable=MagicMock) as mock_boto: + mock_client = mock_boto.return_value + mock_client.invoke_model.return_value = { + "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "Test Response"}]}))) + } + + response = await aws_target._complete_text_async("Test Input") + + assert response == "Test Response" From 5de223dba765d3f7f3febd8d27f3f9b21bfe836d Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Mon, 10 Feb 2025 13:02:43 -0500 Subject: [PATCH 03/26] Add optional aws dependency (boto3) use of aws_bedrock_claude_target.py target class requires the boto3 library --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index ab409e447..3c15c4ac0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,6 +139,7 @@ all = [ "ollama>=0.4.4", "types-PyYAML>=6.0.12.9", ] +aws = ["boto3>=1.36.6"] [tool.pytest.ini_options] pythonpath = ["."] From ac87b282cb4ce0ff11eb03a17e080ddd808b84ff Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Mon, 10 Feb 2025 13:07:15 -0500 Subject: [PATCH 04/26] Update aws_bedrock_claude_target.py Deleting unnecessary commented out line --- pyrit/prompt_target/aws_bedrock_claude_target.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_target.py b/pyrit/prompt_target/aws_bedrock_claude_target.py index ae7f8c8b4..e6aa4a8e3 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_target.py @@ -91,7 +91,6 @@ async def _complete_text_async(self, text: str) -> str: request = json.dumps(native_request) try: - #response = brt.invoke_model(modelId=self._model_id, body=request) response = await asyncio.to_thread(brt.invoke_model, modelId=self._model_id, body=request) except (ClientError, Exception) as e: print(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") From 45785d65adbd397cca35b371043817e4b767a8a5 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Mon, 10 Feb 2025 13:36:01 -0500 Subject: [PATCH 05/26] Adding bedrock claude target class Adding support for multi-turn --- .../aws_bedrock_claude_chat_target.py | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 pyrit/prompt_target/aws_bedrock_claude_chat_target.py diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py new file mode 100644 index 000000000..536bb371f --- /dev/null +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -0,0 +1,154 @@ +import asyncio +import logging +import json +import boto3 +from typing import MutableSequence, Optional + +from botocore.exceptions import ClientError + +from pyrit.chat_message_normalizer import ChatMessageNop, ChatMessageNormalizer +from pyrit.common import net_utility +from pyrit.models import ChatMessage, PromptRequestPiece, PromptRequestResponse, construct_response_from_request, ChatMessageListDictContent +from pyrit.prompt_target import PromptChatTarget, limit_requests_per_minute + +logger = logging.getLogger(__name__) + +class AWSBedrockClaudeChatTarget(PromptChatTarget): + """ + This class initializes an AWS Bedrock target for any of the Anthropic Claude models. + Local AWS credentials (typically stored in ~/.aws) are used for authentication. + See the following for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html + + Parameters: + model_id (str): The model ID for target claude model + max_tokens (int): maximum number of tokens to generate + temperature (float, optional): The amount of randomness injected into the response. + top_p (float, optional): Use nucleus sampling + top_k (int, optional): Only sample from the top K options for each subsequent token + verify (bool, optional): whether or not to perform SSL certificate verification + """ + def __init__( + self, + *, + model_id: str, + max_tokens: int, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + verify: bool = True, + chat_message_normalizer: ChatMessageNormalizer = ChatMessageNop(), + max_requests_per_minute: Optional[int] = None, + ) -> None: + super().__init__(max_requests_per_minute=max_requests_per_minute) + + self._model_id = model_id + self._max_tokens = max_tokens + self._temperature = temperature + self._top_p = top_p + self._top_k = top_k + self._verify = verify + self.chat_message_normalizer = chat_message_normalizer + + self._system_prompt = '' + + @limit_requests_per_minute + async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: + + self._validate_request(prompt_request=prompt_request) + request_piece = prompt_request.request_pieces[0] + + prompt_req_res_entries = self._memory.get_conversation(conversation_id=request_piece.conversation_id) + prompt_req_res_entries.append(prompt_request) + + logger.info(f"Sending the following prompt to the prompt target: {prompt_request}") + + messages = self._build_chat_messages(prompt_req_res_entries) + + response = await self._complete_chat_async(messages=messages) + + response_entry = construct_response_from_request(request=request_piece, response_text_pieces=[response]) + + return response_entry + + def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: + if len(prompt_request.request_pieces) != 1: + raise ValueError("This target only supports a single prompt request piece.") + + if prompt_request.request_pieces[0].converted_value_data_type != "text": + raise ValueError("This target only supports text prompt input.") + + async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) -> str: + brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', verify=self._verify) + + native_request = self._construct_request_body(messages) + + request = json.dumps(native_request) + + try: + response = await asyncio.to_thread(brt.invoke_model, modelId=self._model_id, body=request) + except (ClientError, Exception) as e: + print(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") + exit() + + model_response = json.loads(response["body"].read()) + + answer = model_response["content"][0]["text"] + + logger.info(f'Received the following response from the prompt target "{answer}"') + return answer + + def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptRequestResponse] + ) -> list[ChatMessageListDictContent]: + chat_messages: list[ChatMessageListDictContent] = [] + for prompt_req_resp_entry in prompt_req_res_entries: + prompt_request_pieces = prompt_req_resp_entry.request_pieces + + content = [] + role = None + for prompt_request_piece in prompt_request_pieces: + role = prompt_request_piece.role + if role == "system": + #bedrock doesn't allow a message with role==system, but it does let you specify system role in a param + self._system_prompt = prompt_request_piece.converted_value + elif prompt_request_piece.converted_value_data_type == "text": + entry = {"type": "text", "text": prompt_request_piece.converted_value} + content.append(entry) + else: + raise ValueError( + f"Multimodal data type {prompt_request_piece.converted_value_data_type} is not yet supported." + ) + + if not role: + raise ValueError("No role could be determined from the prompt request pieces.") + + chat_message = ChatMessageListDictContent(role=role, content=content) + chat_messages.append(chat_message) + return chat_messages + + def _construct_request_body(self, messages_list: list[ChatMessageListDictContent]) -> dict: + content = [] + + for message in messages_list: + if message.role != "system": + entry = {"role": message.role, "content": message.content} + content.append(entry) + + data = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": self._max_tokens, + "system": self._system_prompt, + "messages": content + } + + if self._temperature: + data['temperature'] = self._temperature + if self._top_p: + data['top_p'] = self._top_p + if self._top_k: + data['top_k'] = self._top_k + + return(data) + + def is_json_response_supported(self) -> bool: + """Indicates that this target supports JSON response format.""" + return True From f7a876779118857adc73d13c2f8e0d15ad640764 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Mon, 10 Feb 2025 13:39:52 -0500 Subject: [PATCH 06/26] Update __init__.py for new target classes --- pyrit/prompt_target/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 6c7742d18..3ea370322 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -7,7 +7,8 @@ from pyrit.prompt_target.openai.openai_target import OpenAITarget from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget - +from pyrit.prompt_target.aws_bedrock_claude_target import AWSBedrockClaudeTarget +from pyrit.prompt_target.aws_bedrock_claude_chat_target import AWSBedrockClaudeChatTarget from pyrit.prompt_target.azure_blob_storage_target import AzureBlobStorageTarget from pyrit.prompt_target.azure_ml_chat_target import AzureMLChatTarget from pyrit.prompt_target.crucible_target import CrucibleTarget @@ -29,6 +30,8 @@ from pyrit.prompt_target.text_target import TextTarget __all__ = [ + "AWSBedrockClaudeTarget", + "AWSBedrockClaudeChatTarget", "AzureBlobStorageTarget", "AzureMLChatTarget", "CrucibleTarget", From 57252d0a18e473734f7c7c0d16c5a2ff71f94a7a Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Mon, 10 Feb 2025 14:17:50 -0500 Subject: [PATCH 07/26] Unit test for AWSBedrockClaudeChatTarget --- .../test_aws_bedrock_claude_chat_target.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 tests/unit/test_aws_bedrock_claude_chat_target.py diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py new file mode 100644 index 000000000..d7a8fe4a0 --- /dev/null +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -0,0 +1,73 @@ +import pytest +import json +from unittest.mock import AsyncMock, patch, MagicMock + +from pyrit.models import PromptRequestResponse, PromptRequestPiece, ChatMessageListDictContent +from pyrit.prompt_target.aws_bedrock_claude_chat_target import AWSBedrockClaudeChatTarget + +@pytest.fixture +def aws_target() -> AWSBedrockClaudeChatTarget: + return AWSBedrockClaudeChatTarget( + model_id="anthropic.claude-v2", + max_tokens=100, + temperature=0.7, + top_p=0.9, + top_k=50, + verify=True, + ) + +@pytest.fixture +def mock_prompt_request(): + request_piece = PromptRequestPiece( + role="user", + original_value="Hello, Claude!", + converted_value="Hello, how are you?" + ) + return PromptRequestResponse(request_pieces=[request_piece]) + +@pytest.mark.asyncio +async def test_send_prompt_async(aws_target, mock_prompt_request): + with patch("boto3.client", new_callable=MagicMock) as mock_boto: + mock_client = mock_boto.return_value + mock_client.invoke_model.return_value = { + "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "I'm good, thanks!"}]}))) + } + + response = await aws_target.send_prompt_async(prompt_request=mock_prompt_request) + + assert response.request_pieces[0].converted_value == "I'm good, thanks!" + +@pytest.mark.asyncio +async def test_validate_request_valid(aws_target, mock_prompt_request): + aws_target._validate_request(prompt_request=mock_prompt_request) + +@pytest.mark.asyncio +async def test_validate_request_invalid_multiple_pieces(aws_target): + request_pieces = [ + PromptRequestPiece(role="user", original_value="test", converted_value="Text 1", converted_value_data_type="text"), + PromptRequestPiece(role="user", original_value="test", converted_value="Text 2", converted_value_data_type="text") + ] + invalid_request = PromptRequestResponse(request_pieces=request_pieces) + + with pytest.raises(ValueError, match="This target only supports a single prompt request piece."): + aws_target._validate_request(prompt_request=invalid_request) + +@pytest.mark.asyncio +async def test_validate_request_invalid_data_type(aws_target): + request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="image_path")] + invalid_request = PromptRequestResponse(request_pieces=request_pieces) + + with pytest.raises(ValueError, match="This target only supports text prompt input."): + aws_target._validate_request(prompt_request=invalid_request) + +@pytest.mark.asyncio +async def test_complete_chat_async(aws_target): + with patch("boto3.client", new_callable=MagicMock) as mock_boto: + mock_client = mock_boto.return_value + mock_client.invoke_model.return_value = { + "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "Test Response"}]}))) + } + + response = await aws_target._complete_chat_async(messages=[ChatMessageListDictContent(role="user", content=[{"type":"text", "text":"Test input"}])]) + + assert response == "Test Response" From f254145654ea6e24946a3226d9d3bae4339dae7a Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:47:47 -0500 Subject: [PATCH 08/26] Delete pyrit/prompt_target/aws_bedrock_claude_target.py Replacing this with aws_bedrock_claude_chat_target which can handle multi-turn orchestration --- .../aws_bedrock_claude_target.py | 104 ------------------ 1 file changed, 104 deletions(-) delete mode 100644 pyrit/prompt_target/aws_bedrock_claude_target.py diff --git a/pyrit/prompt_target/aws_bedrock_claude_target.py b/pyrit/prompt_target/aws_bedrock_claude_target.py deleted file mode 100644 index e6aa4a8e3..000000000 --- a/pyrit/prompt_target/aws_bedrock_claude_target.py +++ /dev/null @@ -1,104 +0,0 @@ -import logging -import json -import boto3 -from typing import Optional -import asyncio - -from botocore.exceptions import ClientError - -from pyrit.models import PromptRequestResponse, construct_response_from_request -from pyrit.prompt_target import PromptTarget, limit_requests_per_minute - -logger = logging.getLogger(__name__) - -class AWSBedrockClaudeTarget(PromptTarget): - """ - This class initializes an AWS Bedrock target for any of the Anthropic Claude models. - Local AWS credentials (typically stored in ~/.aws) are used for authentication. - See the following for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html - - Parameters: - model_id (str): The model ID for target claude model - max_tokens (int): maximum number of tokens to generate - temperature (float, optional): The amount of randomness injected into the response. - top_p (float, optional): Use nucleus sampling - top_k (int, optional): Only sample from the top K options for each subsequent token - verify (bool, optional): whether or not to perform SSL certificate verification - """ - - def __init__( - self, - *, - model_id: str, - max_tokens: int, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - top_k: Optional[int] = None, - verify: bool = True, - max_requests_per_minute: Optional[int] = None, - ) -> None: - super().__init__(max_requests_per_minute=max_requests_per_minute) - - self._model_id = model_id - self._max_tokens = max_tokens - self._temperature = temperature - self._top_p = top_p - self._top_k = top_k - self._verify = verify - - @limit_requests_per_minute - async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: - - self._validate_request(prompt_request=prompt_request) - request = prompt_request.request_pieces[0] - - logger.info(f"Sending the following prompt to the prompt target: {request}") - - response = await self._complete_text_async(request.converted_value) - - response_entry = construct_response_from_request(request=request, response_text_pieces=[response]) - - return response_entry - - def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: - if len(prompt_request.request_pieces) != 1: - raise ValueError("This target only supports a single prompt request piece.") - - if prompt_request.request_pieces[0].converted_value_data_type != "text": - raise ValueError("This target only supports text prompt input.") - - async def _complete_text_async(self, text: str) -> str: - brt = boto3.client(service_name="bedrock-runtime", verify=self._verify) - - native_request = { - "anthropic_version": "bedrock-2023-05-31", - "max_tokens": self._max_tokens, - "messages": [ - { - "role": "user", - "content": text - } - ] - } - - if self._temperature: - native_request['temperature'] = self._temperature - if self._top_p: - native_request['top_p'] = self._top_p - if self._top_k: - native_request['top_k'] = self._top_k - - request = json.dumps(native_request) - - try: - response = await asyncio.to_thread(brt.invoke_model, modelId=self._model_id, body=request) - except (ClientError, Exception) as e: - print(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") - exit() - - model_response = json.loads(response["body"].read()) - - answer = model_response["content"][0]["text"] - - logger.info(f'Received the following response from the prompt target "{answer}"') - return answer From f0bc2bca60d48e855818f72ba7d3baf1e289fbab Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:48:24 -0500 Subject: [PATCH 09/26] Update __init__.py --- pyrit/prompt_target/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 3ea370322..5b598e984 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -7,7 +7,6 @@ from pyrit.prompt_target.openai.openai_target import OpenAITarget from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget -from pyrit.prompt_target.aws_bedrock_claude_target import AWSBedrockClaudeTarget from pyrit.prompt_target.aws_bedrock_claude_chat_target import AWSBedrockClaudeChatTarget from pyrit.prompt_target.azure_blob_storage_target import AzureBlobStorageTarget from pyrit.prompt_target.azure_ml_chat_target import AzureMLChatTarget @@ -30,7 +29,6 @@ from pyrit.prompt_target.text_target import TextTarget __all__ = [ - "AWSBedrockClaudeTarget", "AWSBedrockClaudeChatTarget", "AzureBlobStorageTarget", "AzureMLChatTarget", From 2ebb5193b64c2f3423dde08d6e6de8726799746b Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:48:52 -0500 Subject: [PATCH 10/26] Delete tests/unit/test_aws_bedrock_claude_target.py --- tests/unit/test_aws_bedrock_claude_target.py | 73 -------------------- 1 file changed, 73 deletions(-) delete mode 100644 tests/unit/test_aws_bedrock_claude_target.py diff --git a/tests/unit/test_aws_bedrock_claude_target.py b/tests/unit/test_aws_bedrock_claude_target.py deleted file mode 100644 index 6baf855f4..000000000 --- a/tests/unit/test_aws_bedrock_claude_target.py +++ /dev/null @@ -1,73 +0,0 @@ -import pytest -import json -from unittest.mock import AsyncMock, patch, MagicMock - -from pyrit.models import PromptRequestResponse, PromptRequestPiece -from pyrit.prompt_target.aws_bedrock_claude_target import AWSBedrockClaudeTarget - -@pytest.fixture -def aws_target() -> AWSBedrockClaudeTarget: - return AWSBedrockClaudeTarget( - model_id="anthropic.claude-v2", - max_tokens=100, - temperature=0.7, - top_p=0.9, - top_k=50, - verify=True, - ) - -@pytest.fixture -def mock_prompt_request(): - request_piece = PromptRequestPiece( - role="user", - original_value="Hello, Claude!", - converted_value="Hello, how are you?" - ) - return PromptRequestResponse(request_pieces=[request_piece]) - -@pytest.mark.asyncio -async def test_send_prompt_async(aws_target, mock_prompt_request): - with patch("boto3.client", new_callable=MagicMock) as mock_boto: - mock_client = mock_boto.return_value - mock_client.invoke_model.return_value = { - "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "I'm good, thanks!"}]}))) - } - - response = await aws_target.send_prompt_async(prompt_request=mock_prompt_request) - - assert response.request_pieces[0].converted_value == "I'm good, thanks!" - -@pytest.mark.asyncio -async def test_validate_request_valid(aws_target, mock_prompt_request): - aws_target._validate_request(prompt_request=mock_prompt_request) - -@pytest.mark.asyncio -async def test_validate_request_invalid_multiple_pieces(aws_target): - request_pieces = [ - PromptRequestPiece(role="user", original_value="test", converted_value="Text 1", converted_value_data_type="text"), - PromptRequestPiece(role="user", original_value="test", converted_value="Text 2", converted_value_data_type="text") - ] - invalid_request = PromptRequestResponse(request_pieces=request_pieces) - - with pytest.raises(ValueError, match="This target only supports a single prompt request piece."): - aws_target._validate_request(prompt_request=invalid_request) - -@pytest.mark.asyncio -async def test_validate_request_invalid_data_type(aws_target): - request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="image_path")] - invalid_request = PromptRequestResponse(request_pieces=request_pieces) - - with pytest.raises(ValueError, match="This target only supports text prompt input."): - aws_target._validate_request(prompt_request=invalid_request) - -@pytest.mark.asyncio -async def test_complete_text_async(aws_target): - with patch("boto3.client", new_callable=MagicMock) as mock_boto: - mock_client = mock_boto.return_value - mock_client.invoke_model.return_value = { - "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "Test Response"}]}))) - } - - response = await aws_target._complete_text_async("Test Input") - - assert response == "Test Response" From 258f287a7ec259c39134523dc485c1fa831115a1 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:49:53 -0500 Subject: [PATCH 11/26] Update aws_bedrock_claude_chat_target.py changed "verify" param to "enable_ssl_verification" --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 536bb371f..5a2befca6 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -25,7 +25,7 @@ class AWSBedrockClaudeChatTarget(PromptChatTarget): temperature (float, optional): The amount of randomness injected into the response. top_p (float, optional): Use nucleus sampling top_k (int, optional): Only sample from the top K options for each subsequent token - verify (bool, optional): whether or not to perform SSL certificate verification + enable_ssl_verification (bool, optional): whether or not to perform SSL certificate verification """ def __init__( self, @@ -35,7 +35,7 @@ def __init__( temperature: Optional[float] = None, top_p: Optional[float] = None, top_k: Optional[int] = None, - verify: bool = True, + enable_ssl_verification: bool = True, chat_message_normalizer: ChatMessageNormalizer = ChatMessageNop(), max_requests_per_minute: Optional[int] = None, ) -> None: @@ -46,7 +46,7 @@ def __init__( self._temperature = temperature self._top_p = top_p self._top_k = top_k - self._verify = verify + self._enable_ssl_verification = enable_ssl_verification self.chat_message_normalizer = chat_message_normalizer self._system_prompt = '' @@ -78,7 +78,7 @@ def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: raise ValueError("This target only supports text prompt input.") async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) -> str: - brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', verify=self._verify) + brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', enable_ssl_verification=self._enable_ssl_verification) native_request = self._construct_request_body(messages) From 408c3089da642cf9de542e780116dab24020637a Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:51:38 -0500 Subject: [PATCH 12/26] Update test_aws_bedrock_claude_chat_target.py --- tests/unit/test_aws_bedrock_claude_chat_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index d7a8fe4a0..cdb596269 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -13,7 +13,7 @@ def aws_target() -> AWSBedrockClaudeChatTarget: temperature=0.7, top_p=0.9, top_k=50, - verify=True, + enable_ssl_verification=True, ) @pytest.fixture From 5865ac690fa92ab16a80a9a32be166d4bc3f15d8 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:54:27 -0500 Subject: [PATCH 13/26] Update pyproject.toml --- pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3c15c4ac0..ffee4a012 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,8 @@ playwright = [ "ollama>=0.4.4" ] +aws = ["boto3>=1.36.6"] + all = [ "accelerate==0.34.2", "azureml-mlflow==1.57.0", @@ -138,8 +140,8 @@ all = [ "flask>=3.1.0", "ollama>=0.4.4", "types-PyYAML>=6.0.12.9", + "boto3>=1.36.6" ] -aws = ["boto3>=1.36.6"] [tool.pytest.ini_options] pythonpath = ["."] From 01addd1ec06978809c6784f90584e772df61f079 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:58:23 -0500 Subject: [PATCH 14/26] Update aws_bedrock_claude_chat_target.py --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 5a2befca6..d67f4d6a4 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -87,8 +87,7 @@ async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) try: response = await asyncio.to_thread(brt.invoke_model, modelId=self._model_id, body=request) except (ClientError, Exception) as e: - print(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") - exit() + raise ValueError(f"ERROR: Can't invoke '{self._model_id}'. Reason: {e}") model_response = json.loads(response["body"].read()) From 627396a3f6151e1bc2a1d7f7bc018c0b2b6edf41 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 10:59:04 -0500 Subject: [PATCH 15/26] Update pyrit/prompt_target/aws_bedrock_claude_chat_target.py Co-authored-by: Roman Lutz --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index d67f4d6a4..07265161c 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -107,7 +107,7 @@ def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptReq for prompt_request_piece in prompt_request_pieces: role = prompt_request_piece.role if role == "system": - #bedrock doesn't allow a message with role==system, but it does let you specify system role in a param + # Bedrock doesn't allow a message with role==system, but it does let you specify system role in a param self._system_prompt = prompt_request_piece.converted_value elif prompt_request_piece.converted_value_data_type == "text": entry = {"type": "text", "text": prompt_request_piece.converted_value} From 59dcb7efc269bddaebdce6b435b35735c0ee7a67 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 11:03:31 -0500 Subject: [PATCH 16/26] Update pyrit/prompt_target/aws_bedrock_claude_chat_target.py Co-authored-by: Roman Lutz --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 07265161c..8d5dfa80a 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -146,7 +146,7 @@ def _construct_request_body(self, messages_list: list[ChatMessageListDictContent if self._top_k: data['top_k'] = self._top_k - return(data) + return data def is_json_response_supported(self) -> bool: """Indicates that this target supports JSON response format.""" From b5c4924d4c1643f89c5345c0b4770ba64dfa6397 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 11:04:34 -0500 Subject: [PATCH 17/26] Update aws_bedrock_claude_chat_target.py --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 8d5dfa80a..3c7791732 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -150,4 +150,4 @@ def _construct_request_body(self, messages_list: list[ChatMessageListDictContent def is_json_response_supported(self) -> bool: """Indicates that this target supports JSON response format.""" - return True + return False From 5d8d7e0b6e30e35d4208e630c17dd8b2f3c84e58 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 16:22:06 -0500 Subject: [PATCH 18/26] Update aws_bedrock_claude_chat_target.py --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 3c7791732..8d9fd18a7 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -78,7 +78,7 @@ def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: raise ValueError("This target only supports text prompt input.") async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) -> str: - brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', enable_ssl_verification=self._enable_ssl_verification) + brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', verify=self._enable_ssl_verification) native_request = self._construct_request_body(messages) From 185bcffd75e0a3182399b516bf441076fd1e800a Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 16:48:20 -0500 Subject: [PATCH 19/26] Update aws_bedrock_claude_chat_target.py Added support for images --- .../aws_bedrock_claude_chat_target.py | 33 +++++++++++++++---- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 8d9fd18a7..c964dfa3f 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -3,6 +3,7 @@ import json import boto3 from typing import MutableSequence, Optional +import base64 from botocore.exceptions import ClientError @@ -51,6 +52,8 @@ def __init__( self._system_prompt = '' + self._valid_image_types = ['jpeg', 'png', 'webp', 'gif'] + @limit_requests_per_minute async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: @@ -62,7 +65,7 @@ async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> P logger.info(f"Sending the following prompt to the prompt target: {prompt_request}") - messages = self._build_chat_messages(prompt_req_res_entries) + messages = await self._build_chat_messages(prompt_req_res_entries) response = await self._complete_chat_async(messages=messages) @@ -71,11 +74,13 @@ async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> P return response_entry def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: - if len(prompt_request.request_pieces) != 1: - raise ValueError("This target only supports a single prompt request piece.") + converted_prompt_data_types = [ + request_piece.converted_value_data_type for request_piece in prompt_request.request_pieces + ] - if prompt_request.request_pieces[0].converted_value_data_type != "text": - raise ValueError("This target only supports text prompt input.") + for prompt_data_type in converted_prompt_data_types: + if prompt_data_type not in ["text", "image_path"]: + raise ValueError("This target only supports text and image_path.") async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) -> str: brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', verify=self._enable_ssl_verification) @@ -96,7 +101,12 @@ async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) logger.info(f'Received the following response from the prompt target "{answer}"') return answer - def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptRequestResponse] + def _convert_local_image_to_base64(self, image_path: str) -> str: + with open(image_path, 'rb') as image_file: + encoded_string = base64.b64encode(image_file.read()) + return encoded_string + + async def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptRequestResponse] ) -> list[ChatMessageListDictContent]: chat_messages: list[ChatMessageListDictContent] = [] for prompt_req_resp_entry in prompt_req_res_entries: @@ -112,6 +122,17 @@ def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptReq elif prompt_request_piece.converted_value_data_type == "text": entry = {"type": "text", "text": prompt_request_piece.converted_value} content.append(entry) + elif prompt_request_piece.converted_value_data_type == "image_path": + image_type = prompt_request_piece.converted_value.split('.')[-1] + if image_type not in self._valid_image_types: + raise ValueError(f"Image file {prompt_request_piece.converted_value} must have valid extension of .jpeg, .png, .webp, or .gif") + + data_base64_encoded = self._convert_local_image_to_base64( + prompt_request_piece.converted_value + ) + media_type = "image/"+image_type + entry = {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": data_base64_encoded.decode()}} + content.append(entry) else: raise ValueError( f"Multimodal data type {prompt_request_piece.converted_value_data_type} is not yet supported." From 2630b43468d6d017078051a426dd98937d01d344 Mon Sep 17 00:00:00 2001 From: kmarsh77 Date: Wed, 12 Feb 2025 16:53:04 -0500 Subject: [PATCH 20/26] Update test_aws_bedrock_claude_chat_target.py --- tests/unit/test_aws_bedrock_claude_chat_target.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index cdb596269..3644651d7 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -41,23 +41,13 @@ async def test_send_prompt_async(aws_target, mock_prompt_request): async def test_validate_request_valid(aws_target, mock_prompt_request): aws_target._validate_request(prompt_request=mock_prompt_request) -@pytest.mark.asyncio -async def test_validate_request_invalid_multiple_pieces(aws_target): - request_pieces = [ - PromptRequestPiece(role="user", original_value="test", converted_value="Text 1", converted_value_data_type="text"), - PromptRequestPiece(role="user", original_value="test", converted_value="Text 2", converted_value_data_type="text") - ] - invalid_request = PromptRequestResponse(request_pieces=request_pieces) - - with pytest.raises(ValueError, match="This target only supports a single prompt request piece."): - aws_target._validate_request(prompt_request=invalid_request) @pytest.mark.asyncio async def test_validate_request_invalid_data_type(aws_target): - request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="image_path")] + request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="video")] invalid_request = PromptRequestResponse(request_pieces=request_pieces) - with pytest.raises(ValueError, match="This target only supports text prompt input."): + with pytest.raises(ValueError, match="This target only supports text and image_path."): aws_target._validate_request(prompt_request=invalid_request) @pytest.mark.asyncio From 0e0e300f1cc8217841b46757ecfa857028dea806 Mon Sep 17 00:00:00 2001 From: Kassidy Marsh Date: Thu, 20 Feb 2025 16:45:17 -0500 Subject: [PATCH 21/26] Updates to address complaints from pre-commit hooks --- .../aws_bedrock_claude_chat_target.py | 69 ++++++++++++------- .../test_aws_bedrock_claude_chat_target.py | 47 +++++++++---- 2 files changed, 76 insertions(+), 40 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index c964dfa3f..cc6607b46 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -1,24 +1,32 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + import asyncio -import logging +import base64 import json -import boto3 +import logging from typing import MutableSequence, Optional -import base64 +import boto3 from botocore.exceptions import ClientError from pyrit.chat_message_normalizer import ChatMessageNop, ChatMessageNormalizer -from pyrit.common import net_utility -from pyrit.models import ChatMessage, PromptRequestPiece, PromptRequestResponse, construct_response_from_request, ChatMessageListDictContent +from pyrit.models import ( + ChatMessageListDictContent, + PromptRequestResponse, + construct_response_from_request, +) from pyrit.prompt_target import PromptChatTarget, limit_requests_per_minute logger = logging.getLogger(__name__) + class AWSBedrockClaudeChatTarget(PromptChatTarget): """ This class initializes an AWS Bedrock target for any of the Anthropic Claude models. Local AWS credentials (typically stored in ~/.aws) are used for authentication. - See the following for more information: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html + See the following for more information: + https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html Parameters: model_id (str): The model ID for target claude model @@ -28,6 +36,7 @@ class AWSBedrockClaudeChatTarget(PromptChatTarget): top_k (int, optional): Only sample from the top K options for each subsequent token enable_ssl_verification (bool, optional): whether or not to perform SSL certificate verification """ + def __init__( self, *, @@ -50,13 +59,13 @@ def __init__( self._enable_ssl_verification = enable_ssl_verification self.chat_message_normalizer = chat_message_normalizer - self._system_prompt = '' + self._system_prompt = "" - self._valid_image_types = ['jpeg', 'png', 'webp', 'gif'] + self._valid_image_types = ["jpeg", "png", "webp", "gif"] @limit_requests_per_minute async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: - + self._validate_request(prompt_request=prompt_request) request_piece = prompt_request.request_pieces[0] @@ -83,7 +92,9 @@ def _validate_request(self, *, prompt_request: PromptRequestResponse) -> None: raise ValueError("This target only supports text and image_path.") async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) -> str: - brt = boto3.client(service_name="bedrock-runtime", region_name='us-east-1', verify=self._enable_ssl_verification) + brt = boto3.client( + service_name="bedrock-runtime", region_name="us-east-1", verify=self._enable_ssl_verification + ) native_request = self._construct_request_body(messages) @@ -102,11 +113,12 @@ async def _complete_chat_async(self, messages: list[ChatMessageListDictContent]) return answer def _convert_local_image_to_base64(self, image_path: str) -> str: - with open(image_path, 'rb') as image_file: + with open(image_path, "rb") as image_file: encoded_string = base64.b64encode(image_file.read()) - return encoded_string + return encoded_string.decode() - async def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[PromptRequestResponse] + async def _build_chat_messages( + self, prompt_req_res_entries: MutableSequence[PromptRequestResponse] ) -> list[ChatMessageListDictContent]: chat_messages: list[ChatMessageListDictContent] = [] for prompt_req_resp_entry in prompt_req_res_entries: @@ -117,21 +129,26 @@ async def _build_chat_messages(self, prompt_req_res_entries: MutableSequence[Pro for prompt_request_piece in prompt_request_pieces: role = prompt_request_piece.role if role == "system": - # Bedrock doesn't allow a message with role==system, but it does let you specify system role in a param + # Bedrock doesn't allow a message with role==system, + # but it does let you specify system role in a param self._system_prompt = prompt_request_piece.converted_value elif prompt_request_piece.converted_value_data_type == "text": entry = {"type": "text", "text": prompt_request_piece.converted_value} content.append(entry) elif prompt_request_piece.converted_value_data_type == "image_path": - image_type = prompt_request_piece.converted_value.split('.')[-1] + image_type = prompt_request_piece.converted_value.split(".")[-1] if image_type not in self._valid_image_types: - raise ValueError(f"Image file {prompt_request_piece.converted_value} must have valid extension of .jpeg, .png, .webp, or .gif") - - data_base64_encoded = self._convert_local_image_to_base64( - prompt_request_piece.converted_value - ) - media_type = "image/"+image_type - entry = {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": data_base64_encoded.decode()}} + raise ValueError( + f"""Image file {prompt_request_piece.converted_value} must + have valid extension of .jpeg, .png, .webp, or .gif""" + ) + + data_base64_encoded = self._convert_local_image_to_base64(prompt_request_piece.converted_value) + media_type = "image/" + image_type + entry = { + "type": "image", + "source": {"type": "base64", "media_type": media_type, "data": data_base64_encoded}, + } content.append(entry) else: raise ValueError( @@ -157,15 +174,15 @@ def _construct_request_body(self, messages_list: list[ChatMessageListDictContent "anthropic_version": "bedrock-2023-05-31", "max_tokens": self._max_tokens, "system": self._system_prompt, - "messages": content + "messages": content, } if self._temperature: - data['temperature'] = self._temperature + data["temperature"] = self._temperature if self._top_p: - data['top_p'] = self._top_p + data["top_p"] = self._top_p if self._top_k: - data['top_k'] = self._top_k + data["top_k"] = self._top_k return data diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index 3644651d7..60e92836a 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -1,9 +1,20 @@ -import pytest +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + import json -from unittest.mock import AsyncMock, patch, MagicMock +from unittest.mock import MagicMock, patch + +import pytest + +from pyrit.models import ( + ChatMessageListDictContent, + PromptRequestPiece, + PromptRequestResponse, +) +from pyrit.prompt_target.aws_bedrock_claude_chat_target import ( + AWSBedrockClaudeChatTarget, +) -from pyrit.models import PromptRequestResponse, PromptRequestPiece, ChatMessageListDictContent -from pyrit.prompt_target.aws_bedrock_claude_chat_target import AWSBedrockClaudeChatTarget @pytest.fixture def aws_target() -> AWSBedrockClaudeChatTarget: @@ -16,15 +27,15 @@ def aws_target() -> AWSBedrockClaudeChatTarget: enable_ssl_verification=True, ) + @pytest.fixture def mock_prompt_request(): request_piece = PromptRequestPiece( - role="user", - original_value="Hello, Claude!", - converted_value="Hello, how are you?" + role="user", original_value="Hello, Claude!", converted_value="Hello, how are you?" ) return PromptRequestResponse(request_pieces=[request_piece]) + @pytest.mark.asyncio async def test_send_prompt_async(aws_target, mock_prompt_request): with patch("boto3.client", new_callable=MagicMock) as mock_boto: @@ -32,11 +43,12 @@ async def test_send_prompt_async(aws_target, mock_prompt_request): mock_client.invoke_model.return_value = { "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "I'm good, thanks!"}]}))) } - + response = await aws_target.send_prompt_async(prompt_request=mock_prompt_request) - + assert response.request_pieces[0].converted_value == "I'm good, thanks!" + @pytest.mark.asyncio async def test_validate_request_valid(aws_target, mock_prompt_request): aws_target._validate_request(prompt_request=mock_prompt_request) @@ -44,12 +56,17 @@ async def test_validate_request_valid(aws_target, mock_prompt_request): @pytest.mark.asyncio async def test_validate_request_invalid_data_type(aws_target): - request_pieces = [PromptRequestPiece(role="user", original_value="test", converted_value="ImageData", converted_value_data_type="video")] + request_pieces = [ + PromptRequestPiece( + role="user", original_value="test", converted_value="ImageData", converted_value_data_type="video" + ) + ] invalid_request = PromptRequestResponse(request_pieces=request_pieces) - + with pytest.raises(ValueError, match="This target only supports text and image_path."): aws_target._validate_request(prompt_request=invalid_request) + @pytest.mark.asyncio async def test_complete_chat_async(aws_target): with patch("boto3.client", new_callable=MagicMock) as mock_boto: @@ -57,7 +74,9 @@ async def test_complete_chat_async(aws_target): mock_client.invoke_model.return_value = { "body": MagicMock(read=MagicMock(return_value=json.dumps({"content": [{"text": "Test Response"}]}))) } - - response = await aws_target._complete_chat_async(messages=[ChatMessageListDictContent(role="user", content=[{"type":"text", "text":"Test input"}])]) - + + response = await aws_target._complete_chat_async( + messages=[ChatMessageListDictContent(role="user", content=[{"type": "text", "text": "Test input"}])] + ) + assert response == "Test Response" From 6d531d50a905a8265a78a099af937224284819f4 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Wed, 26 Feb 2025 15:04:09 -0800 Subject: [PATCH 22/26] Update pyrit/prompt_target/aws_bedrock_claude_chat_target.py --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index cc6607b46..ae8232ed1 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -147,7 +147,11 @@ async def _build_chat_messages( media_type = "image/" + image_type entry = { "type": "image", - "source": {"type": "base64", "media_type": media_type, "data": data_base64_encoded}, + "source": { + "type": "base64", + "media_type": media_type, + "data": data_base64_encoded, + }, # type: ignore } content.append(entry) else: From e7c3c5422b0e5741d74d792a7a4cd6f9a2d8a23a Mon Sep 17 00:00:00 2001 From: Kassidy Marsh Date: Thu, 27 Feb 2025 13:44:24 -0500 Subject: [PATCH 23/26] Adding exceptions for when boto3 isn't installed --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 10 +++++++--- tests/unit/test_aws_bedrock_claude_chat_target.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index cc6607b46..8fdf5b5be 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -7,9 +7,6 @@ import logging from typing import MutableSequence, Optional -import boto3 -from botocore.exceptions import ClientError - from pyrit.chat_message_normalizer import ChatMessageNop, ChatMessageNormalizer from pyrit.models import ( ChatMessageListDictContent, @@ -63,6 +60,13 @@ def __init__( self._valid_image_types = ["jpeg", "png", "webp", "gif"] + try: + import boto3 + from botocore.exceptions import ClientError + except ModuleNotFoundError as e: + logger.error("Could not import boto. You may need to install it via 'pip install pyrit[all]'") + raise e + @limit_requests_per_minute async def send_prompt_async(self, *, prompt_request: PromptRequestResponse) -> PromptRequestResponse: diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index 60e92836a..66d0df3e0 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -16,6 +16,15 @@ ) +def is_boto3_installed(): + try: + import boto3 + + return True + except ModuleNotFoundError: + return False + + @pytest.fixture def aws_target() -> AWSBedrockClaudeChatTarget: return AWSBedrockClaudeChatTarget( @@ -36,6 +45,7 @@ def mock_prompt_request(): return PromptRequestResponse(request_pieces=[request_piece]) +@pytest.mark.skipif(not is_boto3_installed(), reason="boto3 is not installed") @pytest.mark.asyncio async def test_send_prompt_async(aws_target, mock_prompt_request): with patch("boto3.client", new_callable=MagicMock) as mock_boto: From 8ddb59609050e753a4f6748c03e24b2144dba999 Mon Sep 17 00:00:00 2001 From: Kassidy Marsh Date: Thu, 27 Feb 2025 13:46:26 -0500 Subject: [PATCH 24/26] Adding exceptions for when boto3 isn't installed --- tests/unit/test_aws_bedrock_claude_chat_target.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index 66d0df3e0..4c83ff20a 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -77,6 +77,7 @@ async def test_validate_request_invalid_data_type(aws_target): aws_target._validate_request(prompt_request=invalid_request) +@pytest.mark.skipif(not is_boto3_installed(), reason="boto3 is not installed") @pytest.mark.asyncio async def test_complete_chat_async(aws_target): with patch("boto3.client", new_callable=MagicMock) as mock_boto: From b772a9cace4ed0f0f98580e5f8834f37763ce2ba Mon Sep 17 00:00:00 2001 From: Kassidy Marsh Date: Fri, 28 Feb 2025 11:45:29 -0500 Subject: [PATCH 25/26] Adding noqa statements to pass pre-commit checks --- pyrit/prompt_target/aws_bedrock_claude_chat_target.py | 10 +++++++--- tests/unit/test_aws_bedrock_claude_chat_target.py | 2 +- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py index 5e5da7f89..0e5a268f2 100644 --- a/pyrit/prompt_target/aws_bedrock_claude_chat_target.py +++ b/pyrit/prompt_target/aws_bedrock_claude_chat_target.py @@ -5,7 +5,7 @@ import base64 import json import logging -from typing import MutableSequence, Optional +from typing import TYPE_CHECKING, MutableSequence, Optional from pyrit.chat_message_normalizer import ChatMessageNop, ChatMessageNormalizer from pyrit.models import ( @@ -17,6 +17,10 @@ logger = logging.getLogger(__name__) +if TYPE_CHECKING: + import boto3 + from botocore.exceptions import ClientError + class AWSBedrockClaudeChatTarget(PromptChatTarget): """ @@ -61,8 +65,8 @@ def __init__( self._valid_image_types = ["jpeg", "png", "webp", "gif"] try: - import boto3 - from botocore.exceptions import ClientError + import boto3 # noqa: F401 + from botocore.exceptions import ClientError # noqa: F401 except ModuleNotFoundError as e: logger.error("Could not import boto. You may need to install it via 'pip install pyrit[all]'") raise e diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index 4c83ff20a..0f17e4bbf 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -18,7 +18,7 @@ def is_boto3_installed(): try: - import boto3 + import boto3 # noqa: F401 return True except ModuleNotFoundError: From d88919acf237f3e64b47bee43c2db5b179446078 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Fri, 28 Feb 2025 14:10:22 -0800 Subject: [PATCH 26/26] Update tests/unit/test_aws_bedrock_claude_chat_target.py --- tests/unit/test_aws_bedrock_claude_chat_target.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/test_aws_bedrock_claude_chat_target.py b/tests/unit/test_aws_bedrock_claude_chat_target.py index 0f17e4bbf..0d9a86630 100644 --- a/tests/unit/test_aws_bedrock_claude_chat_target.py +++ b/tests/unit/test_aws_bedrock_claude_chat_target.py @@ -59,11 +59,13 @@ async def test_send_prompt_async(aws_target, mock_prompt_request): assert response.request_pieces[0].converted_value == "I'm good, thanks!" +@pytest.mark.skipif(not is_boto3_installed(), reason="boto3 is not installed") @pytest.mark.asyncio async def test_validate_request_valid(aws_target, mock_prompt_request): aws_target._validate_request(prompt_request=mock_prompt_request) +@pytest.mark.skipif(not is_boto3_installed(), reason="boto3 is not installed") @pytest.mark.asyncio async def test_validate_request_invalid_data_type(aws_target): request_pieces = [