{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5f84108f",
   "metadata": {},
   "source": [
    "# Microsoft Azure OpenAI On Your Data with Elasticsearch\n",
    "\n",
    "In this notebook we'll use Elasticsearch indices to provide grounding data for queries to Azure OpenAI models using the Azure OpenAI On Your Data service.\n",
    "\n",
    "The Azure OpenAI On Your Data service currently supports three search scenarios for retrieval of documents that will be sent to the LLM for processing:\n",
    "\n",
    "1) full text search\n",
    "2) vector search using Elasticsearch Machine Learning models\n",
    "3) vector search using embeddings generated using Azure OpenAI (Ada).\n",
    "\n",
    "Each of these examples will be covered in the following sections."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26cd732",
   "metadata": {},
   "source": [
    "## Requirements\n",
    "\n",
    "For this example, you will need:\n",
    "* Python 3.6 or later\n",
    "* An Elastic deployment meeting the following criteria: with machine learning node\n",
    "    * API version 8.x\n",
    "    * A machine learning node for following the example for vector search using an Elasticsearch text embedding model\n",
    "* An Azure OpenAI Resource\n",
    "    * At minimum, one chat model should be deployed for your resource to enable chatting about your data.\n",
    "    * Optionally, if you would like to try out vector search using the Azure OpenAI Ada model, you will also need to deploy an Ada model to your resource.  The examples below will assume you are using the model `text-embedding-ada-002`, but can be updated to suit your needs.\n",
    "* The [Elastic Python client](https://www.elastic.co/guide/en/elasticsearch/client/python-api/current/installation.html)\n",
    "* The [OpenAI Python Client](https://platform.openai.com/docs/api-reference/introduction?lang=python)\n",
    "\n",
    "### Create Elastic Deployment\n",
    "\n",
    "If you don't have an Elastic deployment, you can read more about how to get started here in the official [\"Getting Started\" guide for Elastic Cloud](https://www.elastic.co/getting-started).\n",
    "\n",
    "\n",
    "### Configure Azure OpenAI Resource\n",
    "\n",
    "If you don't have an Azure OpenAI resource, detailed information about how to obtain one can be found in the [official documentation](https://learn.microsoft.com/en-us/azure/ai-services/openai/use-your-data-quickstart?tabs=command-line&pivots=programming-language-python) for the Azure OpenAI On Your Data service."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfff496a",
   "metadata": {},
   "source": [
    "## Install packages and initialize environment\n",
    "\n",
    "The first step is to use `pip` to install all of the packages we need to connect to Elasticsearch and Azure OpenAI services."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65760c02",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qU elasticsearch openai==0.28.1 requests"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a5580b",
   "metadata": {},
   "source": [
    "Here we will also define a few helper classes and functions that will be reused throughout the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f940364d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import statistics\n",
    "from io import StringIO\n",
    "from time import sleep, perf_counter\n",
    "from urllib.request import urlopen\n",
    "\n",
    "\n",
    "class SampleCsvDatasets(object):\n",
    "    datasets = {\n",
    "        \"msmarco\": {\n",
    "            \"url\": \"https://github.com/mayya-sharipova/msmarco/raw/main/msmarco-passagetest2019-unique.tsv\",\n",
    "            \"fieldnames\": [\"id\", \"text\"],\n",
    "            \"delimiter\": \"\\t\"\n",
    "        }\n",
    "    }\n",
    "    def __init__(self):\n",
    "        self._values = {}\n",
    "        for name, config in self.datasets.items():\n",
    "            url = config.get(\"url\")\n",
    "            response = urlopen(url)\n",
    "            self._values[name] = self._create_dataframe(name, response.read().decode(\"utf-8\"))\n",
    "    \n",
    "    def _create_dataframe(self, dataset_name, value):\n",
    "        return pd.read_csv(\n",
    "            StringIO(value),\n",
    "            delimiter=self.datasets.get(dataset_name).get(\"delimiter\"),\n",
    "            names=self.datasets.get(dataset_name).get(\"fieldnames\")\n",
    "        )\n",
    "    \n",
    "    def get_dataframe(self, dataset_name, add_aoai_embeddings=False):\n",
    "        return self._values[name]\n",
    "    \n",
    "    def get_dataframe_with_aoai_embeddings(\n",
    "        self,\n",
    "        dataset_name,\n",
    "        embedding_model_deployment_name,\n",
    "        text_column_name=\"text\",\n",
    "        embeddings_column_name=\"text_embedding.aoai_predicted_value\",\n",
    "        model_name=\"text-embedding-ada-002\"\n",
    "    ):\n",
    "        df = self._values[dataset_name]\n",
    "        \n",
    "        # 16 is the current maximum batch size for embeddings, so we'll partition the dataframe by 16 rows.\n",
    "        maximum_batch_size = 16\n",
    "        \n",
    "        # This value can be manipulated if rate limiting errors occur while producing embeddings.\n",
    "        throttling_interval = 0.5\n",
    "        \n",
    "        # Occasionally the server will become overloaded, so retrying can ensure success.\n",
    "        retry_count = 30\n",
    "        \n",
    "        number_of_buckets = math.ceil(len(df) / maximum_batch_size)\n",
    "        partitions = np.array_split(df, number_of_buckets)\n",
    "        embeddings = []\n",
    "        request_times = []\n",
    "        \n",
    "        start = perf_counter()\n",
    "        for partition in partitions:\n",
    "            success = False\n",
    "            for i in range(retry_count):\n",
    "                try:\n",
    "                    request_start = perf_counter()\n",
    "                    text_embedding_response = openai.Embedding.create(\n",
    "                        deployment_id=embedding_model_deployment_name,\n",
    "                        input=partition[text_column_name].tolist(),\n",
    "                        api_version=\"2023-03-15-preview\"\n",
    "                    )\n",
    "                    embeddings.extend([result[\"embedding\"] for result in text_embedding_response[\"data\"]])\n",
    "                    sleep(throttling_interval)\n",
    "                    request_end = perf_counter()\n",
    "                    request_times.append(request_end-request_start)\n",
    "                    success = True\n",
    "\n",
    "                except Exception as e:\n",
    "                    print(f\"Encountered the following exception when creating embeddings (retrying): {str(e)}\")\n",
    "                    sleep(5)\n",
    "            \n",
    "            if not success:\n",
    "                raise Exception(\"Unable to generate embeddings (see previous error output for more information.)\")\n",
    "        \n",
    "        end = perf_counter()\n",
    "        df[text_embeddings_column_name] == embeddings\n",
    "        print(f\"Average request time for {len(partition)} records: {statistics.mean(request_times)}\")\n",
    "        print(f\"Total time: {end-start}\")\n",
    "\n",
    "def create_index(\n",
    "    elastic_client,\n",
    "    index_name: str,\n",
    "    **properties\n",
    "):\n",
    "    mappings = {\n",
    "        \"properties\": properties\n",
    "    }\n",
    "\n",
    "    # Clean up any previously created index with the same name\n",
    "    elastic_client.indices.delete(index=index_name, ignore_unavailable=True)\n",
    "\n",
    "    # Create the index\n",
    "    elastic_client.indices.create(index=index_name, mappings=mappings)\n",
    "\n",
    "\n",
    "def index_data(\n",
    "    elastic_client,\n",
    "    dataset_reader,\n",
    "    index_name\n",
    "):\n",
    "    # Use the bulk API to index data in chunks of 10000 documents\n",
    "    operations_chunks = []\n",
    "    chunk = []\n",
    "    max_operations_chunk_size = 10000\n",
    "    chunk_size = 0\n",
    "\n",
    "    for row in dataset_reader:\n",
    "\n",
    "        if chunk_size < max_operations_chunk_size:\n",
    "            chunk.append({\"index\": {\"_index\": index_name}})\n",
    "            chunk.append(row)\n",
    "            chunk_size += 1\n",
    "            \n",
    "        else:\n",
    "            operations_chunks.append(chunk)\n",
    "            chunk = []\n",
    "            chunk_size = 0\n",
    "\n",
    "    if len(chunk) > 0:\n",
    "        operations_chunks.append(chunk)\n",
    "\n",
    "    for chunk in operations_chunks:\n",
    "        elastic_client.bulk(index=index_name, operations=chunk, refresh=True)\n",
    "\n",
    "def chat_with_my_data(\n",
    "    chat_query,\n",
    "    aoai_deployment_name,\n",
    "    elasticsearch_endpoint,\n",
    "    elasticsearch_api_key,\n",
    "    index_name,\n",
    "    **embedding_config\n",
    "):\n",
    "    elasticsearch_embedding_model = embedding_config.get(\"elasticsearch_embedding_model\", None)\n",
    "    aoai_embedding_model = embedding_config.get(\"aoai_embedding_model\", None)\n",
    "    aoai_embedding_key = embedding_config.get(\"aoai_embedding_key\", None)\n",
    "    aoai_embedding_endpoint = None if not aoai_embedding_model and not aoai_embedding_key \\\n",
    "        else  f\"{openapi.api_base}/openai/deployments/{aoai_embedding_model}/embeddings?api-version=2023-03-15-preview\"\n",
    "    \n",
    "    query_type = \"vector\" if (aoai_embedding_endpoint and aoai_embedding_key) or elasticsearch_embedding_model else \"simple\"\n",
    "    \n",
    "    completion = openai.ChatCompletion.create(\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": chat_query\n",
    "            }\n",
    "        ],\n",
    "        dataSources=[\n",
    "            {\n",
    "                \"type\": \"Elasticsearch\",\n",
    "                \"parameters\": {\n",
    "                    \"endpoint\": elasticsearch_endpoint,\n",
    "                    \"encodedApiKey\": elasticsearch_api_key,\n",
    "                    \"indexName\": index_name,\n",
    "                    \"queryType\": query_type,\n",
    "                    \"embeddingModelId\": elasticsearch_embedding_model,\n",
    "                    \"embeddingEndpoint\":  aoai_embedding_endpoint,\n",
    "                    \"embeddingKey\": aoai_embedding_key\n",
    "                }\n",
    "            }\n",
    "        ],\n",
    "        deployment_id=aoai_deployment_name\n",
    "    )\n",
    "    print(completion.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c2fdccc",
   "metadata": {},
   "source": [
    "Next, let's set some variables for configuring access to the services."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36f42661",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import openai\n",
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "# Elasticsearch Configuration\n",
    "elasticsearch_endpoint = input(\"Elasticsearch endpoint: \")\n",
    "elasticsearch_api_key = getpass.getpass(\"Elasticsearch API Key: \")\n",
    "\n",
    "# Azure OpenAI Configuration\n",
    "openai.api_base = input(\"Azure OpenAI resource endpoint: \")\n",
    "openai.api_key = getpass.getpass(\"Azure OpenAI resource key: \")\n",
    "chat_model_deployment_name = input(\"Azure OpenAI chat model deployment name: \")\n",
    "embedding_model_deployment_name = input(\"Azure OpenAI embedding model deployment name (enter 'None' if not applicable): \" )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5879fadd",
   "metadata": {},
   "source": [
    "We'll also initialize our dataset manager for easy data access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2baf659f",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets = SampleCsvDatasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e04781fb",
   "metadata": {},
   "source": [
    "Next, we'll configure the Elasticsearch client, which we will use to index some data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4050818b",
   "metadata": {},
   "outputs": [],
   "source": [
    "elastic_client = Elasticsearch(\n",
    "    f\"{elasticsearch_endpoint}:443\",\n",
    "    api_key=elasticsearch_api_key\n",
    ")\n",
    "\n",
    "print(elastic_client.info())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b3ea191",
   "metadata": {},
   "source": [
    "Next, we'll configure the OpenAI client to make requests to our specific Azure OpenAI resource endpoint.  We need to create an instance of `requests.adapters.HTTPAdapter` to ensure that the request session URL used by the OpenAI client points to the correct Azure OpenAI resource URL for each request."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13a59211",
   "metadata": {},
   "outputs": [],
   "source": [
    "from requests import Session\n",
    "from requests.adapters import HTTPAdapter\n",
    "\n",
    "class OnYourDataAdapter(HTTPAdapter):\n",
    "    def send(self, request, **kwargs):\n",
    "        request.url = f\"{openai.api_base}/openai/deployments/{chat_model_deployment_name}/extensions/chat/completions?api-version={openai.api_version}\"\n",
    "        return super().send(request, **kwargs)\n",
    "    \n",
    "\n",
    "session = Session()\n",
    "session.mount(\n",
    "    prefix=f\"{openai.api_base}/openai/deployments/{chat_model_deployment_name}\",\n",
    "    adapter=OnYourDataAdapter()\n",
    ")\n",
    "\n",
    "openai.requestssession = session"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "315211dc",
   "metadata": {},
   "source": [
    "## Example 1: Grounding ChatGPT with data retrieved from a full-text search query \n",
    "\n",
    "\n",
    "### Create Elasticsearch index with required mappings\n",
    "\n",
    "We need to create an index which contains some text data for using full-text search for retrieving documents to ground the responses from Azure OpenAI.  First, we'll create the index mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902e3f19",
   "metadata": {},
   "outputs": [],
   "source": [
    "elasticsearch_index_name = \"msmarco-passagetest2019-unique\"\n",
    "mapping_properties = {\n",
    "    \"id\": {\"type\": \"keyword\"},\n",
    "    \"text\": {\"type\": \"text\"}\n",
    "}\n",
    "\n",
    "create_index(\n",
    "    elastic_client,\n",
    "    elasticsearch_index_name,\n",
    "    **mapping_properties\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be0e040e",
   "metadata": {},
   "source": [
    "#### Download dataset\n",
    "\n",
    "For this example, we'll a subset of the MS MARCO Passage Ranking dataset.  We'll use the dataset manager to get a dataframe of our data to use for indexing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee9cab58",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe = datasets.get_dataframe(\"msmarco\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67981bbd",
   "metadata": {},
   "source": [
    "#### Index documents\n",
    "\n",
    "We'll call the helper function `index_data` to handle the indexing.  This function uses the `bulk` API to index data in batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f8eaa4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_data(\n",
    "    elastic_client,\n",
    "    dataframe,\n",
    "    elasticsearch_index_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a072892",
   "metadata": {},
   "source": [
    "### Chat about your dataset\n",
    "\n",
    "Now that we have some data available in our Elasticsearch cluster, we can use Azure OpenAI On Your Data to ask questions about it.  The `chat_with_my_data` function call below will use full-text search by default, since we are not passing in any additional configuration about embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e34f9221",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_query = \"How's the weather in Jamaica?\"\n",
    "chat_with_my_data(\n",
    "    chat_query,\n",
    "    chat_model_deployment_name,\n",
    "    elasticsearch_endpoint,\n",
    "    elasticsearch_api_key,\n",
    "    elasticsearch_index_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdf209c6",
   "metadata": {},
   "source": [
    "## Example 2: Grounding ChatGPT with data retrieved from a kNN search query using an Elasticsearch machine learning model for embeddings\n",
    "\n",
    "The next example will show how to chat with your data using a vector search query using embeddings produced by a model deployed to your Elasticsearch cluster.  This example is adapted from the Elastic Search Labs guide [\"How to deploy NLP: Text Embeddings and Vector Search\"](https://www.elastic.co/search-labs/how-to-deploy-nlp-text-embeddings-and-vector-search), and uses the same dataset that we already downloaded in the previous example.\n",
    "\n",
    "First, we will need to install the `eland` package, which we will use to deploy an embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02fdf05d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install eland[pytorch]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d6c7581",
   "metadata": {},
   "source": [
    "Next, we will deploy the model `msmarco-MiniLM-L-12-v3` model from Hugging Face, which is trained on a superset of the data we have already indexed in the previous example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b31d2e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!eland_import_hub_model \\\n",
    "    --url {elasticsearch_endpoint}:443 \\\n",
    "    --es-api-key {elasticsearch_api_key} \\\n",
    "    --hub-model-id sentence-transformers/msmarco-MiniLM-L-12-v3 \\\n",
    "    --task-type text_embedding \\\n",
    "    --start"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "570de4ea",
   "metadata": {},
   "source": [
    "Next, we can reindex our data with the embeddings from the model to enable kNN search using those embeddings.  In the next cell, we'll create an ingest pipeline to calculate embeddings from the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49109c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "from elasticsearch.client import IngestClient\n",
    "\n",
    "pipeline_id = \"msmarco-minilm-l-12-v3\"\n",
    "ingest_processors = [\n",
    "    {\n",
    "      \"inference\": {\n",
    "        \"model_id\": \"sentence-transformers__msmarco-minilm-l-12-v3\",\n",
    "        \"target_field\": \"text_embedding\",\n",
    "        \"field_map\": {\n",
    "          \"text\": \"text_field\"\n",
    "        }\n",
    "      }\n",
    "    }\n",
    "]\n",
    "\n",
    "on_failure = [\n",
    "    {\n",
    "      \"set\": {\n",
    "        \"description\": \"Index document to 'failed-'\",\n",
    "        \"field\": \"_index\",\n",
    "        \"value\": \"failed-{{{_index}}}\"\n",
    "      }\n",
    "    },\n",
    "    {\n",
    "      \"set\": {\n",
    "        \"description\": \"Set error message\",\n",
    "        \"field\": \"ingest.failure\",\n",
    "        \"value\": \"{{_ingest.on_failure_message}}}\"\n",
    "      }\n",
    "    }\n",
    "]\n",
    "\n",
    "ingest_client = IngestClient(elastic_client)\n",
    "\n",
    "ingest_client.put_pipeline(\n",
    "    id=pipeline_id,\n",
    "    on_failure=on_failure,\n",
    "    processors=ingest_processors\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63da95e8",
   "metadata": {},
   "source": [
    "Then, we'll reindex the `msmarco-passagetest2019-unique` index using the pipeline we just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140fc17d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the new index\n",
    "elasticsearch_embeddings_index_name = f\"msmarco-passagetest2019-unique-{pipeline_id}\"\n",
    "mapping_properties = {\n",
    "    \"text_embedding.predicted_value\": {\n",
    "        \"type\": \"dense_vector\",\n",
    "        \"dims\": 384,\n",
    "        \"index\": True,\n",
    "        \"similarity\": \"cosine\"\n",
    "      },\n",
    "      \"text\": {\n",
    "        \"type\": \"text\"\n",
    "      }\n",
    "}\n",
    "\n",
    "create_index(\n",
    "    elastic_client,\n",
    "    elasticsearch_embeddings_index_name,\n",
    "    **mapping_properties)\n",
    "\n",
    "# Reindex source data into target index\n",
    "reindex_source = {\n",
    "    \"index\": elasticsearch_index_name\n",
    "}\n",
    "reindex_dest = {\n",
    "    \"index\": elasticsearch_embeddings_index_name,\n",
    "    \"pipeline\": pipeline_id\n",
    "}\n",
    "\n",
    "response = elastic_client.reindex(\n",
    "    source=reindex_source,\n",
    "    dest=reindex_dest,\n",
    "    wait_for_completion=False\n",
    ")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe3900f7",
   "metadata": {},
   "source": [
    "The reindexing operation may take some time, so we can poll for the task status before moving to the next step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5bf795",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "\n",
    "is_completed = False\n",
    "while not is_completed:\n",
    "    sleep(1)\n",
    "    task = elastic_client.tasks.get(task_id=response.get(\"task\"))\n",
    "    is_completed = task.get(\"completed\")\n",
    "    \n",
    "print(\"Reindexing is complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "648d6e4a",
   "metadata": {},
   "source": [
    "Now that we have our embeddings index ready to go, we can use this new index with embeddings with the Azure OpenAI On Your Data service.  Note that this call to `chat_with_my_data` passes in some configuration for our embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818e393c",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_query = \"What's the weather like in Jamaica?\"\n",
    "\n",
    "chat_with_my_data(\n",
    "    chat_query,\n",
    "    chat_model_deployment_name,\n",
    "    elasticsearch_endpoint,\n",
    "    elasticsearch_api_key,\n",
    "    elasticsearch_embeddings_index_name,\n",
    "    elasticsearch_embedding_model=\"sentence-transformers__msmarco-minilm-l-12-v3\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2125fbee",
   "metadata": {},
   "source": [
    "## Example 3: Grounding ChatGPT with data retrieved from a kNN search query using Azure OpenAI embeddings\n",
    "\n",
    "We also have the option to create embeddings to use for kNN search using Azure OpenAI.  \n",
    "\n",
    "For this example, in addition to the chat model deployment used in the previous two examples for chatting with your data, you will also need an embedding model deployment added to your Azure OpenAI resource.  We will first use this embedding model to generate embeddings on your data, and then later refer to this deployment in our requests to the Azure OpenAI On Your Data service, which will call the model to generate the embeddings for your query as a part of the chat request.\n",
    "\n",
    "First, we will create a new index with the correct mappings for embeddings produced using the Azure OpenAI Ada embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5df5e380",
   "metadata": {},
   "outputs": [],
   "source": [
    "aoai_embeddings_index_name = f\"msmarco-passagetest2019-unique-{embedding_model_deployment_name}\"\n",
    "mapping_properties = {\n",
    "    \"id\": {\"type\": \"keyword\"},\n",
    "    \"text\": {\"type\": \"text\"},\n",
    "    \"text_embedding.aoai_predicted_value\": {\n",
    "        \"type\": \"dense_vector\",\n",
    "        \"dims\": 1536,\n",
    "        \"index\": True,\n",
    "        \"similarity\": \"cosine\"\n",
    "    }\n",
    "}\n",
    "\n",
    "create_index(\n",
    "    elastic_client,\n",
    "    aoai_embeddings_index_name,\n",
    "    **mapping_properties\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7882d2ac",
   "metadata": {},
   "source": [
    "Next, we will calculate embeddings for our dataset.  Note that this step may take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9091deb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe = datasets.get_dataframe_with_aoai_embeddings(\"msmarco\", embedding_model_deployment_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44b2b5e2",
   "metadata": {},
   "source": [
    "Then, we will index the data with embeddings we just calculated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9767db99",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_data(\n",
    "    elastic_client,\n",
    "    dataframe,\n",
    "    aoai_embeddings_index_name,\n",
    "    add_aoai_embeddings=True,\n",
    "    aoai_embedding_model_deployment=embedding_model_deployment_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c505efc9",
   "metadata": {},
   "source": [
    "Now that our index has been created with Azure Open AI embeddings, we can use the new index to chat with our data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea9580a",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_query = \"What's the weather like in Jamaica?\"\n",
    "chat_with_my_data(\n",
    "    chat_query,\n",
    "    chat_model_deployment_name,\n",
    "    elasticsearch_endpoint,\n",
    "    elasticsearch_api_key,\n",
    "    aoai_embeddings_index_name,\n",
    "    aoai_embedding_model=embedding_model_deployment_name,\n",
    "    aoai_embedding_key=openai.api_key\n",
    ")"
   ]
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}