Basic_Samples/AAD_Integration/aad_integration_example_restapi.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Azure OpenAI Service API with AAD authentication\n",
    "\n",
    "The reference for the AOAI REST API `/completions` method is [here](https://docs.microsoft.com/en-us/azure/cognitive-services/openai/reference/api-reference#completions).\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "1. Setup for Azure Active Directory (AAD) authentication.\n",
    "    * See [Setup to use AAD and test with CLI](setup_aad.md).\n",
    "2. A Python environment setup with all the requirements.  \n",
    "    * See the [setup_python_env.md](setup_python_env.md) page for instructions on setting up your environment.  You don't need the `openai` package for this notebook, but you do need the `azure-identity` package.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the API endpoint URL for your team"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import needed libraries\n",
    "import os\n",
    "import requests\n",
    "import json\n",
    "import datetime\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Parameters\n",
    "\n",
    "\n",
    "Here we will load the configurations from _config.json_ file to setup deployment_name, base_url and openai_api_version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load config values\n",
    "with open(r'config.json') as config_file:\n",
    "    config_details = json.load(config_file)\n",
    "\n",
    "# Setting up the deployment name\n",
    "deployment_name = config_details['COMPLETIONS_MODEL']\n",
    "\n",
    "# The base URL for your Azure OpenAI resource. e.g. \"https://<your resource name>.openai.azure.com\"\n",
    "base_url = config_details['OPENAI_API_BASE']\n",
    "\n",
    "# Currently OPENAI API have the following versions available: 2022-12-01\n",
    "openai_api_version = config_details['OPENAI_API_VERSION']\n",
    "\n",
    "# Define the API endpoint URL\n",
    "request_url = base_url + \"/openai/deployments/\"+ deployment_name + \"/completions?api-version=\" + openai_api_version"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up AAD authentication\n",
    "\n",
    "`DefaultAzureCredential` can read your credentials from the environment after doing `az login`. \n",
    "\n",
    "In VS Code, you can use the Azure Account extension to log in with your Azure account.  If you are running this notebook in VS Code, be sure to restart VS Code after you do `az login`.\n",
    "\n",
    "This article gives details on what identity `DefaultAzureCredential` uses: https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python\n",
    "\n",
    "If you get an error similar to the following, you can try using `AzureCliCredential` instead of `DefaultAzureCredential`:\n",
    "\n",
    "```\n",
    "DefaultAzureCredential failed to retrieve a token from the included credentials. Attempted credentials:     EnvironmentCredential: EnvironmentCredential authentication unavailable. Environment variables are not fully configured. Visit https://aka.ms/azsdk/python/identity/environmentcredential/troubleshoot to troubleshoot.this issue.     ManagedIdentityCredential: ManagedIdentityCredential authentication unavailable, no response from the IMDS endpoint.     SharedTokenCacheCredential: Azure Active Directory error '(invalid_scope) AADSTS70011: The provided request must include a 'scope' input parameter. The provided value for the input parameter 'scope' is not valid. The scope https://cognitiveservices.azure.com is not valid. The scope format is invalid. Scope must be in a valid URI form <https://example/scope> or a valid Guid <guid/scope>. \n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.identity import DefaultAzureCredential, AzureCliCredential #DefaultAzureCredential should work but you may need AzureCliCredential to make the authentication work\n",
    "default_credential = AzureCliCredential()\n",
    "#default_credential = DefaultAzureCredential()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define method to get a prompt completion using AAD authentication\n",
    "\n",
    "The `refresh_token` function below is used to get a new token when the current token expires.  The `refresh_token` method is called \n",
    "by the `get_completion` to get the token if it is not already set or if the token has expired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "token = None\n",
    "\n",
    "def refresh_token():\n",
    "    global token\n",
    "    # Check if Azure token is still valid\n",
    "    if not token or datetime.datetime.fromtimestamp(token.expires_on) < datetime.datetime.now():\n",
    "        token = default_credential.get_token(\"https://cognitiveservices.azure.com\")\n",
    "\n",
    "def get_completion(payload):\n",
    "    # Refresh token\n",
    "    refresh_token()\n",
    "    \n",
    "    # Yes this can be optimized to only set Authorization header when token is refreshed :D\n",
    "    headers = {\n",
    "        \"Authorization\": f\"Bearer {token.token}\",\n",
    "        \"Content-Type\": \"application/json\"\n",
    "    }\n",
    "\n",
    "    r = requests.post(\n",
    "        request_url, \n",
    "        headers=headers,\n",
    "        json = payload\n",
    "        )\n",
    "    \n",
    "    return r"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the API\n",
    "\n",
    "The payload parameters in the `get_completion` call are for illustration only and can be changed for your use case as described in the [reference](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status Code: 200\n",
      "Headers: \n",
      "{'Cache-Control': 'no-cache, must-revalidate', 'Content-Length': '1239', 'Content-Type': 'application/json', 'access-control-allow-origin': '*', 'openai-model': 'text-davinci-003', 'apim-request-id': 'be5722a6-ffad-4bcb-b988-d70971ff0c84', 'openai-processing-ms': '4552.6285', 'x-content-type-options': 'nosniff', 'x-accel-buffering': 'no', 'x-request-id': '1f720ba0-6fe8-40e7-a55e-de8cb05f6e2b', 'Strict-Transport-Security': 'max-age=31536000; includeSubDomains; preload', 'x-ms-region': 'East US', 'Date': 'Wed, 29 Mar 2023 11:00:51 GMT'} \n",
      "\n",
      "{'id': 'cmpl-6zNpmCskuKaLc0heqY9DIOnzYAzUz', 'object': 'text_completion', 'created': 1680087646, 'model': 'text-davinci-003', 'choices': [{'text': ' there was a young girl named Alice. She lived in a small village in the middle of the forest. One day, Alice was walking through the woods when she stumbled upon a white rabbit. She followed the rabbit deep into the forest until she came to a strange looking tree. The tree had a door at the bottom with a sign that said \"Drink Me\".\\n\\nAlice was curious, so she opened the door and drank the liquid inside. Suddenly, Alice began to shrink and she soon found herself in a strange world. She encountered many strange creatures, including talking animals and a strange caterpillar who smoked a hookah.\\n\\nAlice eventually found her way back home, but she was forever changed by her experience in Wonderland. She never forgot the strange creatures and adventures she experienced.\\n\\nAlice\\'s journey was filled with lessons about friendship, courage, and the power of imagination. It taught her that nothing is impossible and that life can be filled with exciting possibilities.', 'index': 0, 'finish_reason': 'stop', 'logprobs': None}], 'usage': {'completion_tokens': 195, 'prompt_tokens': 4, 'total_tokens': 199}}\n"
     ]
    }
   ],
   "source": [
    "# Example prompt for request payload\n",
    "prompt = \"Once upon a time\"\n",
    "\n",
    "payload = {\n",
    "            \"prompt\": prompt,\n",
    "            \"max_tokens\": 300,\n",
    "            \"temperature\": 0.7,\n",
    "            \"top_p\": 1\n",
    "          }\n",
    "\n",
    "r = get_completion(payload)\n",
    "\n",
    "# printing the status code & headers\n",
    "print(f'Status Code: {r.status_code}\\nHeaders: \\n{r.headers} \\n') \n",
    "\n",
    "response = r.json()\n",
    "\n",
    "# printing the response\n",
    "print(f'Response:\\n {response}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "5453039e1c3beadcd7c6bc838a2f5d2d849c1bfa9a5a5e848cf27873534572b2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}