feat: Guardrails for Amazon Bedrock

backend/app/bedrock.py

@@ -12,7 +12,7 @@
 from app.repositories.models.conversation import MessageModel
 from app.repositories.models.custom_bot import GenerationParamsModel
 from app.routes.schemas.conversation import type_model_name
-from app.utils import convert_dict_keys_to_camel_case, get_bedrock_client
+from app.utils import convert_dict_keys_to_camel_case, get_bedrock_runtime_client
 
 logger = logging.getLogger(__name__)
 
@@ -24,7 +24,14 @@
     else DEFAULT_CLAUDE_GENERATION_CONFIG
 )
 
-client = get_bedrock_client()
+client = get_bedrock_runtime_client()
+
+
+class GuardrailConfig(TypedDict):
+    guardrailIdentifier: str
+    guardrailVersion: str
+    trace: str
+    streamProcessingMode: str
 
 
 class ConverseApiRequest(TypedDict):
@@ -34,6 +41,7 @@ class ConverseApiRequest(TypedDict):
     messages: list[dict]
     stream: bool
     system: list[dict]
+    guardrailConfig: GuardrailConfig | None
 
 
 class ConverseApiResponseMessageContent(TypedDict):
@@ -188,23 +196,151 @@ def compose_args_for_converse_api(
     return args
 
 
+def compose_args_for_converse_api_with_guardrail(
+    messages: list[MessageModel],
+    model: type_model_name,
+    instruction: str | None = None,
+    stream: bool = False,
+    generation_params: GenerationParamsModel | None = None,
+    grounding_source: dict | None = None,
+    guardrail: dict | None = None,
+) -> ConverseApiRequest:
+    """Compose arguments for Converse API.
+    Ref: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse_stream.html
+    """
+    arg_messages = []
+    for message in messages:
+        if message.role not in ["system", "instruction"]:
+            content_blocks = []
+            for c in message.content:
+                if c.content_type == "text":
+                    if message.role == "user":
+                        if guardrail and guardrail["grounding_threshold"] > 0:
+                            content_blocks.append({"guardContent": grounding_source})
+                        content_blocks.append(
+                            {
+                                "guardContent": {
+                                    "text": {"text": c.body, "qualifiers": ["query"]},
+                                }
+                            }
+                        )
+                    elif message.role == "assistant":
+                        content_blocks.append(
+                            {
+                                "text": (
+                                    {"content": c.body}
+                                    if isinstance(c.body, str)
+                                    else None
+                                )
+                            }
+                        )
+
+                elif c.content_type == "image":
+                    # e.g. "image/png" -> "png"
+                    format = (
+                        c.media_type.split("/")[1]
+                        if c.media_type is not None
+                        else "unknown"
+                    )
+
+                    content_blocks.append(
+                        {
+                            "image": {
+                                "format": format,
+                                # decode base64 encoded image
+                                "source": {"bytes": base64.b64decode(c.body)},
+                            }
+                        }
+                    )
+
+                elif c.content_type == "textAttachment":
+                    content_blocks.append(
+                        {
+                            "document": {
+                                "format": _get_converse_supported_format(
+                                    Path(c.file_name).suffix[
+                                        1:
+                                    ],  # e.g. "document.txt" -> "txt"
+                                ),
+                                "name": Path(
+                                    c.file_name
+                                ).stem,  # e.g. "document.txt" -> "document"
+                                # encode text attachment body
+                                "source": {"bytes": c.body.encode("utf-8")},
+                            }
+                        }
+                    )
+                else:
+                    raise NotImplementedError()
+            arg_messages.append({"role": message.role, "content": content_blocks})
+
+    inference_config = {
+        **DEFAULT_GENERATION_CONFIG,
+        **(
+            {
+                "maxTokens": generation_params.max_tokens,
+                "temperature": generation_params.temperature,
+                "topP": generation_params.top_p,
+                "stopSequences": generation_params.stop_sequences,
+            }
+            if generation_params
+            else {}
+        ),
+    }
+
+    # `top_k` is configured in `additional_model_request_fields` instead of `inference_config`
+    additional_model_request_fields = {"top_k": inference_config["top_k"]}
+    del inference_config["top_k"]
+
+    args: ConverseApiRequest = {
+        "inference_config": convert_dict_keys_to_camel_case(inference_config),
+        "additional_model_request_fields": additional_model_request_fields,
+        "model_id": get_model_id(model),
+        "messages": arg_messages,
+        "stream": stream,
+        "system": [],
+        "guardrailConfig": None,  # Initialize with None
+    }
+
+    if instruction:
+        args["system"].append({"text": instruction})
+
+    if guardrail and "guardrail_arn" in guardrail and "guardrail_version" in guardrail:
+        args["guardrailConfig"] = {  # Update the value
+            "guardrailIdentifier": guardrail["guardrail_arn"],
+            "guardrailVersion": guardrail["guardrail_version"],
+            "trace": "enabled",
+            "streamProcessingMode": "async",
+        }
+
+    return args
+
+
 def call_converse_api(args: ConverseApiRequest) -> ConverseApiResponse:
-    client = get_bedrock_client()
+    client = get_bedrock_runtime_client()
     messages = args["messages"]
     inference_config = args["inference_config"]
     additional_model_request_fields = args["additional_model_request_fields"]
     model_id = args["model_id"]
     system = args["system"]
 
-    response = client.converse(
-        modelId=model_id,
-        messages=messages,
-        inferenceConfig=inference_config,
-        system=system,
-        additionalModelRequestFields=additional_model_request_fields,
-    )
-
-    return response
+    if args and "guardrailConfig" in args:
+        return client.converse(
+            modelId=model_id,
+            messages=messages,
+            inferenceConfig=inference_config,
+            system=system,
+            additionalModelRequestFields=additional_model_request_fields,
+            guardrailConfig=args["guardrailConfig"],
+        )
+    else:
+        return client.converse(
+            modelId=model_id,
+            messages=messages,
+            inferenceConfig=inference_config,
+            system=system,
+            additionalModelRequestFields=additional_model_request_fields,
+        )
 
 
 def calculate_price(

backend/app/bot_remove.py

@@ -15,8 +15,9 @@
 
 DB_SECRETS_ARN = os.environ.get("DB_SECRETS_ARN", "")
 DOCUMENT_BUCKET = os.environ.get("DOCUMENT_BUCKET", "documents")
+BEDROCK_REGION = os.environ.get("BEDROCK_REGION", "us-east-1")
 
-s3_client = boto3.client("s3")
+s3_client = boto3.client("s3", BEDROCK_REGION)
 
 
 def delete_from_postgres(bot_id: str):

backend/app/guardrails.py

@@ -0,0 +1,41 @@
+from botocore.exceptions import ClientError
+from app.repositories.common import _get_table_client, compose_bot_id
+from app.utils import get_bedrock_client
+
+
+class BotNotFoundException(Exception):
+    pass
+
+
+class GuardrailArnRetrievalError(Exception):
+    pass
+
+
+def get_guardrail_arn(user_id: str, bot_id: str) -> str:
+    table = _get_table_client(user_id)
+    try:
+        response = table.get_item(
+            Key={"PK": user_id, "SK": compose_bot_id(user_id, bot_id)},
+            ConsistentRead=True,
+        )
+        guardrail_arn = (
+            response["Item"]["GuardrailsParams"]["guardrail_arn"]
+            if "Item" in response
+            and "GuardrailsParams" in response["Item"]
+            and "guardrail_arn" in response["Item"]["GuardrailsParams"]
+            else ""
+        )
+        print(guardrail_arn)
+        return guardrail_arn
+    except ClientError as e:
+        if e.response["Error"]["Code"] == "ConditionalCheckFailedException":
+            raise BotNotFoundException(f"Bot with id {bot_id} not found")
+        else:
+            raise GuardrailArnRetrievalError(
+                f"Error getting guardrail_arn for bot: {bot_id}: {e}"
+            )
+
+
+def delete_guardrail(guardrail_id):
+    client = get_bedrock_client()
+    client.delete_guardrail(guardrailIdentifier=guardrail_id)

backend/app/repositories/conversation.py

@@ -27,11 +27,13 @@
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
-s3_client = boto3.client("s3")
 
 THRESHOLD_LARGE_MESSAGE = 300 * 1024  # 300KB
 LARGE_MESSAGE_BUCKET = os.environ.get("LARGE_MESSAGE_BUCKET")
 
+BEDROCK_REGION = os.environ.get("BEDROCK_REGION", "us-east-1")
+s3_client = boto3.client("s3", BEDROCK_REGION)
+
 
 def store_conversation(
     user_id: str, conversation: ConversationModel, threshold=THRESHOLD_LARGE_MESSAGE