diff --git a/00_Prerequisites/bedrock_basics.ipynb b/00_Prerequisites/bedrock_basics.ipynb index df4a2b89..2f121e8f 100644 --- a/00_Prerequisites/bedrock_basics.ipynb +++ b/00_Prerequisites/bedrock_basics.ipynb @@ -611,7 +611,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "id": "bd800ef5-78f7-43d4-b73e-c37c7609cb40", "metadata": { "tags": [] @@ -627,7 +627,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "id": "8ef24e0c-dda8-4f5d-bbcb-db389e67e713", "metadata": { "tags": [] @@ -664,7 +664,8 @@ "from IPython.display import clear_output, display, display_markdown, Markdown\n", "\n", "body = json.dumps(messages_API_body)\n", - "modelId = \"anthropic.claude-v2\" # (Change this to try different model versions)\n", + "# modelId = \"anthropic.claude-v2\" # (Change this to try different model versions)\n", + "modelId = 'anthropic.claude-3-sonnet-20240229-v1:0' # change this to use a different version from the model provider\n", "accept = \"application/json\"\n", "contentType = \"application/json\"\n", "\n", @@ -1427,9 +1428,9 @@ ], "instance_type": "ml.t3.medium", "kernelspec": { - "display_name": "Python 3 (Data Science 3.0)", + "display_name": ".venv", "language": "python", - "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/sagemaker-data-science-310-v1" + "name": "python3" }, "language_info": { "codemirror_mode": { @@ -1441,7 +1442,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.6" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/07_Responsible_AI/responsible_ai-bedrock-guardrails_text_to_image_with_claude3.ipynb b/07_Responsible_AI/responsible_ai-bedrock-guardrails_text_to_image_with_claude3.ipynb new file mode 100644 index 00000000..9925ebc6 --- /dev/null +++ b/07_Responsible_AI/responsible_ai-bedrock-guardrails_text_to_image_with_claude3.ipynb @@ -0,0 +1,644 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Responsible AI \n", + "This notebook illustrates input and output guardrails as well as the AI safety mechanisms implemented in the base LLM. As an example, we will assume a simple text to image scneario. For input guardrails, we implement Bedrock Guardrails. If the input prompt appears to pass, we will forward along to the image generating model. In this example, we will use Titan Image Generator. Once the image is generated, we will further run the output through the Claude3 Haiku multi-modal LLM to ensure an additonal layer of safety is implemented before sending the resulting image back to the user. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Send in a text prompt to attempt to generate an Image\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Install Dependencies" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "# Install latest version of boto3\n", + "# !pip install boto3 s3transfer termcolor ipython --quiet\n", + "# !pip install matplotlib --quiet\n", + "import boto3\n", + "import pprint" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "1.34.95\n" + ] + } + ], + "source": [ + "# !pip install --upgrade --force-reinstall boto3 botocore\n", + "print(boto3.__version__)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "# Built-in libraries\n", + "import base64\n", + "import io\n", + "import json\n", + "import os\n", + "import sys\n", + "\n", + "# External dependencies\n", + "from PIL import Image\n", + "import botocore\n", + "\n", + "bedrock_rt = boto3.client('bedrock-runtime')\n", + "bedrock_client = boto3.client('bedrock')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create the Bedrock Guardrails" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "guardrail_config = {\n", + " 'name': 'no-warfare-2',\n", + " 'description': 'disallow input prompts for generating images related to military or warfare',\n", + " 'topicPolicyConfig': {\n", + " 'topicsConfig': [\n", + " {\n", + " 'name': 'deny military or warfare related prompts',\n", + " 'definition': 'Asking for generating images of mititary equipment like tanks, guns, mitilary drones, bullets, bombs, bunkers, shields etc are not allowed.',\n", + " 'examples': [\n", + " 'Generate an image depicting a military tank manufacturing plant',\n", + " 'Image showing bullets being inserted into a gun',\n", + " 'Image showing a mitiary bunker or base',\n", + " 'military drones',\n", + " ],\n", + " 'type': 'DENY'\n", + " }\n", + " ]\n", + " },\n", + " 'wordPolicyConfig': {\n", + " 'wordsConfig': [\n", + " {\n", + " 'text': 'tank'\n", + " },\n", + " {\n", + " 'text': 'bullet'\n", + " },\n", + " {\n", + " 'text': 'drone'\n", + " },\n", + " {\n", + " 'text': 'fighter aircrafts'\n", + " },\n", + " {\n", + " 'text': 'guns'\n", + " }\n", + " ],\n", + " },\n", + " 'blockedInputMessaging': 'I apologize, but I am not able to generate images related to the military. Please try again with another topic',\n", + " 'blockedOutputsMessaging': 'I apologize, but I am not able to generate images related to the military. Please try again with another topic',\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# create_response = bedrock_client.create_guardrail(**guardrail_config)\n", + "# pprint.pprint(create_response)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Let's now get all of our updates \n", + "# get_response = bedrock_client.get_guardrail(\n", + "# guardrailIdentifier=create_response['guardrailId'],\n", + "# guardrailVersion='DRAFT'\n", + "# )\n", + "# pprint.pprint(get_response)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# To list the DRAFT version of all your guardrails, don’t specify the guardrailIdentifier field. To list all versions of a guardrail, specify the ARN of the guardrail in the guardrailIdentifier field.\n", + "# list_guardrails_response = bedrock_client.list_guardrails(\n", + "# guardrailIdentifier=create_response['guardrailArn'],\n", + "# maxResults=5)\n", + "\n", + "# pprint.pprint(list_guardrails_response)" + ] + }, + { + "cell_type": "code", + "execution_count": 38, + "metadata": {}, + "outputs": [], + "source": [ + "prompt = \"a beautiful lake surrounded by trees with a mountain range at the distance\"\n", + "# prompt = \"military drones\"\n", + "negative_prompts = \"poorly rendered, poor background details, poorly drawn mountains, disfigured mountain features\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run the text promt through the input guardrails\n", + "\n", + "The Amazon Bedrock `InvokeModel` provides access to Amazon Titan Image Generator by setting the right model ID, and returns a JSON response including a [Base64 encoded string](https://en.wikipedia.org/wiki/Base64) that represents the (PNG) image.\n", + "\n", + "When making an `InvokeModel` request, we need to fill the `body` field with a JSON object that varies depending on the task (`taskType`) you wish to perform viz. text to image, image variation, inpainting or outpainting. The Amazon Titan models supports the following parameters:\n", + "* `cfgscale` - determines how much the final image reflects the prompt\n", + "* `seed` - a number used to initialize the generation, using the same seed with the same prompt + settings combination will produce the same results\n", + "* `numberOfImages` - the number of times the image is sampled and produced\n", + "* `quality` - determines the output image quality (`standard` or `premium`)\n", + "\n", + "> ☝️ For more information on available input parameters for the model, refer to the [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-image.html#model-parameters-titan-img-request-body) (Inference parameters > Amazon Titan image models > Model invocation request body fields).\n", + "\n", + "The cell below invokes the Amazon Titan Image Generator model through Amazon Bedrock to create an initial image:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### If safe, genreate the image" + ] + }, + { + "cell_type": "code", + "execution_count": 39, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Output: iVBORw0KGgoAAAANSUhEUgAABAAAAAQACAIAAADwf7zUAAEAAElEQVR4nFT9S7Jk27IrhgFwn7H5XkWm...\n" + ] + } + ], + "source": [ + "# Create payload\n", + "body = json.dumps(\n", + " {\n", + " \"taskType\": \"TEXT_IMAGE\",\n", + " \"textToImageParams\": {\n", + " \"text\": prompt, # Required\n", + " \"negativeText\": negative_prompts # Optional\n", + " },\n", + " \"imageGenerationConfig\": {\n", + " \"numberOfImages\": 1, # Range: 1 to 5 \n", + " \"quality\": \"standard\", # Options: standard or premium\n", + " \"height\": 1024, # Supported height list in the docs \n", + " \"width\": 1024, # Supported width list in the docs\n", + " \"cfgScale\": 7.5, # Range: 1.0 (exclusive) to 10.0\n", + " \"seed\": 42 # Range: 0 to 214783647\n", + " }\n", + " }\n", + ")\n", + "\n", + "# Make model request\n", + "response = bedrock_rt.invoke_model(\n", + " body=body,\n", + " modelId=\"amazon.titan-image-generator-v1\",\n", + " accept=\"application/json\", \n", + " contentType=\"application/json\"\n", + ")\n", + "\n", + "# Process the image\n", + "response_body = json.loads(response.get(\"body\").read())\n", + "img1_b64 = response_body[\"images\"][0]\n", + "\n", + "# Debug\n", + "print(f\"Output: {img1_b64[0:80]}...\")" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [], + "source": [ + "# Some utility functions\n", + "\n", + "def save_base64_image(img_base64, filename, directory=\"data/titan\"):\n", + " \"\"\"\n", + " Saves a base64 encoded image to the specified directory.\n", + "\n", + " Args:\n", + " img_base64 (str): Base64 encoded image data.\n", + " filename (str): Filename to save the image as.\n", + " directory (str, optional): Directory to save the image in. Defaults to \"data/titan\".\n", + "\n", + " Returns:\n", + " PIL.Image.Image: The saved image object.\n", + " \"\"\"\n", + " # Create the directory if it doesn't exist\n", + " os.makedirs(directory, exist_ok=True)\n", + "\n", + " # Decode and save the image\n", + " img_bytes = base64.decodebytes(bytes(img_base64, \"utf-8\"))\n", + " img = Image.open(io.BytesIO(img_bytes))\n", + " img.save(os.path.join(directory, filename))\n", + "\n", + " return img\n", + "\n", + "def get_base64_image(image_path):\n", + " \"\"\"\n", + " Generates the base64 encoded string representation of an image file.\n", + "\n", + " Args:\n", + " image_path (str): Path to the image file.\n", + "\n", + " Returns:\n", + " str: Base64 encoded string representation of the image.\n", + " \"\"\"\n", + " with open(image_path, \"rb\") as image_file:\n", + " encoded_string = base64.b64encode(image_file.read())\n", + "\n", + " return encoded_string.decode(\"utf-8\")\n", + "\n", + "def resize_and_save_base64(input_data, max_size, quality=90, is_base64=False):\n", + " \"\"\"\n", + " Resizes an image to fit within the specified maximum dimensions and saves it as a base64 encoded string.\n", + "\n", + " Args:\n", + " input_data (str or PIL.Image.Image): The image data as a base64 string or a PIL Image object.\n", + " max_size (tuple): The maximum width and height (max_width, max_height).\n", + " quality (int, optional): The quality of the saved image (0-100). Defaults to 90.\n", + " is_base64 (bool, optional): Whether the input_data is a base64 string or an image path. Defaults to False.\n", + "\n", + " Returns:\n", + " str: Base64 encoded string representation of the resized image.\n", + " \"\"\"\n", + " # Convert input data to a PIL Image object\n", + " if is_base64:\n", + " image_bytes = base64.decodebytes(bytes(input_data, \"utf-8\"))\n", + " image = Image.open(io.BytesIO(image_bytes))\n", + " else:\n", + " image = Image.open(input_data)\n", + "\n", + " # Resize the image if necessary\n", + " if image.size[0] > max_size[0] or image.size[1] > max_size[1]:\n", + " image.thumbnail(max_size, Image.Resampling.LANCZOS)\n", + "\n", + " # Save the resized image to a BytesIO object\n", + " image_data = io.BytesIO()\n", + " image.save(image_data, format='PNG', optimize=True, quality=quality)\n", + "\n", + " # Encode the image data as a base64 string\n", + " image_data.seek(0)\n", + " base64_data = base64.b64encode(image_data.getvalue()).decode('utf-8')\n", + "\n", + " return base64_data\n" + ] + }, + { + "cell_type": "code", + "execution_count": 35, + "metadata": {}, + "outputs": [], + "source": [ + "filename = \"image2.png\"\n", + "directory=\"data/titan\"" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Output: iVBORw0KGgoAAAANSUhEUgAABAAAAAQACAIAAADwf7zUAAEAAElEQVR4nEz9vbJlzbIrAEnKGrP3DRwi...\n" + ] + } + ], + "source": [ + "# Debug\n", + "print(f\"Output: {img1_b64[0:80]}...\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "filename = \"image2.png\"\n", + "directory=\"data/titan\"\n", + "\n", + "save_base64_image(img1_b64, filename, directory)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Base64 encoded image\n", + "# from fileinput import filename\n", + "base64_image_data = img1_b64\n", + "max_size = (1024, 1024)\n", + "resized_base64 = resize_and_save_base64(base64_image_data, max_size, is_base64=True)\n", + "\n", + "save_base64_image(resized_base64, filename, directory)\n", + "\n", + "# Image path\n", + "# image_path = os.path.join(directory, filename)\n", + "# max_size = (1024, 1024)\n", + "# resized_base64 = resize_and_save_base64(image_path, max_size)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "# image_path = os.path.join(directory, filename)\n", + "# img1_b64_temp2 = get_base64_image(image_path)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 44, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Output: iVBORw0KGgoAAAANSUhEUgAABAAAAAQACAIAAADwf7zUAAEAAElEQVR4nFT9S7Jk27IrhgFwn7H5XkWm...\n" + ] + } + ], + "source": [ + "# Debug\n", + "print(f\"Output: {img1_b64[0:80]}...\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Run the generated image through multi-modal model to check for safety" + ] + }, + { + "cell_type": "code", + "execution_count": 45, + "metadata": {}, + "outputs": [], + "source": [ + "import base64, re,json\n", + "import boto3\n", + "from termcolor import colored\n", + "import matplotlib.pyplot as plt\n", + "import matplotlib.image as mpimg\n", + "from matplotlib import rcParams\n", + "from IPython.display import display, Image\n", + "from PIL import Image\n", + "import io\n", + "import base64" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```python\n", + "messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"image\",\n", + " \"source\": {\n", + " \"type\": \"base64\",\n", + " \"media_type\": image1_media_type,\n", + " \"data\": image1_data,\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"Describe this image.\"\n", + " }\n", + " ],\n", + " }\n", + " ],\n", + "```\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 46, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_vision_answer(bedrock_rt:boto3.client,messages:list, model_id:str, claude_config:dict,system_prompt:str):\n", + " \"\"\"\n", + " Generates a vision answer using the specified model and configuration.\n", + " \n", + " Parameters:\n", + " - bedrock_rt (boto3.client): The Bedrock runtime client.\n", + " - messages (list): A list of messages.\n", + " - model_id (str): The ID of the model to use.\n", + " - claude_config (dict): The configuration for Claude.\n", + " - system_prompt (str): The system prompt.\n", + " \n", + " Returns:\n", + " - str: The formatted response.\n", + " \"\"\"\n", + " \n", + " body={'messages': [messages],**claude_config, \"system\": system_prompt}\n", + " \n", + " response = bedrock_rt.invoke_model(modelId=model_id, body=json.dumps(body)) \n", + " response = json.loads(response['body'].read().decode('utf-8'))\n", + " formated_response= response['content'][0]['text']\n", + " \n", + " return formated_response" + ] + }, + { + "cell_type": "code", + "execution_count": 47, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "CPU times: user 4 µs, sys: 1 µs, total: 5 µs\n", + "Wall time: 6.91 µs\n" + ] + } + ], + "source": [ + "%%time\n", + "\n", + "# Model id and claude config\n", + "# from https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html#model-ids-arns\n", + "model_id = \"anthropic.claude-3-sonnet-20240229-v1:0\"\n", + "claude_config = {\n", + " 'max_tokens': 6000, \n", + " 'temperature': 0, \n", + " 'anthropic_version': '', \n", + " 'top_p': 1, \n", + " 'stop_sequences': ['Human:']\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 48, + "metadata": {}, + "outputs": [], + "source": [ + "encoded_pngs = []\n", + "\n", + "# we can use the base64 version of the image directly since Titan generated it for us\n", + "# Else use the utility function provided above to encode the image\n", + "encoded_pngs.append(img1_b64)\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 49, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[32mSAFE\n", + "\n", + "The image depicts a beautiful natural landscape with a lake reflecting the snow-capped mountains in the background. There are no military or warfare-related elements present in the scene. The image shows trees, a body of water, and mountainous terrain, which are all allowed based on the provided guardrail configuration.\u001b[0m\n" + ] + } + ], + "source": [ + "content = [{\"type\": \"image\", \"source\": {\"type\": \"base64\", \"media_type\": \"image/png\", \"data\": encoded_png}} for encoded_png in encoded_pngs[:20]]\n", + "safety_prompt = \"\\n\" + json.dumps(guardrail_config) + \"\" + \"\\n\\nAs per the guardrail config provided, the input image is\"\n", + "content.append({\"type\": \"text\", \"text\": safety_prompt})\n", + "\n", + "# Create message with the prompt and the base64 encoded image\n", + "messages={\"role\": \"user\", \"content\": content }\n", + "system_prompt=\" You are a safety guard model. You will expect in a guardrail configuration json enclosed by tags. Reply UNSAFE if the image has any items denied in the guardrail_config. Else reply SAFE\" \n", + "\n", + "\n", + "# Generate answer\n", + "answer= generate_vision_answer(bedrock_rt, messages, model_id, claude_config, system_prompt)\n", + "print(colored(answer, \"green\"))" + ] + }, + { + "cell_type": "code", + "execution_count": 50, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "('\\n'\n", + " '{\"name\": \"no-warfare-2\", \"description\": \"disallow input '\n", + " 'prompts for generating images related to military or warfare\", '\n", + " '\"topicPolicyConfig\": {\"topicsConfig\": [{\"name\": \"deny military or warfare '\n", + " 'related prompts\", \"definition\": \"Asking for generating images of mititary '\n", + " 'equipment like tanks, guns, mitilary drones, bullets, bombs, bunkers, '\n", + " 'shields etc are not allowed.\", \"examples\": [\"Generate an image depicting a '\n", + " 'military tank manufacturing plant\", \"Image showing bullets being inserted '\n", + " 'into a gun\", \"Image showing a mitiary bunker or base\", \"military drones\"], '\n", + " '\"type\": \"DENY\"}]}, \"wordPolicyConfig\": {\"wordsConfig\": [{\"text\": \"tank\"}, '\n", + " '{\"text\": \"bullet\"}, {\"text\": \"drone\"}, {\"text\": \"fighter aircrafts\"}, '\n", + " '{\"text\": \"guns\"}]}, \"blockedInputMessaging\": \"I apologize, but I am not able '\n", + " 'to generate images related to the military. Please try again with another '\n", + " 'topic\", \"blockedOutputsMessaging\": \"I apologize, but I am not able to '\n", + " 'generate images related to the military. Please try again with another '\n", + " 'topic\"}\\n'\n", + " '\\n'\n", + " 'As per the guardrail config provided, the input image is')\n" + ] + } + ], + "source": [ + "# Sanity Check\n", + "pprint.pprint(safety_prompt)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 00000000..b03ad2c1 --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,21 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""General helper utilities the workshop notebooks""" +# Python Built-Ins: +from io import StringIO +import sys +import textwrap + + +def print_ww(*args, width: int = 100, **kwargs): + """Like print(), but wraps output to `width` characters (default 100)""" + buffer = StringIO() + try: + _stdout = sys.stdout + sys.stdout = buffer + print(*args, **kwargs) + output = buffer.getvalue() + finally: + sys.stdout = _stdout + for line in output.splitlines(): + print("\n".join(textwrap.wrap(line, width=width))) diff --git a/utils/bedrock.py b/utils/bedrock.py new file mode 100644 index 00000000..fb558af2 --- /dev/null +++ b/utils/bedrock.py @@ -0,0 +1,79 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Helper utilities for working with Amazon Bedrock from Python notebooks""" +# Python Built-Ins: +import os +from typing import Optional + +# External Dependencies: +import boto3 +from botocore.config import Config + + +def get_bedrock_client( + assumed_role: Optional[str] = None, + region: Optional[str] = None, + runtime: Optional[bool] = True, +): + """Create a boto3 client for Amazon Bedrock, with optional configuration overrides + + Parameters + ---------- + assumed_role : + Optional ARN of an AWS IAM role to assume for calling the Bedrock service. If not + specified, the current active credentials will be used. + region : + Optional name of the AWS Region in which the service should be called (e.g. "us-east-1"). + If not specified, AWS_REGION or AWS_DEFAULT_REGION environment variable will be used. + runtime : + Optional choice of getting different client to perform operations with the Amazon Bedrock service. + """ + if region is None: + target_region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION")) + else: + target_region = region + + print(f"Create new client\n Using region: {target_region}") + session_kwargs = {"region_name": target_region} + client_kwargs = {**session_kwargs} + + profile_name = os.environ.get("AWS_PROFILE") + if profile_name: + print(f" Using profile: {profile_name}") + session_kwargs["profile_name"] = profile_name + + retry_config = Config( + region_name=target_region, + retries={ + "max_attempts": 10, + "mode": "standard", + }, + ) + session = boto3.Session(**session_kwargs) + + if assumed_role: + print(f" Using role: {assumed_role}", end='') + sts = session.client("sts") + response = sts.assume_role( + RoleArn=str(assumed_role), + RoleSessionName="langchain-llm-1" + ) + print(" ... successful!") + client_kwargs["aws_access_key_id"] = response["Credentials"]["AccessKeyId"] + client_kwargs["aws_secret_access_key"] = response["Credentials"]["SecretAccessKey"] + client_kwargs["aws_session_token"] = response["Credentials"]["SessionToken"] + + if runtime: + service_name='bedrock-runtime' + else: + service_name='bedrock' + + bedrock_client = session.client( + service_name=service_name, + config=retry_config, + **client_kwargs + ) + + print("boto3 Bedrock client successfully created!") + print(bedrock_client._endpoint) + return bedrock_client