FEAT new target class for AWS Bedrock Anthropic Claude models

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 = ["."]

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",

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

pyrit/prompt_target/aws_bedrock_claude_target.py

@@ -0,0 +1,104 @@
+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

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"