bedrock_inference_example.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyO3uf+GyNU0sjk8wVwa9W0j",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/detsutut/bedrock-inference/blob/main/bedrock_inference_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Bedrock-inference package: how to use\n"
      ],
      "metadata": {
        "id": "p4mDxhJdoixy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 0. Install the bedrock-inference package from Pip"
      ],
      "metadata": {
        "id": "2xM-W7s8gaDO"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4MHMfy3bcRTH"
      },
      "outputs": [],
      "source": [
        "!pip install bedrock-inference"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 1. Establish a MFA authenticated session\n",
        "\n",
        "Here, secrets are retrieved from the Colab Secrets vault."
      ],
      "metadata": {
        "id": "ot9iGrpeep2G"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import userdata\n",
        "from bedrock_inference import bedrock\n",
        "\n",
        "session = bedrock.aws_login_mfa(token=\"087914\",\n",
        "                                arn=userdata.get('AWS_ARN_MFA_DEVICE'),\n",
        "                                aws_access_key_id=userdata.get('AWS_ACCESS_KEY_ID'),\n",
        "                                aws_secret_access_key=userdata.get('AWS_SECRET_ACCESS_KEY'))"
      ],
      "metadata": {
        "id": "Cq_t5UencVYr"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Check that the Bedrock service is available in the established session"
      ],
      "metadata": {
        "id": "mdJdkxTGgH-7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "assert 'bedrock' in session.get_available_services()\n",
        "assert 'bedrock-runtime' in session.get_available_services()"
      ],
      "metadata": {
        "id": "Gm_dYUR3fiV4"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2. Instantiate the Bedrock caller"
      ],
      "metadata": {
        "id": "tB7wBFfUeiW7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "caller = bedrock.Bedrock(session=session, region=\"us-west-2\")"
      ],
      "metadata": {
        "id": "TMimXMKFef_M"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Check the list of available models for your region and retrieve the target model id.\n",
        "In this case, we want to use Anthropic's Claude 3 Sonnet, but we don't know the correspondent AWS modelId."
      ],
      "metadata": {
        "id": "28rD6prVknTI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "models_df = caller.list_models()\n",
        "models_df[models_df[\"modelName\"].str.contains('claude 3 sonnet', case=False)]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 135
        },
        "id": "mr1vB86Ogx6v",
        "outputId": "1a3d3618-660b-4f31-f9ca-dd74b340de44"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                             modelArn  \\\n",
              "18  arn:aws:bedrock:us-west-2::foundation-model/an...   \n",
              "\n",
              "                                    modelId        modelName providerName  \\\n",
              "18  anthropic.claude-3-sonnet-20240229-v1:0  Claude 3 Sonnet    Anthropic   \n",
              "\n",
              "   inputModalities outputModalities responseStreamingSupported  \\\n",
              "18   [TEXT, IMAGE]           [TEXT]                       True   \n",
              "\n",
              "   customizationsSupported inferenceTypesSupported        modelLifecycle  \n",
              "18                      []             [ON_DEMAND]  {'status': 'ACTIVE'}  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-e547c671-3932-456d-beee-d50f98abbd13\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>modelArn</th>\n",
              "      <th>modelId</th>\n",
              "      <th>modelName</th>\n",
              "      <th>providerName</th>\n",
              "      <th>inputModalities</th>\n",
              "      <th>outputModalities</th>\n",
              "      <th>responseStreamingSupported</th>\n",
              "      <th>customizationsSupported</th>\n",
              "      <th>inferenceTypesSupported</th>\n",
              "      <th>modelLifecycle</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>18</th>\n",
              "      <td>arn:aws:bedrock:us-west-2::foundation-model/an...</td>\n",
              "      <td>anthropic.claude-3-sonnet-20240229-v1:0</td>\n",
              "      <td>Claude 3 Sonnet</td>\n",
              "      <td>Anthropic</td>\n",
              "      <td>[TEXT, IMAGE]</td>\n",
              "      <td>[TEXT]</td>\n",
              "      <td>True</td>\n",
              "      <td>[]</td>\n",
              "      <td>[ON_DEMAND]</td>\n",
              "      <td>{'status': 'ACTIVE'}</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-e547c671-3932-456d-beee-d50f98abbd13')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-e547c671-3932-456d-beee-d50f98abbd13 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-e547c671-3932-456d-beee-d50f98abbd13');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "summary": "{\n  \"name\": \"models_df[models_df[\\\"modelName\\\"]\",\n  \"rows\": 1,\n  \"fields\": [\n    {\n      \"column\": \"modelArn\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"modelId\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"anthropic.claude-3-sonnet-20240229-v1:0\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"modelName\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"Claude 3 Sonnet\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"providerName\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          \"Anthropic\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"inputModalities\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"outputModalities\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"responseStreamingSupported\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 1,\n        \"samples\": [\n          true\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"customizationsSupported\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"inferenceTypesSupported\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"modelLifecycle\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}"
            }
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. Invoke the target model and get the generated answer"
      ],
      "metadata": {
        "id": "Z3WrK_gJk341"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "answer = caller.invoke_model(prompt=\"What's the conceptual opposite of 'Hello World'?\",\n",
        "                             model_id=\"anthropic.claude-3-sonnet-20240229-v1:0\",\n",
        "                             maxTokens = 200)\n",
        "print(answer)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZVxKNhmrj-GH",
        "outputId": "107784cf-0174-4fcc-f2d0-71b6c1f0e9f2"
      },
      "execution_count": 39,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "The conceptual opposite of the famous \"Hello World\" program, which is often used to introduce programming concepts and verify that a language or environment is set up correctly, could be considered \"Goodbye World.\"\n",
            "\n",
            "\"Hello World\" is typically the first program that new programmers write, as it serves as a simple introduction to printing output to the console or displaying a message. It's a way of saying \"hello\" to the world of programming.\n",
            "\n",
            "On the other hand, \"Goodbye World\" could symbolize the end of a programming journey, a farewell to the world of coding, or the termination of a program's execution. It's the conceptual antithesis of \"Hello World,\" representing the opposite of starting something new and instead signifying the end or completion of a process.\n",
            "\n",
            "While \"Hello World\" is a warm greeting and a starting point, \"Goodbye World\" could be seen as a final message or a symbolic way to conclude a program or a programming\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "---------------------\n",
        "\n",
        "## Advanced Usage\n",
        "\n",
        "`bedrock-inference` handles API requests for you in a nice and user-friendly way.\n",
        "\n",
        "However, this is just a nice wrapping of the original API requests.\n",
        "\n",
        "If you want to take control of what's happening under the hood, you can use the original Bedrock clients to send custom requests."
      ],
      "metadata": {
        "id": "cJUf53WXlFZs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In the following example, we want to retrieve the raw request output instead of just the textual answer generated by the model.\n",
        "\n",
        "Check the [boto3 documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime.html) for further details.\n"
      ],
      "metadata": {
        "id": "B2KraoiWmhY2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "client = caller._runtime_client_"
      ],
      "metadata": {
        "id": "rPeAMZ8Zm4Pi"
      },
      "execution_count": 40,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response = client.converse(\n",
        "    modelId=\"anthropic.claude-3-sonnet-20240229-v1:0\",\n",
        "    messages=[\n",
        "            {\n",
        "                \"role\": \"user\",\n",
        "                \"content\": [{\"text\": \"What's the conceptual opposite of 'Hello World'?\"}],\n",
        "            }\n",
        "        ],\n",
        "    inferenceConfig={\n",
        "        'maxTokens': 200,\n",
        "    }\n",
        ")"
      ],
      "metadata": {
        "id": "9XyGULaBm9zZ"
      },
      "execution_count": 42,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "response"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "KPqTv42roFnC",
        "outputId": "7150f8f4-6854-4614-b3d7-d5257393106b"
      },
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'ResponseMetadata': {'RequestId': 'c801c2ed-d475-4ef1-a8d4-640d8a0cae9b',\n",
              "  'HTTPStatusCode': 200,\n",
              "  'HTTPHeaders': {'date': 'Mon, 30 Sep 2024 15:50:31 GMT',\n",
              "   'content-type': 'application/json',\n",
              "   'content-length': '1213',\n",
              "   'connection': 'keep-alive',\n",
              "   'x-amzn-requestid': 'c801c2ed-d475-4ef1-a8d4-640d8a0cae9b'},\n",
              "  'RetryAttempts': 0},\n",
              " 'output': {'message': {'role': 'assistant',\n",
              "   'content': [{'text': 'The conceptual opposite of the famous \"Hello World\" program, which is often used as a simple introduction to programming, could be considered \"Goodbye World.\"\\n\\n\"Hello World\" is typically one of the first programs a beginner programmer writes when learning a new programming language. It\\'s a simple program that outputs the text \"Hello World\" to the console or user interface, serving as a way to verify that the development environment is set up correctly and the programmer can successfully compile and run basic code.\\n\\nIn contrast, \"Goodbye World\" could be seen as a symbolic representation of the end of a programming journey or the termination of a program\\'s execution. It might be used as a final program or exercise to signal the completion of a course or tutorial on programming basics.\\n\\nWhile \"Hello World\" marks the beginning and the programmer\\'s entry into the world of coding, \"Goodbye World\" could signify the exit or the conclusion of that initial learning phase, indicating that the programmer'}]}},\n",
              " 'stopReason': 'max_tokens',\n",
              " 'usage': {'inputTokens': 20, 'outputTokens': 200, 'totalTokens': 220},\n",
              " 'metrics': {'latencyMs': 6666}}"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    }
  ]
}