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",
    "### 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 = \"uksouth\"\n",
    "\n",
    "apim_sku = 'Basicv2'\n",
    "\n",
    "openai_resources = [ {\"name\": \"openai1\", \"location\": \"uksouth\"}]\n",
    "\n",
    "\n",
    "openai_model_name = \"gpt-4o-mini\"\n",
    "openai_model_version = \"2024-07-18\"\n",
    "openai_model_sku = \"GlobalStandard\"\n",
    "openai_deployment_name = \"gpt-4o-mini\"\n",
    "openai_api_version = \"2024-10-21\"\n",
    "\n",
    "utils.print_ok('Notebook initialized')"
   ]
  },
  {
   "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",
    "\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": [
    "<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(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",
    "# Write the parameters to the params.json 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\")"
   ]
  },
  {
   "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",
    "\n",
    "if output.success and output.json_data:\n",
    "    apim_service_id = utils.get_deployment_output(output, 'apimServiceId', 'APIM Service Id')\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='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": [
    "import json, requests, 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",
    "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",
    "api_runs = []\n",
    "\n",
    "# Initialize a session for connection pooling and set any default headers\n",
    "session = requests.Session()\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",
    "        utils.print_response_code(response)\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(f\"{response.text}\\n\")\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",
    "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(\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",
    "        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": [
    "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",
    "\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 matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.figsize'] = [15, 7]\n",
    "if df.empty:\n",
    "    print(\"No data to plot\")\n",
    "else:\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",
    "    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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}