poc-to-prod/inference-profiles/inference-profile-basics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amazon Bedrock Application Inference Profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook demonstrates how organizations can implement, test, validate, and operationalize Amazon Bedrock application inference profiles. The aim is to provide a comprehensive understanding of how to manage and utilize application inference profiles effectively."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Overview\n",
    "\n",
    "Amazon Bedrock application inference profiles enable organizations to tag all Bedrock base foundation models with Cost Allocation Tags, making it possible to categorize usage by organizational taxonomies like cost centers, business units, teams, and applications. This scalable, programmatic approach to managing AI spend across multiple workloads reduces reliance on manual processes, lowers the risk of cost overruns, and ensures that critical applications receive priority. With enhanced visibility and control over AI-related expenses, organizations can optimize their GenAI investments and drive innovation more efficiently. This notebook demonstrates the creation of an application inference profile and its use in invoking models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Architecture \n",
    "![Architecture](./images/architecture.png)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The preceding architecture illustrates the configuration of Amazon Bedrock’s application inference profiles, which enable granular control over model usage and cost allocation. An application inference profile is associated with a specific model ID and AWS region, and can be created by copying a System-Defined Inference Profile template and adding custom tags for detailed tracking and management. While System-Defined Inference Profiles support cross-region inference, automatically routing requests across regions using a unified model identifier, these predefined profiles cannot be tagged, limiting visibility for managing costs and performance. To address this, Amazon Bedrock has introduced application inference profiles, a new capability that empowers organizations to optimize foundation model management across regions. With application inference profiles, organizations can create custom profiles with metadata tailored to tenants, such as teams, projects, and workloads, streamlining resource allocation and monitoring expenses across diverse AI applications."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use case \n",
    "\n",
    "This notebook contains implementation, test, and validation steps for various Amazon Bedrock application inference profile API functionalities, including:\n",
    "- Create Inference Profile\n",
    "- Get Inference Profile\n",
    "- List Inference Profiles\n",
    "- Invoke model with Inference Profile using Converse API\n",
    "- Invoke model with Inference Profile using ConverseStream API\n",
    "- Invoke model with Inference Profile using InvokeModel API\n",
    "- Invoke model with Inference Profile using InvokeModelWithResponseStream API\n",
    "- Tag Resource\n",
    "- List Tags For Resource\n",
    "- Untag Resource\n",
    "- Cleanup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "1. Ensure an AWS account has been created: [Link](https://docs.aws.amazon.com/accounts/latest/reference/manage-acct-creating.html)\n",
    "2. Ensure the user has permissions to access the correct models in Amazon Bedrock: [Link](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html)\n",
    "3. This Notebook was created in Amazon SageMaker in the us-west-2 region. If using this notebook in an outside environment ensure the AWS credentials are set correctly. If using in a different region, ensure the region variable is changed and that Amazon Bedrock application inference profiles are available in that region."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Code with Comments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up the Bedrock clients\n",
    "\n",
    "In the first cell, we import the necessary libraries and create a `boto3` client for Amazon Bedrock. \n",
    "We then call the `list_models` method to retrieve a list of available Amazon Bedrock base foundation models. This is the first step in understanding which models can be used for inference in your applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "import os\n",
    "import json\n",
    "import boto3\n",
    "\n",
    "print(f\"Boto3 Version: {boto3.__version__}\")\n",
    "\n",
    "# Get credentials and region information\n",
    "session = boto3.Session()\n",
    "credentials = session.get_credentials().get_frozen_credentials()\n",
    "region = session.region_name\n",
    "\n",
    "# Initialize Boto3 clients for Bedrock build time, runtime, agent, and agent runtime\n",
    "bedrock = boto3.client(\n",
    "    'bedrock',\n",
    "    region_name=region\n",
    ")\n",
    "bedrock_runtime = boto3.client(\n",
    "    'bedrock-runtime',\n",
    "    region_name=region\n",
    ")\n",
    "\n",
    "# Create a Bedrock client\n",
    "bedrock_client = boto3.client('bedrock')\n",
    "\n",
    "# Define filter parameters\n",
    "filter_params = {}\n",
    "# {\n",
    "#     'byProvider': 'your_provider_name',  # Replace with the desired model provider\n",
    "#     'byCustomizationType': 'FINE_TUNING',  # Options: FINE_TUNING or CONTINUED_PRE_TRAINING\n",
    "#     'byOutputModality': 'TEXT',  # Options: TEXT, IMAGE, or EMBEDDING\n",
    "#     'byInferenceType': 'ON_DEMAND'  # Options: ON_DEMAND or PROVISIONED\n",
    "# }\n",
    "\n",
    "# Displaying the available models in Amazon Bedrock with filters\n",
    "response = bedrock_client.list_foundation_models(**filter_params)\n",
    "model_summaries = response.get('modelSummaries', [])\n",
    "\n",
    "# Extract and print only the model ARNs\n",
    "model_arns = [model['modelArn'] for model in model_summaries]\n",
    "print(json.dumps(model_arns, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Inference Profile \n",
    "This cell calls the `create_inference_profile` method to create an application inference profile. By specifying a base foundation model ARN and custom tags, the application inference profile provides a way to organize and manage model configurations for the department's specific use case. The created inference profile ARN will be stored and can be used in place of default model IDs when making API calls to Bedrock, enabling more tailored interaction with the chosen model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def create_inference_profile(profile_name, model_arn, tags):\n",
    "    \"\"\"Create Inference Profile using base model ARN\"\"\"\n",
    "    response = bedrock.create_inference_profile(\n",
    "        inferenceProfileName=profile_name,\n",
    "        description=\"test\",\n",
    "        modelSource={'copyFrom': model_arn},\n",
    "        tags=tags\n",
    "    )\n",
    "    print(\"CreateInferenceProfile Response:\", response['ResponseMetadata']['HTTPStatusCode']),\n",
    "    print(f\"{response}\\n\")\n",
    "    return response\n",
    "\n",
    "# Create Inference Profile\n",
    "print(\"Testing CreateInferenceProfile...\")\n",
    "tags = [{'key': 'dept', 'value': 'claims'}]\n",
    "base_model_arn = \"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0\"\n",
    "claims_dept_claude_3_sonnet_profile = create_inference_profile(\"claims_dept_claude_3_sonnet_profile\", base_model_arn, tags)\n",
    "\n",
    "# Extracting the ARN and retrieving Inference Profile ID\n",
    "claims_dept_claude_3_sonnet_profile_arn = claims_dept_claude_3_sonnet_profile['inferenceProfileArn']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Inference Profile \n",
    "This cell retrieves metadata for the application inference profile created earlier by using its ARN. The `get_inference_profile` method takes the Inference Profile's ARN as input and calls the Bedrock API to fetch detailed metadata about the profile, such as configuration settings, model information, and associated tags. This retrieval allows you to verify and inspect the properties and setup of the profile, ensuring it meets the intended specifications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_inference_profile(inference_profile_arn):\n",
    "    \"\"\"Get Inference Profile by ARN\"\"\"\n",
    "    response = bedrock.get_inference_profile(\n",
    "        inferenceProfileIdentifier=inference_profile_arn\n",
    "    )\n",
    "    print(\"GetInferenceProfile Response:\", response['ResponseMetadata']['HTTPStatusCode']),\n",
    "    print(response)\n",
    "    return response\n",
    "\n",
    "print(\"Testing GetInferenceProfile...\")\n",
    "profile_response = get_inference_profile(claims_dept_claude_3_sonnet_profile_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### List Inference Profiles\n",
    "This cell utilizes the `list_inference_profiles` method to retrieve and display all inference profiles filtered by type: `SYSTEM_DEFINED` or `APPLICATION`. The response provides an overview of each profile's details, allowing you to view all configured profiles and their metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_inference_profiles():\n",
    "    \"\"\"List Inference Profiles filtered by type\"\"\"\n",
    "    response = bedrock.list_inference_profiles(\n",
    "        typeEquals=\"APPLICATION\"  # Filter for APPLICATION or SYSTEM_DEFINED\n",
    "    )\n",
    "    print(\"ListInferenceProfiles Response:\", response['ResponseMetadata']['HTTPStatusCode'])\n",
    "    print(response)\n",
    "    return response\n",
    "\n",
    "# List Inference Profiles\n",
    "print(\"Testing LisInferenceProfiles...\")\n",
    "profile_list_response = list_inference_profiles()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoke model with the application inference profile using Converse API \n",
    "The `Converse` API in Amazon Bedrock is a unified interface designed for engaging with large language models (LLMs), supporting features like chat history, automated prompt formatting specific to each model, and simplified model testing or swapping. This flexibility allows for easy substitution of Inference Profiles, regardless of the LLM powering them.\n",
    "\n",
    "This cell uses the `parse_converse_response` function processes the response received from the Converse API. It extracts various components from the response, including -\n",
    "\n",
    "- Role: The role of the message sender (e.g., user or model).\n",
    "- Text Content: Any text messages in the response, compiled into a list.\n",
    "- Images: Information about images, including format and byte data.\n",
    "- Documents: Details of any documents returned, such as format, name, and byte data.\n",
    "- Tool Usages and Results: Any tool usages or results included in the response.\n",
    "- Guardrail Content: Information related to guardrails, if present.\n",
    "- Stop Reason: The reason the conversation may have stopped.\n",
    "- Usage and Metrics: Additional usage statistics and performance metrics.\n",
    "\n",
    "The function returns a structured dictionary containing all extracted information, allowing for easy access to relevant data after conversing with the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_converse_response(response):\n",
    "    \"\"\"Parse Converse API response\"\"\"\n",
    "    output = response.get('output', {})\n",
    "    message = output.get('message', {})\n",
    "    role = message.get('role')\n",
    "    contents = message.get('content', [])\n",
    "\n",
    "    # Extract the text content if available\n",
    "    text_content = [item.get('text') for item in contents if 'text' in item]\n",
    "    \n",
    "    # Extract image data if available\n",
    "    images = [\n",
    "        {\n",
    "            'format': item['image']['format'],\n",
    "            'bytes': item['image']['source']['bytes']\n",
    "        }\n",
    "        for item in contents if 'image' in item\n",
    "    ]\n",
    "\n",
    "    # Extract document data if available\n",
    "    documents = [\n",
    "        {\n",
    "            'format': item['document']['format'],\n",
    "            'name': item['document']['name'],\n",
    "            'bytes': item['document']['source']['bytes']\n",
    "        }\n",
    "        for item in contents if 'document' in item\n",
    "    ]\n",
    "\n",
    "    # Extract tool use and tool results if present\n",
    "    tool_uses = [\n",
    "        item.get('toolUse') for item in contents if 'toolUse' in item\n",
    "    ]\n",
    "    tool_results = [\n",
    "        item.get('toolResult') for item in contents if 'toolResult' in item\n",
    "    ]\n",
    "\n",
    "    # Extract guardrail information if available\n",
    "    guard_content = [\n",
    "        item['guardContent'] for item in contents if 'guardContent' in item\n",
    "    ]\n",
    "\n",
    "    # Parse stop reason\n",
    "    stop_reason = response.get('stopReason')\n",
    "\n",
    "    # Parse token usage and metrics\n",
    "    usage = response.get('usage', {})\n",
    "    metrics = response.get('metrics', {})\n",
    "    \n",
    "    return {\n",
    "        'role': role,\n",
    "        'text_content': text_content,\n",
    "        'images': images,\n",
    "        'documents': documents,\n",
    "        'tool_uses': tool_uses,\n",
    "        'tool_results': tool_results,\n",
    "        'guard_content': guard_content,\n",
    "        'stop_reason': stop_reason,\n",
    "        'usage': usage,\n",
    "        'metrics': metrics\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Converse API\n",
    "The `converse` method enables conversational interaction by sending a series of messages to the specified model and is equipped with the `parse_converse_response` function to format and extract relevant output data from the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def converse(model_id, messages):\n",
    "    \"\"\"Use the Converse API to engage in a conversation with the specified model\"\"\"\n",
    "    response = bedrock_runtime.converse(\n",
    "        modelId=model_id,\n",
    "        messages=messages,\n",
    "        inferenceConfig={\n",
    "            'maxTokens': 300,  # Specify max tokens if needed\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"Converse Response:\", status_code)\n",
    "    parsed_response = parse_converse_response(response)\n",
    "    print(parsed_response)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Converse test\n",
    "This block demonstrates how to use the converse function in a practical test. A message list is created to structure the user input, specifying the role as `user` and including the `content` derived from the variable prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of converse\n",
    "print(\"\\nTesting Converse...\")\n",
    "prompt = \"\\n\\nHuman: Tell me about Amazon Bedrock.\\n\\nAssistant:\"\n",
    "messages = [{\"role\": \"user\", \"content\": [{\"text\": prompt}]}]\n",
    "response = converse(claims_dept_claude_3_sonnet_profile_arn, messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoke model with Inference Profile using ConverseStream API \n",
    "The `parse_converse_stream_response` function is designed to handle and interpret the responses from the `ConverseStream` API. It initializes an empty list to store parsed events and retrieves the event stream from the API response. The function checks if the stream is present; if not, it logs an appropriate message and returns an empty list. \n",
    "\n",
    "The function then iterates through each event, identifying the event type (such as messageStart, contentBlockStart, contentBlockDelta, contentBlockStop, messageStop, and metadata) and extracting relevant information for each type. Additionally, the function handles potential exceptions, logging pertinent details when encountered. The parsed events are collected into a list, which is returned for further processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_converse_stream_response(response):\n",
    "    \"\"\"Parse the ConverseStream API response\"\"\"\n",
    "    parsed_events = []\n",
    "    \n",
    "    # Access the EventStream directly\n",
    "    stream = response.get('stream')\n",
    "    \n",
    "    # Check if stream is valid\n",
    "    if stream is None:\n",
    "        print(\"No event stream found in the response.\")\n",
    "        return parsed_events\n",
    "    \n",
    "    # Iterate over each event in the EventStream\n",
    "    for event in stream:\n",
    "        print(\"Event:\", event)\n",
    "        event_type = event.get('eventType')\n",
    "        event_data = event.get('eventData', {})\n",
    "        parsed_event = {}\n",
    "        \n",
    "        if event_type == 'messageStart':\n",
    "            parsed_event['type'] = 'messageStart'\n",
    "            parsed_event['role'] = event_data.get('role')\n",
    "        \n",
    "        elif event_type == 'contentBlockStart':\n",
    "            parsed_event['type'] = 'contentBlockStart'\n",
    "            start_info = event_data.get('start', {})\n",
    "            parsed_event['contentBlockIndex'] = event_data.get('contentBlockIndex')\n",
    "            if 'toolUse' in start_info:\n",
    "                parsed_event['toolUseId'] = start_info['toolUse'].get('toolUseId')\n",
    "                parsed_event['toolName'] = start_info['toolUse'].get('name')\n",
    "\n",
    "        elif event_type == 'contentBlockDelta':\n",
    "            parsed_event['type'] = 'contentBlockDelta'\n",
    "            parsed_event['contentBlockIndex'] = event_data.get('contentBlockIndex')\n",
    "            delta = event_data.get('delta', {})\n",
    "            parsed_event['text'] = delta.get('text')\n",
    "            if 'toolUse' in delta:\n",
    "                parsed_event['toolInput'] = delta['toolUse'].get('input')\n",
    "        \n",
    "        elif event_type == 'contentBlockStop':\n",
    "            parsed_event['type'] = 'contentBlockStop'\n",
    "            parsed_event['contentBlockIndex'] = event_data.get('contentBlockIndex')\n",
    "\n",
    "        elif event_type == 'messageStop':\n",
    "            parsed_event['type'] = 'messageStop'\n",
    "            parsed_event['stopReason'] = event_data.get('stopReason')\n",
    "            parsed_event['additionalModelResponseFields'] = event_data.get('additionalModelResponseFields')\n",
    "\n",
    "        elif event_type == 'metadata':\n",
    "            parsed_event['type'] = 'metadata'\n",
    "            parsed_event['usage'] = event_data.get('usage', {})\n",
    "            parsed_event['metrics'] = event_data.get('metrics', {})\n",
    "            parsed_event['trace'] = event_data.get('trace', {})\n",
    "\n",
    "        # Handle guardrail assessments and policies within metadata.trace.guardrail\n",
    "        if event_type == 'metadata' and 'trace' in event_data:\n",
    "            guardrail = event_data['trace'].get('guardrail', {})\n",
    "            if 'inputAssessment' in guardrail:\n",
    "                parsed_event['inputAssessment'] = guardrail['inputAssessment']\n",
    "            if 'outputAssessments' in guardrail:\n",
    "                parsed_event['outputAssessments'] = guardrail['outputAssessments']\n",
    "        \n",
    "        # Handle exceptions\n",
    "        elif event_type in ['internalServerException', 'modelStreamErrorException', \n",
    "                            'validationException', 'throttlingException', 'serviceUnavailableException']:\n",
    "            parsed_event['type'] = 'exception'\n",
    "            parsed_event['exceptionType'] = event_type\n",
    "            parsed_event['message'] = event_data.get('message')\n",
    "            if event_type == 'modelStreamErrorException':\n",
    "                parsed_event['originalStatusCode'] = event_data.get('originalStatusCode')\n",
    "                parsed_event['originalMessage'] = event_data.get('originalMessage')\n",
    "\n",
    "        # Add the parsed event to the list of events\n",
    "        parsed_events.append(parsed_event)\n",
    "\n",
    "    return parsed_events"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ConverseStream API\n",
    "The `converse_stream` method leverages the `ConverseStream` API to facilitate a real-time conversational interaction with a specified model. The function initiates a streaming response by sending user messages to the model, with a configuration to limit the maximum token usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def converse_stream(model_id, messages):\n",
    "    \"\"\"Use the ConverseStream API to engage in a streaming response conversation with the specified model\"\"\"\n",
    "    response = bedrock_runtime.converse_stream(\n",
    "        modelId=model_id,\n",
    "        messages=messages,\n",
    "        inferenceConfig={\n",
    "            'maxTokens': 300,\n",
    "        }\n",
    "    )\n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"ConverseStream Response:\", status_code)\n",
    "    parsed_response = parse_converse_stream_response(response)\n",
    "    print(parsed_response)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ConverseStream test\n",
    "This cell demonstrates a test of the `converse_stream` functionality by passing in the application inference profile ARN and the user messages. The output from this interaction will showcase how the model responds to the user's input in a streaming manner, allowing for a dynamic conversational experience. The results of the conversation will be displayed in the console, providing insights into the model's performance and response handling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of converse stream\n",
    "print(\"\\nTesting ConverseStream...\")\n",
    "messages = [{\"role\": \"user\", \"content\": [{\"text\": prompt}]}]\n",
    "response = converse_stream(claims_dept_claude_3_sonnet_profile_arn, messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoke model with Inference Profile using InvokeModel API \n",
    "The `parse_converse_stream_response` function handles the response from the Amazon Bedrock `InvokeModel` API call. It checks if the response includes a body and, if so, processes the body stream to extract and decode the JSON content. The parsed output is returned as a Python dictionary, making it easier to work with the model's response data in subsequent code. If no body is found in the response, it prints a warning message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_invoke_model_response(response):\n",
    "    \"\"\"Parse the InvokeModel API response\"\"\"\n",
    "    \n",
    "    # Check if the response contains a body\n",
    "    if 'body' in response:\n",
    "        # Read the body stream\n",
    "        body_stream = response['body']\n",
    "        body_content = body_stream.read()  # This reads the streaming body\n",
    "        \n",
    "        # Decode the bytes to a string\n",
    "        body_str = body_content.decode('utf-8')\n",
    "        \n",
    "        # Parse the JSON string into a Python dictionary\n",
    "        output = json.loads(body_str)\n",
    "        \n",
    "        return output\n",
    "    else:\n",
    "        print(\"No body found in the response.\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### InvokeModel API\n",
    "The `invoke_model` method utilizes the Amazon Bedrock `InvokeModel` API to run inference on a specified model using provided input data. This function takes in the model's ARN and the input body, converts the body to a JSON-encoded byte format, and sends it to the API. After invoking the model, it prints the HTTP status code and the full response for debugging purposes, before returning the response for further processing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def invoke_model(model_id, body):\n",
    "    \"\"\"Use the InvokeModel API to invoke the specified model to run inference using the provided input\"\"\"\n",
    "    body_bytes = json.dumps(body).encode('utf-8')\n",
    "    response = bedrock_runtime.invoke_model(\n",
    "        modelId=model_id,\n",
    "        body=body_bytes\n",
    "    )\n",
    "    print(\"InvokeModel Response:\", response['ResponseMetadata']['HTTPStatusCode'])\n",
    "    print(response)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### InvokeModel test\n",
    "In this code block, an example invocation of the Bedrock model is demonstrated. It begins by defining a prompt and the corresponding input data. Next, it creates an inference profile using the specified base model ARN, which sets up the model for use. After retrieving the inference profile ARN, it calls the `invoke_model` function with the input data. Finally, it parses the model's response using the `parse_invoke_model_response` function, and if successful, prints the parsed output. This block showcases the complete workflow of setting up and invoking a model, along with handling the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example of invoking model\n",
    "print(\"Testing InvokeModel...\")\n",
    "prompt = \"Tell me about Amazon Bedrock\"\n",
    "input_data = {\n",
    "    \"prompt\": prompt\n",
    "}\n",
    "\n",
    "# Create Inference Profile\n",
    "print(\"Testing CreateInferenceProfile...\")\n",
    "tags = [{'key': 'dept', 'value': 'claims'}]\n",
    "base_model_arn = \"arn:aws:bedrock:us-east-1::foundation-model/meta.llama3-70b-instruct-v1:0\"\n",
    "claims_dept_llama3_70b_profile = create_inference_profile(\"claims_dept_llama3_70b_profile\", base_model_arn, tags)\n",
    "\n",
    "# Extracting the ARN and retrieving Inference Profile ID\n",
    "claims_dept_llama3_70b_profile_arn = claims_dept_llama3_70b_profile['inferenceProfileArn']\n",
    "\n",
    "response = invoke_model(claims_dept_llama3_70b_profile_arn, input_data)\n",
    "parsed_output = parse_invoke_model_response(response)\n",
    "\n",
    "if parsed_output:\n",
    "    print(\"Parsed Model Output:\", parsed_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Invoke model with Inference Profile using InvokeModelWithResponseStream API \n",
    "The `parse_invoke_model_with_stream` function is designed to process the response stream received from the `InvokeModelWithResponseStream` API. It checks for the presence of a 'body' in the response and iterates through the events within the EventStream. For each event that contains a 'chunk', it decodes the bytes and parses the JSON data, accumulating the results in a list. If no body is found in the response, it outputs a message indicating this. The function returns a list of parsed event data or `None` if the body is absent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_invoke_model_with_stream(response):\n",
    "    \"\"\"Parse the response stream from the InvokeModelWithResponseStream API call\"\"\"\n",
    "    output = []\n",
    "    if 'body' in response:\n",
    "        body_stream = response['body']\n",
    "        \n",
    "        # Iterate through the events in the EventStream\n",
    "        for event in body_stream:\n",
    "            # Each event is a dictionary, extract the data you need\n",
    "            if 'chunk' in event:\n",
    "                # Decode the chunk and parse it into JSON\n",
    "                event_data = event['chunk']['bytes']\n",
    "                output.append(json.loads(event_data.decode('utf-8'))) \n",
    "        \n",
    "        return output\n",
    "    else:\n",
    "        print(\"No body found in the response.\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### InvokeModelWithResponseStream API\n",
    "In this code block, the `invoke_model_with_stream` method invokes a specified model using the `InvokeModelWithResponseStream` API. It takes a model ID and an input body. After invoking the model, it extracts the HTTP status code from the response metadata and prints it for reference, returning the entire response object for further processing.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def invoke_model_with_stream(model_id, body):\n",
    "    \"\"\"Invoke the model with response stream using the specified model ID and input body\"\"\"\n",
    "    # Ensure the body contains the correct prompt key\n",
    "    body_bytes = json.dumps(body).encode('utf-8')\n",
    "    response = bedrock_runtime.invoke_model_with_response_stream(\n",
    "        modelId=model_id,\n",
    "        body=body_bytes,\n",
    "        contentType='application/json'\n",
    "    )\n",
    "    \n",
    "    # Extracting the HTTP status code from the response\n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"InvokeModelWithResponseStream Response:\", status_code)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### InvokeModelWithResponseStream test\n",
    "This block contains a simple test to demonstrate the functionality of the model invocation and response parsing. It initializes a prompt asking for information about Amazon Bedrock and creates an input data dictionary containing this prompt. The `invoke_model_with_stream` function is called with the specified model ID (e.g., `claims_dept_llama3_70b_profile_arn`) and the input data. The response from the model is then parsed using the `parse_invoke_model_with_stream` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\nTesting InvokeModelWithResponseStream...\")\n",
    "prompt = \"Tell me about Amazon Bedrock.\"\n",
    "input_data = {\"prompt\": prompt}\n",
    "\n",
    "# Example of invoking model with response stream\n",
    "response = invoke_model_with_stream(claims_dept_llama3_70b_profile_arn, input_data)\n",
    "parsed_output = parse_invoke_model_with_stream(response)\n",
    "\n",
    "print(\"Parsed Model Output:\", parsed_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tagging\n",
    "#### Create Tag\n",
    "The `tag_resource` method tags a specified resource in Amazon Bedrock. It takes a resource's ARN and a list of tags as input. In the provided example, two resources are tagged: one with the department tag \"claims\" and another with the department tag \"underwriting.\" This tagging process helps in organizing and managing resources based on specific departments, enhancing resource identification and categorization within the Amazon Bedrock environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def tag_resource(resource_arn, tags):\n",
    "    \"\"\"Tag a specified resource.\"\"\"\n",
    "    response = bedrock.tag_resource(\n",
    "        resourceARN=resource_arn,\n",
    "        tags=tags\n",
    "    )\n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"TagResource Response:\", status_code)\n",
    "    print(f\"{response}\\n\")\n",
    "    return response\n",
    "\n",
    "# Example of tagging resources\n",
    "print(\"\\nTesting TagResource...\")\n",
    "claims_tag = [{\"key\": \"dept\", \"value\": \"claims\"}]\n",
    "claims_response = tag_resource(claims_dept_claude_3_sonnet_profile_arn, claims_tag)\n",
    "\n",
    "underwriting_tag = [{\"key\": \"dept\", \"value\": \"underwriting\"}]\n",
    "underwriting_response = tag_resource(claims_dept_llama3_70b_profile_arn, underwriting_tag)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### List Tags\n",
    "In this code block, the `list_tags_for_resource` method retrieves and displays the tags associated with a specified Bedrock resourc. The function takes the resource ARN as its parameter. The example demonstrates how to list tags for the previously tagged resource, providing insight into its metadata and organization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def list_tags(resource_arn):\n",
    "    \"\"\"List tags for a specified resource.\"\"\"\n",
    "    response = bedrock.list_tags_for_resource(\n",
    "        resourceARN=resource_arn\n",
    "    )\n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"ListTagsForResource Response:\", status_code)\n",
    "    print(f\"{response}\\n\")\n",
    "    return response\n",
    "\n",
    "# Example of listing tags for a resource\n",
    "print(\"\\nTesting ListTagsForResource...\")\n",
    "claims_response = list_tags(claims_dept_claude_3_sonnet_profile_arn)\n",
    "underwriting_response = list_tags(claims_dept_llama3_70b_profile_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Remove Tag\n",
    "This block calls the `untag_resource` method, which removes specified tags from a given resource in Amazon Bedrock. The function accepts a resource ARN and a list of tag keys to be removed. The provided example illustrates the untagging process for a resource by removing a tag associated with the department, demonstrating how to manage and modify resource tagging effectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def untag_resource(resource_arn, tag_keys):\n",
    "    \"\"\"Untag a specified resource.\"\"\"\n",
    "    response = bedrock.untag_resource(\n",
    "        resourceARN=resource_arn,\n",
    "        tagKeys=tag_keys\n",
    "    )\n",
    "    status_code = response.get('ResponseMetadata', {}).get('HTTPStatusCode')\n",
    "    print(\"UntagResource Response:\", status_code)\n",
    "    print(f\"{response}\\n\")\n",
    "    return response\n",
    "\n",
    "# Example of untagging resources\n",
    "print(\"\\nTesting UntagResource...\")\n",
    "tag_keys = [\"dept\"]\n",
    "untag_response = untag_resource(claims_dept_claude_3_sonnet_profile_arn, tag_keys)\n",
    "claims_response = list_tags(claims_dept_claude_3_sonnet_profile_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up section\n",
    "The `delete_inference_profile` function accepts a Bedrock ARN as an argument to delete the corresponding application inference profile. The function is called twice to delete all inference profiles that were created earlier in the notebook: `claims_dept_llama3_70b_profile_arn` and `claims_dept_claude_3_sonnet_profile_arn`. This cleanup process ensures that resources are properly decommissioned, helping maintain a tidy and organized environment by preventing unused resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_inference_profile(inference_profile_arn):\n",
    "    \"\"\"Delete a specified inference profile by its ARN.\"\"\"\n",
    "    response = bedrock.delete_inference_profile(\n",
    "        inferenceProfileIdentifier=inference_profile_arn\n",
    "    )\n",
    "    print(\"DeleteInferenceProfile Response:\", response['ResponseMetadata']['HTTPStatusCode'])\n",
    "    print(response)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initiating the deletion of inference profiles\n",
    "print(\"Testing DeleteInferenceProfiles...\")\n",
    "delete_inference_profile(claims_dept_llama3_70b_profile_arn)\n",
    "delete_inference_profile(claims_dept_claude_3_sonnet_profile_arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# END OF NOTEBOOK "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}