labs/zero-to-production/zero-to-production.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# APIM ❤️ OpenAI\n",
    "\n",
    "## Zero-to-Production lab\n",
    "\n",
    "Playground to create a combination of several policies in an iterative approach. We will accomplish three successive policy additions:\n",
    "\n",
    "1) Add load balancing across multiple region\n",
    "1) Add token emitting to observe token usage\n",
    "1) Apply token rate limiting to avoid runaway token usage scenarios\n",
    "\n",
    "Each of these sets of policies is derived from other labs in this repo.\n",
    "\n",
    "### Prerequisites\n",
    "- [Python 3.12 or later version](https://www.python.org/) installed\n",
    "- [Pandas Library](https://pandas.pydata.org/) and matplotlib installed\n",
    "- [VS Code](https://code.visualstudio.com/) installed with the [Jupyter notebook extension](https://marketplace.visualstudio.com/items?itemName=ms-toolsai.jupyter) enabled\n",
    "- [Azure CLI](https://learn.microsoft.com/cli/azure/install-azure-cli) installed\n",
    "- [An Azure Subscription](https://azure.microsoft.com/free/) with Contributor permissions\n",
    "- [Access granted to Azure OpenAI](https://aka.ms/oai/access) or just enable the mock service\n",
    "- [# to Azure with Azure CLI](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize notebook variables\n",
    "\n",
    "- Resources will be suffixed by a unique string based on your subscription id.\n",
    "- Change the location parameters according your preferences and on the [product availability by Azure region.](https://azure.microsoft.com/explore/global-infrastructure/products-by-region/?cdn=disable&products=cognitive-services,api-management)\n",
    "- Adjust the OpenAI model, version according the [availability by region.](https://learn.microsoft.com/azure/ai-services/openai/concepts/models) \n",
    "- Experiment with the priority, weight, and capacity OpenAI parameters to affect the load balancing\n",
    "- `capacity` is set intentionally low - the unit is tokens per minute - to trigger the retry logic in the load balancer (transparent to the user) as well as the priority failover from priority 1 to 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, json\n",
    "sys.path.insert(1, '../../shared')  # add the shared directory to the Python path\n",
    "import utils\n",
    "\n",
    "deployment_name = os.path.basename(os.path.dirname(globals()['__vsc_ipynb_file__']))\n",
    "resource_group_name = f\"lab-{deployment_name}\" # change the name to match your naming style\n",
    "resource_group_location = \"eastus2\"\n",
    "\n",
    "apim_sku = 'Basicv2'\n",
    "\n",
    "# Prioritize East US until exhaustion (simulate PTU with TPM), then equally distribute between Sweden and West US (consumption fallback)\n",
    "openai_resources = [\n",
    "    {\"name\": \"openai1\", \"location\": \"eastus\", \"priority\": 1, \"weight\": 100, \"capacity\": 4},\n",
    "    {\"name\": \"openai2\", \"location\": \"swedencentral\", \"priority\": 2, \"weight\": 50, \"capacity\": 8 },\n",
    "    {\"name\": \"openai3\", \"location\": \"westus\", \"priority\": 2, \"weight\": 50, \"capacity\": 8}\n",
    "]\n",
    "\n",
    "openai_deployment_name = \"gpt-4o-mini\"\n",
    "openai_model_name = \"gpt-4o-mini\"\n",
    "openai_model_version = \"2024-07-18\"\n",
    "openai_model_sku = 'Standard'\n",
    "openai_api_version = \"2024-02-01\"\n",
    "\n",
    "backend_id = 'openai-backend-pool' if len(openai_resources) > 1 else openai_resources[0]['name']\n",
    "\n",
    "utils.print_ok('Notebook initialized')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verify the Azure CLI and the connected Azure subscription\n",
    "\n",
    "The following commands ensure that you have the latest version of the Azure CLI and that the Azure CLI is connected to your Azure subscription."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = utils.run(\"az account show\", \"Retrieved az account\", \"Failed to get the current az account\")\n",
    "\n",
    "if output.success and output.json_data:\n",
    "    current_user = output.json_data['user']['name']\n",
    "    tenant_id = output.json_data['tenantId']\n",
    "    subscription_id = output.json_data['id']\n",
    "\n",
    "    utils.print_info(f\"Current user: {current_user}\")\n",
    "    utils.print_info(f\"Tenant ID: {tenant_id}\")\n",
    "    utils.print_info(f\"Subscription ID: {subscription_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy 1 - Load Balancing\n",
    "\n",
    "This lab uses [Bicep](https://learn.microsoft.com/azure/azure-resource-manager/bicep/overview?tabs=bicep) to declaratively define all resources that will be deployed in the specified resource group. Change the parameters or the [main.bicep](main.bicep) directly to try different configurations.\n",
    "\n",
    "#### Create deployment using 🦾 Bicep\n",
    "\n",
    "The `retry-count` parameter should have a value that represents one less than the total number of backends. For example, if we have three defined Azure OpenAI backends, we want to try initially, then have up to two retries, so long as we have remaining, active backends. This ensures that we cover all available backends."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_xml_file = \"policy-1.xml\"\n",
    "bicep_parameters_file = \"params-1.json\"\n",
    "\n",
    "# Create the resource group if doesn't exist\n",
    "utils.create_resource_group(resource_group_name, resource_group_location)\n",
    "\n",
    "# Define the Bicep parameters\n",
    "bicep_parameters = {\n",
    "    \"$schema\": \"https://schema.management.azure.com/schemas/2019-04-01/deploymentParameters.json#\",\n",
    "    \"contentVersion\": \"1.0.0.0\",\n",
    "    \"parameters\": {\n",
    "        \"apimSku\": { \"value\": apim_sku },\n",
    "        \"openAIConfig\": { \"value\": openai_resources },\n",
    "        \"openAIDeploymentName\": { \"value\": openai_deployment_name },\n",
    "        \"openAIModelName\": { \"value\": openai_model_name },\n",
    "        \"openAIModelVersion\": { \"value\": openai_model_version },\n",
    "        \"openAIModelSKU\": { \"value\": openai_model_sku },\n",
    "        \"openAIAPIVersion\": { \"value\": openai_api_version }\n",
    "    }\n",
    "}\n",
    "\n",
    "bicep_parameters = utils.create_bicep_params(policy_xml_file, bicep_parameters_file, bicep_parameters, [\n",
    "    ('{backend-id}', backend_id),\n",
    "    ('{retry-count}', len(openai_resources) - 1)\n",
    "])\n",
    "\n",
    "# Run the deployment\n",
    "output = utils.run(f\"az deployment group create --name {deployment_name} --resource-group {resource_group_name} --template-file main.bicep --parameters {bicep_parameters_file}\",\n",
    "    f\"Deployment '{deployment_name}' succeeded\", f\"Deployment '{deployment_name}' failed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2deploymentoutputs'></a>\n",
    "#### Get the deployment outputs\n",
    "\n",
    "Retrieve the required outputs from the Bicep deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obtain all of the outputs from the deployment\n",
    "output = utils.run(f\"az deployment group show --name {deployment_name} -g {resource_group_name}\", f\"Retrieved deployment: {deployment_name}\", f\"Failed to retrieve deployment: {deployment_name}\")\n",
    "\n",
    "if output.success and output.json_data:\n",
    "    apim_service_id = utils.get_deployment_output(output, 'apimServiceId', 'APIM Service Id')\n",
    "    apim_service_name = utils.get_deployment_output(output, 'apimServiceName', 'APIM Service Name')\n",
    "    apim_resource_gateway_url = utils.get_deployment_output(output, 'apimResourceGatewayURL', 'APIM API Gateway URL')\n",
    "    apim_subscription1_key = utils.get_deployment_output(output, 'apimSubscription1Key', 'APIM Subscription 1 Key (masked)', True)\n",
    "    apim_subscription2_key = utils.get_deployment_output(output, 'apimSubscription2Key', 'APIM Subscription 2 Key (masked)', True)\n",
    "    apim_subscription3_key = utils.get_deployment_output(output, 'apimSubscription3Key', 'APIM Subscription 3 Key (masked)', True)\n",
    "    app_insights_name = utils.get_deployment_output(output, 'applicationInsightsName', 'Application Insights Name')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2sdk'></a>\n",
    "#### 🧪 Test the API using the Azure OpenAI Python SDK\n",
    "\n",
    "Use the OpenAI Python SDK to make requests to API Management and any of the Azure OpenAI backends. Note that we return the `x-ms-region` header to show the frontend which backend was used. You may not want to do that in a production scenario.\n",
    "\n",
    "You will not see HTTP 429s returned as API Management's `retry` policy will select an available backend. If no backends are viable, an HTTP 503 will be returned.\n",
    "\n",
    "Tip: Use the [tracing tool](../../tools/tracing.ipynb) to track the behavior of the backend pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "runs = 20\n",
    "sleep_time_ms = 100\n",
    "total_tokens_all_runs = 0\n",
    "api_runs = []\n",
    "\n",
    "client = AzureOpenAI(\n",
    "    azure_endpoint = apim_resource_gateway_url,\n",
    "    api_key = apim_subscription1_key,\n",
    "    api_version = openai_api_version\n",
    ")\n",
    "\n",
    "for i in range(runs):\n",
    "    print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "    start_time = time.time()\n",
    "    raw_response = client.chat.completions.with_raw_response.create(\n",
    "        model = openai_model_name,\n",
    "        messages = [\n",
    "            {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "        ])\n",
    "    response_time = time.time() - start_time\n",
    "\n",
    "    print(f\"⌚ {response_time:.2f} seconds\")\n",
    "\n",
    "    if \"x-ms-region\" in raw_response.headers:\n",
    "        print(f\"x-ms-region: \\x1b[1;32m{raw_response.headers.get(\"x-ms-region\")}\\x1b[0m\") # this header is useful to determine the region of the backend that served the request\n",
    "        api_runs.append((response_time, raw_response.headers.get(\"x-ms-region\")))\n",
    "\n",
    "    response = raw_response.parse()\n",
    "\n",
    "    if response.usage:\n",
    "        total_tokens_all_runs += response.usage.total_tokens\n",
    "        print(f\"Token usage:\\n   Total tokens: {response.usage.total_tokens}\\n   Prompt tokens: {response.usage.prompt_tokens}\\n   Completion tokens: {response.usage.completion_tokens}\\n   Total tokens all runs: {total_tokens_all_runs}\\n\")\n",
    "\n",
    "\n",
    "    print(f\"💬 {response.choices[0].message.content}\\n\")\n",
    "\n",
    "    time.sleep(sleep_time_ms/1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2plot'></a>\n",
    "#### 🔍 Analyze Load Balancing results\n",
    "\n",
    "The priority 1 backend will be used until TPM exhaustion sets in, then distribution will occur near equally across the two priority 2 backends with 50/50 weights.  \n",
    "\n",
    "Please note that the first request of the lab can take a bit longer and should be discounted statistically as an outlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle as pltRectangle\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.figsize'] = [15, 7]\n",
    "df = pd.DataFrame(api_runs, columns = ['Response Time', 'Region'])\n",
    "df['Run'] = range(1, len(df) + 1)\n",
    "\n",
    "# Define a color map for each region\n",
    "color_map = {'East US': 'lightpink', 'Sweden Central': 'lightyellow', 'West US': 'lightblue'}  # Add more regions and colors as needed\n",
    "\n",
    "# Plot the dataframe with colored bars\n",
    "ax = df.plot(kind = 'bar', x = 'Run', y = 'Response Time', color = [color_map.get(region, 'gray') for region in df['Region']], legend = False)\n",
    "\n",
    "# Add legend\n",
    "legend_labels = [pltRectangle((0, 0), 1, 1, color = color_map.get(region, 'gray')) for region in df['Region'].unique()]\n",
    "ax.legend(legend_labels, df['Region'].unique())\n",
    "\n",
    "plt.title('Load Balancing results')\n",
    "plt.xlabel('Run #')\n",
    "plt.ylabel('Response Time')\n",
    "plt.xticks(rotation = 0)\n",
    "\n",
    "average = df['Response Time'].mean()\n",
    "plt.axhline(y = average, color = 'r', linestyle = '--', label = f'Average: {average:.2f}')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy 2 - Token Emitting\n",
    "\n",
    "We now add token emitting to the existing API policy in order to track token usage by subscriptions. This aids usage and cost analysis and chargeback models inside organizations. You can see the policy to be added in the `policy-2.xml` file in this folder.\n",
    "\n",
    "#### Update the API management policy via the REST API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_xml_file = \"policy-2.xml\"\n",
    "\n",
    "with open(policy_xml_file, 'r') as file:\n",
    "    policy_xml = file.read()\n",
    "    policy_xml = policy_xml.replace('{backend-id}', backend_id).replace('{retry-count}', str(len(openai_resources) - 1))\n",
    "\n",
    "utils.update_api_policy(subscription_id, resource_group_name, apim_service_name, \"openai\", policy_xml)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3sdk'></a>\n",
    "#### 🧪 Execute multiple runs for each subscription using the Azure OpenAI Python SDK\n",
    "\n",
    "We will send requests for each subscription. Adjust the `sleep_time_ms` and the number of `runs` to your test scenario.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "runs = 5\n",
    "sleep_time_ms = 100\n",
    "total_tokens_all_runs = [0, 0, 0]\n",
    "\n",
    "clients = [\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription1_key,\n",
    "        api_version = openai_api_version\n",
    "    ),\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription2_key,\n",
    "        api_version = openai_api_version\n",
    "    ),\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription3_key,\n",
    "        api_version = openai_api_version\n",
    "    )\n",
    "]\n",
    "\n",
    "for i in range(runs):\n",
    "    print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "    for j in range(0, 3):\n",
    "        start_time = time.time()\n",
    "\n",
    "        try:\n",
    "            raw_response = clients[j].chat.completions.with_raw_response.create(\n",
    "                model = openai_model_name,\n",
    "                messages = [\n",
    "                    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "                    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "                ],\n",
    "                extra_headers = {\"x-user-id\": \"alex\"}\n",
    "            )\n",
    "\n",
    "            response_time = time.time() - start_time\n",
    "            print(f\"🔑 Subscription {j+1}\")\n",
    "            print(f\"⌚ {response_time:.2f} seconds\")\n",
    "\n",
    "            if \"x-ms-region\" in raw_response.headers:\n",
    "                print(f\"x-ms-region: \\x1b[1;32m{raw_response.headers.get(\"x-ms-region\")}\\x1b[0m\") # this header is useful to determine the region of the backend that served the request\n",
    "\n",
    "            response = raw_response.parse()\n",
    "\n",
    "            if response.usage:\n",
    "                total_tokens_all_runs[j] += response.usage.total_tokens\n",
    "                print(f\"Token usage:\\n   Total tokens: {response.usage.total_tokens}\\n   Prompt tokens: {response.usage.prompt_tokens}\\n   Completion tokens: {response.usage.completion_tokens}\\n   Total tokens all runs: {total_tokens_all_runs[j]}\\n\")\n",
    "\n",
    "            print(f\"💬 {response.choices[0].message.content}\\n\")\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "\n",
    "    print()\n",
    "\n",
    "    time.sleep(sleep_time_ms/1000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3metricsinportal'></a>\n",
    "### 🔍 See the metrics on the Azure Portal\n",
    "\n",
    "One way to see the newly-added token metrics in the Azure Portal:\n",
    "\n",
    "1) Open the _Application Insights_ resource in the resource group.\n",
    "1) Navigate to the _Metrics_ blade.\n",
    "1) Change the timespan to the last 30 minutes with a 1 minute time granularity.\n",
    "1) Then select the _openai_ metric namespace. \n",
    "1) Choose the _Total Tokens_ metric \n",
    "1) Select the _Sum_ aggregation. \n",
    "1) Apply splitting by _Subscription Id_ to view values for each dimension. For better visibility switch to an area chart.\n",
    "\n",
    "![result](result.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3kql'></a>\n",
    "#### 🔍 Analyze Application Insights custom metrics with a KQL query\n",
    "\n",
    "Alternatively, you can query the custom metrics that were emitted by Azure APIM. **Note that it may take a few minutes for data to become available.** If you see the metrics in the Azure Portal, you don't have to wait for the data here and plotting it to continue. Come back to this section later, if you like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "query = \"\\\"\" + \"customMetrics \\\n",
    "| where name == 'Total Tokens' \\\n",
    "| extend parsedCustomDimensions = parse_json(customDimensions) \\\n",
    "| extend clientIP = tostring(parsedCustomDimensions.['Client IP']) \\\n",
    "| extend apiId = tostring(parsedCustomDimensions.['API ID']) \\\n",
    "| extend apimSubscription = tostring(parsedCustomDimensions.['Subscription ID']) \\\n",
    "| extend UserId = tostring(parsedCustomDimensions.['User ID']) \\\n",
    "| project timestamp, value, clientIP, apiId, apimSubscription, UserId \\\n",
    "| order by timestamp asc\" + \"\\\"\"\n",
    "\n",
    "output = utils.run(f\"az monitor app-insights query --app {app_insights_name} -g {resource_group_name} --analytics-query {query}\",\n",
    "    f\"App Insights query succeeded\", f\"App Insights query  failed\")\n",
    "\n",
    "table = output.json_data['tables'][0]\n",
    "df = pd.DataFrame(table.get(\"rows\"), columns = [col.get(\"name\") for col in table.get('columns')])\n",
    "df['timestamp'] = pd.to_datetime(df['timestamp']).dt.strftime('%H:%M')\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3plot'></a>\n",
    "#### 🔍 Plot the custom metrics results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the results\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.figsize'] = [15, 7]\n",
    "ax = df.plot(kind = 'line', x = 'timestamp', y = 'value', legend = False)\n",
    "plt.title('Total token usage over time')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Tokens')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Policy 3 - Token Rate Limiting\n",
    "\n",
    "Emitting tokens is very valuable to understand usage. Even more valuable is adding rate limiting for subscriptions as a sensible way to limit runaway usage.\n",
    "\n",
    "Please note that results are going to be a bit skewed as we are executing these requests sequentially, but it should get the gist across as to what's happening from a rate-limiting perspective. Threading may be considered in the future.\n",
    "\n",
    "#### Update the API management policy via the REST API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy_xml_file = \"policy-3.xml\"\n",
    "tokens_per_minute = 500\n",
    "\n",
    "with open(policy_xml_file, 'r') as file:\n",
    "    policy_xml = file.read()\n",
    "    policy_xml = policy_xml.replace('{backend-id}', backend_id).replace('{retry-count}', str(len(openai_resources) - 1)).replace('{tpm}', str(tokens_per_minute))\n",
    "\n",
    "utils.update_api_policy(subscription_id, resource_group_name, apim_service_name, \"openai\", policy_xml)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4sdk'></a>\n",
    "#### 🧪 Execute multiple runs for each subscription using the Azure OpenAI Python SDK\n",
    "\n",
    "We will send requests for each subscription. Adjust the `sleep_time_ms` and the number of `runs` to your test scenario. You should be able to observe a significant pause after a few runs as the requests will hit the tokens-per-minute limit that we defined earlier. This is expected and validation of the policies working.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "runs = 3\n",
    "sleep_time_ms = 100\n",
    "total_tokens_all_runs = [0, 0, 0]\n",
    "api_runs = []\n",
    "\n",
    "clients = [\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription1_key,\n",
    "        api_version = openai_api_version\n",
    "    ),\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription2_key,\n",
    "        api_version = openai_api_version\n",
    "    ),\n",
    "    AzureOpenAI(\n",
    "        azure_endpoint = apim_resource_gateway_url,\n",
    "        api_key = apim_subscription3_key,\n",
    "        api_version = openai_api_version\n",
    "    )\n",
    "]\n",
    "\n",
    "for i in range(runs):\n",
    "    print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "    for j in range(0, 3):\n",
    "        start_time = time.time()\n",
    "\n",
    "        try:\n",
    "            raw_response = clients[j].chat.completions.with_raw_response.create(\n",
    "                model = openai_model_name,\n",
    "                messages = [\n",
    "                    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "                    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "                ],\n",
    "                extra_headers = {\"x-user-id\": \"alex\"}\n",
    "            )\n",
    "\n",
    "            response_time = time.time() - start_time\n",
    "            print(f\"🔑 Subscription {j+1}\")\n",
    "            print(f\"⌚ {response_time:.2f} seconds\")\n",
    "            api_runs.append((response_time, j+1))\n",
    "\n",
    "            if \"x-ms-region\" in raw_response.headers:\n",
    "                print(f\"x-ms-region: \\x1b[1;32m{raw_response.headers.get(\"x-ms-region\")}\\x1b[0m\") # this header is useful to determine the region of the backend that served the request\n",
    "\n",
    "            response = raw_response.parse()\n",
    "\n",
    "            if response.usage:\n",
    "                total_tokens_all_runs[j] += response.usage.total_tokens\n",
    "                print(f\"Token usage:\\n   Total tokens: {response.usage.total_tokens}\\n   Prompt tokens: {response.usage.prompt_tokens}\\n   Completion tokens: {response.usage.completion_tokens}\\n   Total tokens all runs: {total_tokens_all_runs[j]}\\n\")\n",
    "\n",
    "            print(f\"💬 {response.choices[0].message.content}\\n\")\n",
    "        except e as Exception:\n",
    "            print(e)\n",
    "\n",
    "    print()\n",
    "\n",
    "    time.sleep(sleep_time_ms/1000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle as pltRectangle\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.figsize'] = [15, 7]\n",
    "df = pd.DataFrame(api_runs, columns = ['Response Time', 'Subscription'])\n",
    "df['Run'] = range(1, len(df) + 1)\n",
    "\n",
    "# Define a color map for each subscription\n",
    "color_map = {1: 'lightpink', 2: 'lightyellow', 3: 'lightblue'}  # Add more subscriptions and colors as needed\n",
    "\n",
    "# Plot the dataframe with colored bars\n",
    "ax = df.plot(kind = 'bar', x = 'Run', y = 'Response Time', color = [color_map.get(subscription, 'gray') for subscription in df['Subscription']], legend = False)\n",
    "\n",
    "# Add legend\n",
    "legend_labels = [pltRectangle((0, 0), 1, 1, color = color_map[subscription]) for subscription in color_map.keys()]\n",
    "ax.legend(legend_labels, ['Subscription 1', 'Subscription 2', 'Subscription 3'])\n",
    "\n",
    "plt.title('Token Rate Limiting by Subscription Results (high response times indicate likely throttling)')\n",
    "plt.xlabel('Run #')\n",
    "plt.ylabel('Response Time')\n",
    "plt.xticks(rotation = 0)\n",
    "\n",
    "average = df['Response Time'].mean()\n",
    "plt.axhline(y = average, color = 'r', linestyle = '--', label = f'Average: {average:.2f}')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='clean'></a>\n",
    "### 🗑️ Clean up resources\n",
    "\n",
    "When you're finished with the lab, you should remove all your deployed resources from Azure to avoid extra charges and keep your Azure subscription uncluttered.\n",
    "Use the [clean-up-resources notebook](clean-up-resources.ipynb) for that."
   ]
  }
 ],
 "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}