tools/streaming.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# APIM ❤️ OpenAI\n",
    "\n",
    "## Streaming tool\n",
    "\n",
    "Invoke OpenAI API with stream enabled and returns response in chunks.\n",
    "\n",
    "Notes:\n",
    "- Follow the [APIM guidelines for SSE](https://learn.microsoft.com/azure/api-management/how-to-server-sent-events#guidelines-for-sse) to guarantee that your APIM configuration is compatible with streaming.\n",
    "- This tool reuses the [Cookbook - How to stream completions](https://cookbook.openai.com/examples/how_to_stream_completions) published by OpenAI.\n",
    "\n",
    "## TOC\n",
    "- [Initialize notebook variables](#0)\n",
    "- [Get the deployment outputs](#1)\n",
    "- [Get the APIM authorization debug token](#2)\n",
    "- [🧪 Test the API using a direct HTTP call](#requests)\n",
    "- [🔍 Analyze the API trace from direct HTTP call](#trace1)\n",
    "- [🧪 Test with streaming using the Azure OpenAI Python SDK](#sdk)\n",
    "- [🔍 Analyze the API trace from the SDK call](#trace2)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='0'></a>\n",
    "### Initialize notebook variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deployment_name = \"\" # name of the label that you want to use with  this tool (ex: semantic-caching)\n",
    "resource_group_name = f\"lab-{deployment_name}\"\n",
    "openai_deployment_name = \"gpt-35-turbo\"\n",
    "openai_api_version = \"2024-02-01\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='1'></a>\n",
    "### Get the deployment outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deployment_stdout = ! az deployment group show --name {deployment_name} -g {resource_group_name} --query properties.outputs.apimServiceId.value -o tsv\n",
    "apim_service_id = deployment_stdout.n\n",
    "print(\"👉🏻 APIM Service Id: \", apim_service_id)\n",
    "\n",
    "deployment_stdout = ! az deployment group show --name {deployment_name} -g {resource_group_name} --query properties.outputs.apimSubscriptionKey.value -o tsv\n",
    "apim_subscription_key = deployment_stdout.n\n",
    "deployment_stdout = ! az deployment group show --name {deployment_name} -g {resource_group_name} --query properties.outputs.apimResourceGatewayURL.value -o tsv\n",
    "apim_resource_gateway_url = deployment_stdout.n\n",
    "print(\"👉🏻 API Gateway URL: \", apim_resource_gateway_url)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='2'></a>\n",
    "### Get the APIM authorization debug token\n",
    "\n",
    "This token will be used to trace the API request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "import json\n",
    "token = ! az account get-access-token --query accessToken --output tsv\n",
    "\n",
    "request = {\n",
    "    \"credentialsExpireAfter\": \"PT1H\",\n",
    "    \"apiId\": apim_service_id + \"/apis/openai\",\n",
    "    \"purposes\": [\"tracing\"]\n",
    "}\n",
    "url = \"https://management.azure.com\" + apim_service_id + \"/gateways/managed/listDebugCredentials?api-version=2023-05-01-preview\"\n",
    "\n",
    "response = requests.post(url, headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer ' + token.n}, json = request)\n",
    "\n",
    "if (response.status_code == 200):\n",
    "    data = json.loads(response.text)\n",
    "    apim_debug_authorization = data.get(\"token\")\n",
    "else:\n",
    "    print(response.text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='requests'></a>\n",
    "### 🧪 Test the API using a direct HTTP call\n",
    "The Python requests library has support for [streaming](https://docs.python-requests.org/en/latest/user/advanced/#streaming-requests). We use the direct http call to inspect the response headers. The policy is injecting a header that identifies if it's a streaming request.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = apim_resource_gateway_url + \"/openai/deployments/\" + openai_deployment_name + \"/chat/completions?api-version=\" + openai_api_version\n",
    "payload={\"messages\":[\n",
    "    {\"role\": \"system\", \"content\": \"You are a sarcastic, unhelpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": \"Can you tell me the time, please?\"}\n",
    "],\n",
    "\"stream\": True}\n",
    "response = requests.post(url, headers = {'api-key':apim_subscription_key, 'Apim-Debug-Authorization': apim_debug_authorization}, json = payload)\n",
    "print(\"status code: \", response.status_code)\n",
    "trace_id = response.headers.get(\"Apim-Trace-Id\")\n",
    "print(\"Apim-Trace-Id: \", trace_id) # this header will be used to get API trace details\n",
    "print(\"headers \", response.headers)\n",
    "print(\"x-ms-region: \", response.headers.get(\"x-ms-region\")) # this header is useful to determine the region of the backend that served the request\n",
    "print(\"x-ms-stream: \", response.headers.get(\"x-ms-stream\")) # this header is useful to determine if the response is streamed\n",
    "if (response.status_code == 200):\n",
    "    for chunk in response.iter_lines():\n",
    "        print('chunk:', chunk)    \n",
    "else:\n",
    "    print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='trace1'></a>\n",
    "### 🔍 Analyze the API trace from direct HTTP call\n",
    "\n",
    "With the following request we will get the json with the complete trace information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = {\n",
    "    \"traceId\": trace_id\n",
    "}\n",
    "url = \"https://management.azure.com\" + apim_service_id + \"/gateways/managed/listTrace?api-version=2023-05-01-preview\"\n",
    "response = requests.post(url, headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer ' + token.n}, json = request)\n",
    "\n",
    "if (response.status_code == 200):\n",
    "    data = json.loads(response.text)\n",
    "    print(json.dumps(data, indent=4))\n",
    "else:\n",
    "    print(response.text)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='sdk'></a>\n",
    "### 🧪 Test with streaming using the Azure OpenAI Python SDK\n",
    "With a streaming API call, the response is sent back incrementally in chunks via an [event stream](https://developer.mozilla.org/docs/Web/API/Server-sent_events/Using_server-sent_events#event_stream_format). In Python, you can iterate over these events with a for loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from openai import AzureOpenAI\n",
    "messages=[\n",
    "        {'role': 'user', 'content': 'Count to 100, with a comma between each number and no newlines. E.g., 1, 2, 3, ...'}\n",
    "]\n",
    "start_time = time.time()\n",
    "client = AzureOpenAI(\n",
    "    azure_endpoint=apim_resource_gateway_url,\n",
    "    api_key=apim_subscription_key,\n",
    "    api_version=openai_api_version\n",
    ")\n",
    "response = client.chat.completions.with_raw_response.create(model=openai_deployment_name, messages=messages, extra_headers={'Apim-Debug-Authorization': apim_debug_authorization}, stream=True)\n",
    "trace_id = response.headers.get(\"Apim-Trace-Id\")\n",
    "print(\"Apim-Trace-Id: \", trace_id) # this header will be used to get API trace details\n",
    "print(\"headers \", response.headers)\n",
    "print(\"x-ms-region: \", response.headers.get(\"x-ms-region\")) # this header is useful to determine the region of the backend that served the request\n",
    "print(\"x-ms-stream: \", response.headers.get(\"x-ms-stream\")) # this header is useful to determine if the response is streamed\n",
    "\n",
    "completion = response.parse() \n",
    "\n",
    "# create variables to collect the stream of chunks\n",
    "collected_chunks = []\n",
    "collected_messages = []\n",
    "# iterate through the stream of events\n",
    "for chunk in completion:\n",
    "    chunk_time = time.time() - start_time  # calculate the time delay of the chunk\n",
    "    collected_chunks.append(chunk)  # save the event response\n",
    "    if chunk.choices:\n",
    "        chunk_message = chunk.choices[0].delta.content  # extract the message\n",
    "        collected_messages.append(chunk_message)  # save the message\n",
    "        print(f\"Message received {chunk_time:.2f} seconds after request: {chunk_message}\")  # print the delay and text\n",
    "# print the time delay and text received\n",
    "print(f\"Full response received {chunk_time:.2f} seconds after request\")\n",
    "# clean None in collected_messages\n",
    "collected_messages = [m for m in collected_messages if m is not None]\n",
    "full_reply_content = ''.join(collected_messages)\n",
    "print(f\"Full conversation received: {full_reply_content}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='trace2'></a>\n",
    "### 🔍 Analyze the API trace from the SDK call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request = {\n",
    "    \"traceId\": trace_id\n",
    "}\n",
    "url = \"https://management.azure.com\" + apim_service_id + \"/gateways/managed/listTrace?api-version=2023-05-01-preview\"\n",
    "response = requests.post(url, headers = {'Content-Type': 'application/json', 'Authorization': 'Bearer ' + token.n}, json = request)\n",
    "\n",
    "if (response.status_code == 200):\n",
    "    data = json.loads(response.text)\n",
    "    print(json.dumps(data, indent=4))\n",
    "else:\n",
    "    print(response.text)\n"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}