labs/token-metrics-emitting/token-metrics-emitting.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# APIM ❤️ OpenAI\n",
    "\n",
    "## Token Metrics Emitting lab\n",
    "![flow](../../images/token-metrics-emitting.gif)\n",
    "\n",
    "Playground to try the [emit token metric policy](https://learn.microsoft.com/azure/api-management/azure-openai-emit-token-metric-policy). The policy sends metrics to Application Insights about consumption of large language model tokens through Azure OpenAI Service APIs.\n",
    "\n",
    "Notes:\n",
    "- Token count metrics include: Total Tokens, Prompt Tokens, and Completion Tokens.\n",
    "- This policy supports OpenAI response streaming! Use the [streaming tool](../../tools/streaming.ipynb) to test and troubleshoot response streaming.\n",
    "- Use the [tracing tool](../../tools/tracing.ipynb) to track the behavior and troubleshoot the [policy](policy.xml).\n",
    "\n",
    "[View policy configuration](policy.xml)\n",
    "\n",
    "### Result\n",
    "![result](result.png)\n",
    "\n",
    "### TOC\n",
    "- [0️⃣ Initialize notebook variables](#0)\n",
    "- [1️⃣ Verify the Azure CLI and the connected Azure subscription](#1)\n",
    "- [2️⃣ Create deployment using 🦾 Bicep](#2)\n",
    "- [3️⃣ Get the deployment outputs](#3)\n",
    "- [🧪 Test the API using a direct HTTP call](#requests)\n",
    "- [🧪 Execute multiple runs for each subscription using the Azure OpenAI Python SDK](#sdk)\n",
    "- [🔍 Analyze Application Insights custom metrics with a KQL query](#kql)\n",
    "- [🔍 Plot the custom metrics results](#plot)\n",
    "- [🔍 See the metrics on the Azure Portal](#portal)\n",
    "- [🗑️ Clean up resources](#clean)\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)\n",
    "- [Sign in to Azure with Azure CLI](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='0'></a>\n",
    "### 0️⃣ Initialize notebook variables\n",
    "\n",
    "- Resources will be suffixed by a unique string based on your subscription id.\n",
    "- Adjust 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 and version according the [availability by region.](https://learn.microsoft.com/azure/ai-services/openai/concepts/models) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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 = \"westeurope\"\n",
    "\n",
    "apim_sku = 'Basicv2'\n",
    "\n",
    "openai_resources = [\n",
    "    {\"name\": \"openai1\", \"location\": \"swedencentral\"}\n",
    "]\n",
    "\n",
    "openai_model_name = \"gpt-35-turbo\"\n",
    "openai_model_version = \"0613\"\n",
    "openai_deployment_name = \"gpt-35-turbo\"\n",
    "openai_api_version = \"2024-02-01\"\n",
    "\n",
    "utils.print_ok('Notebook initiaized')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='1'></a>\n",
    "### 1️⃣ 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",
    "if output.success and output.json_data:\n",
    "    current_user = output.json_data['user']['name']\n",
    "    subscription_id = output.json_data['id']\n",
    "    tenant_id = output.json_data['tenantId']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2'></a>\n",
    "### 2️⃣ Create deployment using 🦾 Bicep\n",
    "\n",
    "This lab uses [Bicep](https://learn.microsoft.com/azure/azure-resource-manager/bicep/overview?tabs=bicep) to declarative define all the resources that will be deployed in the specified resource group. Change the parameters or the [main.bicep](main.bicep) directly to try different configurations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the resource group if doesn't exist\n",
    "utils.create_resource_group(True, resource_group_name, resource_group_location)\n",
    "\n",
    "# update the APIM policy file before the deployment\n",
    "policy_xml = None\n",
    "with open(\"policy.xml\", 'r') as policy_xml_file:\n",
    "    policy_template_xml = policy_xml_file.read()\n",
    "    if \"{backend-id}\" in policy_template_xml:\n",
    "        policy_xml = policy_template_xml.replace(\"{backend-id}\", str(\"openai-backend-pool\" if len(openai_resources) > 1 else openai_resources[0].get(\"name\")))                \n",
    "    policy_xml_file.close()\n",
    "if policy_xml is not None:\n",
    "    open(\"policy.xml\", 'w').write(policy_xml)\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",
    "        \"openAIAPIVersion\": { \"value\": openai_api_version }\n",
    "    }\n",
    "}\n",
    "\n",
    "# write the parameters to a file \n",
    "with open('params.json', 'w') as bicep_parameters_file:\n",
    "    bicep_parameters_file.write(json.dumps(bicep_parameters))\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 params.json\", \n",
    "                f\"Deployment '{deployment_name}' succeeded\", f\"Deployment '{deployment_name}' failed\")\n",
    "open(\"policy.xml\", 'w').write(policy_template_xml)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='3'></a>\n",
    "### 3️⃣ 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",
    "if output.success and output.json_data:\n",
    "    apim_service_id = utils.get_deployment_output(output, 'apimServiceId', 'APIM Service Id')\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",
    "    apim_resource_gateway_url = utils.get_deployment_output(output, 'apimResourceGatewayURL', 'APIM API Gateway URL')\n",
    "    app_insights_name = utils.get_deployment_output(output, 'applicationInsightsName', 'Application Insights Name')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='requests'></a>\n",
    "### 🧪 Test the API using a direct HTTP call\n",
    "\n",
    "Tip: Use the [tracing tool](../../tools/tracing.ipynb) to track the behavior and troubleshoot the [policy](policy.xml)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load ../../shared/snippets/api-http-requests.py\n",
    "import json\n",
    "import requests\n",
    "import time\n",
    "\n",
    "runs = 10\n",
    "sleep_time_ms = 100\n",
    "url = f\"{apim_resource_gateway_url}/openai/deployments/{openai_deployment_name}/chat/completions?api-version={openai_api_version}\"\n",
    "api_runs = []\n",
    "messages = {\"messages\": [\n",
    "    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "]}\n",
    "\n",
    "# Initialize a session for connection pooling\n",
    "session = requests.Session()\n",
    "# Set default headers\n",
    "session.headers.update({\n",
    "    'api-key': apim_subscription1_key,\n",
    "    'x-user-id': 'alex'\n",
    "})\n",
    "\n",
    "try:\n",
    "    for i in range(runs):\n",
    "        print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "        start_time = time.time()\n",
    "        response = session.post(url, json = messages)\n",
    "        response_time = time.time() - start_time\n",
    "        print(f\"⌚ {response_time:.2f} seconds\")\n",
    "\n",
    "        # Check the response status code and apply formatting\n",
    "        if 200 <= response.status_code < 300:\n",
    "            status_code_str = f\"\\x1b[1;32m{response.status_code} - {response.reason}\\x1b[0m\" # Bold and green\n",
    "        elif response.status_code >= 400:\n",
    "            status_code_str = f\"\\x1b[1;31m{response.status_code} - {response.reason}\\x1b[0m\" # Bold and red\n",
    "        else:\n",
    "            status_code_str = str(response.status_code)  # No formatting\n",
    "\n",
    "        # Print the response status with the appropriate formatting\n",
    "        print(f\"Response status: {status_code_str}\")\n",
    "        print(f\"Response headers: {json.dumps(dict(response.headers), indent = 4)}\")\n",
    "\n",
    "        if \"x-ms-region\" in response.headers:\n",
    "            print(f\"x-ms-region: \\x1b[1;32m{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, response.headers.get(\"x-ms-region\")))\n",
    "\n",
    "        if (response.status_code == 200):\n",
    "            data = json.loads(response.text)\n",
    "            print(f\"Token usage: {json.dumps(dict(data.get(\"usage\")), indent = 4)}\\n\")\n",
    "            print(f\"💬 {data.get(\"choices\")[0].get(\"message\").get(\"content\")}\\n\")\n",
    "        else:\n",
    "            print(response.text)\n",
    "\n",
    "        time.sleep(sleep_time_ms/1000)\n",
    "finally:\n",
    "    # Close the session to release the connection\n",
    "    session.close()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='sdk'></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 = 10\n",
    "sleep_time_ms = 100\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",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "]\n",
    "\n",
    "for i in range(runs):\n",
    "    print(f\"▶️ Run {i+1}/{runs}:\")\n",
    "\n",
    "    for j in range(0, 3):\n",
    "        response = clients[j].chat.completions.create(model = openai_model_name, messages = messages, extra_headers={\"x-user-id\": \"alex\"}) # type: ignore\n",
    "        print(f\"💬 Subscription {j+1}: {response.choices[0].message.content}\")\n",
    "\n",
    "    print()\n",
    "\n",
    "    time.sleep(sleep_time_ms/1000)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='kql'></a>\n",
    "### 🔍 Analyze Application Insights custom metrics with a KQL query\n",
    "\n",
    "With this query you can get the custom metrics that were emitted by Azure APIM. Note that it may take a few minutes for data to become available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "metadata": {}
   },
   "outputs": [],
   "source": [
    "# type: ignore\n",
    "\n",
    "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",
    "# If the query fails, replace {resource_group_name} with the name of the resource group where the Application Insights resource is deployed.\n",
    "result_stdout = ! az monitor app-insights query --app {app_insights_name} -g {resource_group_name} --analytics-query {query}\n",
    "print(result_stdout)\n",
    "result = json.loads(result_stdout.n)\n",
    "\n",
    "table = result.get('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\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='plot'></a>\n",
    "### 🔍 Plot the custom metrics results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the results\n",
    "import pandas as pd\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": [
    "<a id='portal'></a>\n",
    "### 🔍 See the metrics on the Azure Portal\n",
    "\n",
    "Open the Application Insights resource, navigate to the Metrics blade, then select the defined namespace (openai). Choose the metric \"Total Tokens\" with a Sum aggregation. Then, apply splitting by 'Subscription Id' to view values for each dimension.\n",
    "\n",
    "![result](result.png)\n"
   ]
  },
  {
   "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": "Python 3",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}