labs/response-streaming/response-streaming.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# APIM ❤️ OpenAI\n",
    "\n",
    "## Response streaming lab\n",
    "![flow](../../images/response-streaming.gif)\n",
    "\n",
    "Playground to try response streaming with APIM and Azure OpenAI endpoints to explore the advantages and shortcomings associated with streaming. To get responses sooner, you can 'stream' the completion as it's being generated. This allows you to start printing or processing the beginning of the completion before the full completion is finished. \n",
    "\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",
    "\n",
    "This lab reuses the [Cookbook - How to stream completions](https://cookbook.openai.com/examples/how_to_stream_completions) published by OpenAI.\n",
    "\n",
    "### Notes\n",
    "- To stream completions, set stream=True when calling the chat completions or completions endpoints. This will return an object that streams back the response as data-only server-sent events. Extract chunks from the delta field rather than the message field. \n",
    "- Streaming responses doesn't includes the usage field to tell how many tokens were consumed.\n",
    "- Streaming in a production application makes it more difficult to moderate the content of the completions, as partial completions may be more difficult to evaluate.\n",
    "\n",
    "### TOC\n",
    "- [0️⃣ Initialize notebook variables](#0)\n",
    "- [1️⃣ Create the Azure Resource Group](#1)\n",
    "- [2️⃣ Create deployment using 🦾 Bicep](#2)\n",
    "- [3️⃣ Get the deployment outputs](#3)\n",
    "- [🧪 Test streaming with direct HTTP call](#requests)\n",
    "- [🧪 Test without streaming using the Azure OpenAI Python SDK](#sdk1)\n",
    "- [🧪 Test with streaming using the Azure OpenAI Python SDK](#sdk2)\n",
    "- [🗑️ Clean up resources](#clean)\n",
    "\n",
    "### Prerequisites\n",
    "- [Python 3.9 or later version](https://www.python.org/) 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",
    "- [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",
    "- The ```mock_webapps``` variable sets the list of deployed Web Apps for the mocking functionality. Clean the ```openai_resources``` list to simulate the OpenAI behaviour with the mocking service.\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": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import datetime\n",
    "import requests\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",
    "apim_resource_name = \"apim\"\n",
    "apim_resource_location = \"westeurope\"\n",
    "apim_resource_sku = \"Basicv2\"\n",
    "openai_resources = [ {\"name\": \"openai1\", \"location\": \"swedencentral\"}  ] # list of OpenAI resources to deploy. Clear this list to use only the mock resources\n",
    "openai_resources_sku = \"S0\"\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",
    "openai_specification_url='https://raw.githubusercontent.com/Azure/azure-rest-api-specs/main/specification/cognitiveservices/data-plane/AzureOpenAI/inference/stable/' + openai_api_version + '/inference.json'\n",
    "openai_backend_pool = \"openai-backend-pool\"\n",
    "mock_backend_pool = \"mock-backend-pool\"\n",
    "mock_webapps = [ {\"name\": \"openaimock1\", \"endpoint\": \"https://openaimock1.azurewebsites.net\"}, {\"name\": \"openaimock2\", \"endpoint\": \"https://openaimock2.azurewebsites.net\"} ]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='1'></a>\n",
    "### 1️⃣ Create the Azure Resource Group\n",
    "All resources deployed in this lab will be created in the specified resource group. Skip this step if you want to use an existing resource group."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resource_group_stdout = ! az group create --name {resource_group_name} --location {resource_group_location}\n",
    "if resource_group_stdout.n.startswith(\"ERROR\"):\n",
    "    print(resource_group_stdout)\n",
    "else:\n",
    "    print(\"✅ Azure Resource Group \", resource_group_name, \" created ⌚ \", datetime.datetime.now().time())"
   ]
  },
  {
   "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. Change the parameters or the [main.bicep](main.bicep) directly to try different configurations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "if len(openai_resources) > 0:\n",
    "    backend_id = openai_backend_pool if len(openai_resources) > 1 else openai_resources[0].get(\"name\")\n",
    "elif len(mock_webapps) > 0:\n",
    "    backend_id = mock_backend_pool if len(mock_backend_pool) > 1 else mock_webapps[0].get(\"name\")\n",
    "\n",
    "with open(\"policy.xml\", 'r') as policy_xml_file:\n",
    "    policy_template_xml = policy_xml_file.read()\n",
    "    policy_xml = policy_template_xml.replace(\"{backend-id}\", backend_id)\n",
    "    policy_xml_file.close()\n",
    "open(\"policy.xml\", 'w').write(policy_xml)\n",
    "\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",
    "    \"mockWebApps\": { \"value\": mock_webapps },\n",
    "    \"mockBackendPoolName\": { \"value\": mock_backend_pool },\n",
    "    \"openAIBackendPoolName\": { \"value\": openai_backend_pool },\n",
    "    \"openAIConfig\": { \"value\": openai_resources },\n",
    "    \"openAIDeploymentName\": { \"value\": openai_deployment_name },\n",
    "    \"openAISku\": { \"value\": openai_resources_sku },\n",
    "    \"openAIModelName\": { \"value\": openai_model_name },\n",
    "    \"openAIModelVersion\": { \"value\": openai_model_version },\n",
    "    \"openAIAPISpecURL\": { \"value\": openai_specification_url },\n",
    "    \"apimResourceName\": { \"value\": apim_resource_name},\n",
    "    \"apimResourceLocation\": { \"value\": apim_resource_location},\n",
    "    \"apimSku\": { \"value\": apim_resource_sku}\n",
    "  }\n",
    "}\n",
    "with open('params.json', 'w') as bicep_parameters_file:\n",
    "    bicep_parameters_file.write(json.dumps(bicep_parameters))\n",
    "\n",
    "! az deployment group create --name {deployment_name} --resource-group {resource_group_name} --template-file \"main.bicep\" --parameters \"params.json\"\n",
    "\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",
    "We are now at the stage where we only need to retrieve the gateway URL and the subscription before we are ready for testing."
   ]
  },
  {
   "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.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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='requests'></a>\n",
    "### 🧪 Test streaming with 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}, json = payload)\n",
    "print(\"status code: \", response.status_code)\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='sdk1'></a>\n",
    "### 🧪 Test without streaming using the Azure OpenAI Python SDK\n",
    "With a typical ChatCompletions API call, the response is first computed and then returned all at once."
   ]
  },
  {
   "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.create(model=openai_model_name, messages=messages, stream=False)\n",
    "response_time = time.time() - start_time\n",
    "print(f\"Full response received {response_time:.2f} seconds after request\")\n",
    "print(f\"Full response received:\\n{response}\")\n",
    "\n",
    "print(f\"Response content in full: {response.choices[0].message.content}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='sdk2'></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.create(model=openai_model_name, messages=messages, stream=True)\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 response:\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='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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}