sdk/python/generative-ai/rag/notebooks/db_copilot_with_rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "345e6aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install azure-ai-ml\n",
    "%pip install -U 'azureml-rag[faiss]>=0.1.11'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "038912d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If `import win32file` fails with a DLL error then run the following and restart kernel:\n",
    "# %pip uninstall -y pywin32\n",
    "# %conda install -y --force-reinstall pywin32"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28b26c69",
   "metadata": {},
   "source": [
    "# Create a FAISS based Vector Index for DBCopilot with AzureML\n",
    "We'll walk through setting up an AzureML Pipeline which grounding a DataBase into a LangChain-compatible FAISS Vector Index and create the promptflow to consume this index to serve as a DBCopilot chatbot."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1a99955c",
   "metadata": {},
   "source": [
    "## Get client for AzureML Workspace\n",
    "\n",
    "The workspace is the top-level resource for Azure Machine Learning, providing a centralized place to work with all the artifacts you create when you use Azure Machine Learning. In this section we will connect to the workspace in which the job will be run.\n",
    "\n",
    "Enter your Workspace details below, running this still will write a `workspace.json` file to the current folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63178816",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile workspace.json\n",
    "{\n",
    "    \"subscription_id\": \"<subscription_id>\",\n",
    "    \"resource_group\": \"<resource_group_name>\",\n",
    "    \"workspace_name\": \"<workspace_name>\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb5c8f0",
   "metadata": {},
   "source": [
    "`MLClient` is how you interact with AzureML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1af37c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "from azure.ai.ml import MLClient\n",
    "from azureml.core import Workspace\n",
    "\n",
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    # Check if given credential can get token successfully.\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    # Fall back to InteractiveBrowserCredential in case DefaultAzureCredential not work\n",
    "    credential = InteractiveBrowserCredential()\n",
    "\n",
    "try:\n",
    "    ml_client = MLClient.from_config(credential=credential, path=\"workspace.json\")\n",
    "except Exception as ex:\n",
    "    raise Exception(\n",
    "        \"Failed to create MLClient from config file. Please modify and then run the above cell with your AzureML Workspace details.\"\n",
    "    ) from ex\n",
    "\n",
    "ws = Workspace(\n",
    "    subscription_id=ml_client.subscription_id,\n",
    "    resource_group=ml_client.resource_group_name,\n",
    "    workspace_name=ml_client.workspace_name,\n",
    ")\n",
    "print(ml_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade6080e",
   "metadata": {},
   "source": [
    "## Create a SQL Datastore\n",
    "1. Go to workspace in Azure Portal\n",
    "2. Click **Data** -> **Datastore** -> **+ Create**\n",
    "3. Fill in the datastore form\n",
    "\n",
    "| Field | Value |\n",
    "| --- | --- |\n",
    "| Datastore name | the name for the datastore |\n",
    "| Datastore type | Azure SQL Database |\n",
    "| Account information| could be found in the Azure SQL Database overview page |\n",
    "| Authentication type | SQL Authentication/ Service principal |\n",
    "| Authentication info | SQL Authentication: username and password; Service principal: tenant id, client id, client secret |"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "30906d39",
   "metadata": {},
   "source": [
    "## Azure OpenAI\n",
    "\n",
    "We recommend using gpt-35-turbo model to get good quality QAs. [Follow these instructions](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal) to setup an Azure OpenAI Instance and deploy the model. Once you have the model deployed in AOAI you can specify your Model name and Deployment name below.\n",
    "\n",
    "We will use the automatically created `Default_AzureOpenAI` connection, change `aoai_connection_name` to use your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab3f1c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "aoai_connection_name = \"Default_AzureOpenAI\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129ac0b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.utils.connections import (\n",
    "    get_connection_by_name_v2,\n",
    "    create_connection_v2,\n",
    ")\n",
    "\n",
    "try:\n",
    "    aoai_connection = get_connection_by_name_v2(ws, aoai_connection_name)\n",
    "except Exception as ex:\n",
    "    # Create New Connection\n",
    "    # Modify the details below to match the `Endpoint` and API key of your AOAI resource, these details can be found in Azure Portal\n",
    "    raise RuntimeError(\n",
    "        \"Have you entered your AOAI resource details below? If so, delete me!\"\n",
    "    )\n",
    "    aoai_connection = create_connection_v2(\n",
    "        workspace=ws,\n",
    "        name=aoai_connection,\n",
    "        category=\"AzureOpenAI\",\n",
    "        # 'Endpoint' from Azure OpenAI resource overview\n",
    "        target=\"https://<endpoint_name>.openai.azure.com/\",\n",
    "        auth_type=\"ApiKey\",\n",
    "        credentials={\n",
    "            # Either `Key` from the `Keys and Endpoint` tab of your Azure OpenAI resource, will be stored in your Workspace associated Azure Key Vault.\n",
    "            \"key\": \"<api-key>\"\n",
    "        },\n",
    "        metadata={\"ApiType\": \"azure\", \"ApiVersion\": \"2023-05-15\"},\n",
    "    )\n",
    "\n",
    "aoai_connection_id = aoai_connection[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c4ece4",
   "metadata": {},
   "source": [
    "Now that your Workspace has a connection to Azure OpenAI we will make sure the `gpt-35-turbo` model has been deployed ready for inference. This cell will fail if there is not deployment for the embeddings model, [follow these instructions](https://learn.microsoft.com/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal#deploy-a-model) to deploy a model with Azure OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c045de88",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.utils.deployment import infer_deployment\n",
    "\n",
    "aoai_embedding_model_name = \"text-embedding-ada-002\"\n",
    "try:\n",
    "    aoai_embedding_deployment_name = infer_deployment(\n",
    "        aoai_connection, aoai_embedding_model_name\n",
    "    )\n",
    "    print(\n",
    "        f\"Deployment name in AOAI workspace for model '{aoai_embedding_model_name}' is '{aoai_embedding_deployment_name}'\"\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Deployment name in AOAI workspace for model '{model_name}' is not found.\")\n",
    "    print(\n",
    "        f\"Please create a deployment for this model by following the deploy instructions on the resource page for '{aoai_connection['properties']['target']}' in Azure Portal.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e5cc13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.utils.deployment import infer_deployment\n",
    "\n",
    "aoai_completion_model_name = \"gpt-35-turbo\"\n",
    "\n",
    "try:\n",
    "    aoai_completion_deployment_name = infer_deployment(\n",
    "        aoai_connection, aoai_completion_model_name\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(\n",
    "        f\"Deployment name in AOAI workspace for model '{aoai_completion_model_name}' is not found.\"\n",
    "    )\n",
    "    print(\n",
    "        f\"Please create a deployment for this model by following the deploy instructions on the resource page for '{aoai_connection['properties']['target']}' in Azure Portal.\"\n",
    "    )\n",
    "\n",
    "print(\n",
    "    f\"Deployment name in AOAI workspace for model '{aoai_completion_model_name}' is '{aoai_completion_deployment_name}'\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff40ff04",
   "metadata": {},
   "source": [
    "Finally we will combine the deployment and model information into a uri form which the AzureML embeddings components expect as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dfe810e",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_completion_config = f'{{\"type\":\"azure_open_ai\",\"model_name\":\"{aoai_completion_model_name}\",\"deployment_name\":\"{aoai_completion_deployment_name}\",\"temperature\":0,\"max_tokens\":\"1500\"}}'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "56878876",
   "metadata": {},
   "source": [
    "### Setup Pipeline\n",
    "\n",
    "The Components are published to a [Registry](https://learn.microsoft.com/azure/machine-learning/how-to-manage-registries?view=azureml-api-2&tabs=cli), `azureml`, which should have access to by default, it can be accessed from any Workspace.\n",
    "In the below cell we get the Component Definitions from the `azureml` registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a3752a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_registry = MLClient(credential=credential, registry_name=\"azureml\")\n",
    "\n",
    "db_copilot_component = ml_registry.components.get(\n",
    "    \"llm_ingest_db_to_faiss\", label=\"latest\"\n",
    ")\n",
    "\n",
    "print(db_copilot_component)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "209700d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.dsl import pipeline\n",
    "\n",
    "\n",
    "@pipeline(name=f\"db_copilot_vector_pipeline_faiss\", default_compute=\"serverless\")\n",
    "def db_copilot_vector_pipeline_faiss(\n",
    "    aoai_connection: str,\n",
    "    db_datastore: str,\n",
    "    embeddings_model: str,\n",
    "    chat_aoai_deployment_name: str,\n",
    "    embedding_aoai_deployment_name: str,\n",
    "    mlindex_dataset_name: str,\n",
    "    selected_tables: str = None,\n",
    "    max_sampling_rows: int = 3,\n",
    "):\n",
    "    db_copilot_component(\n",
    "        db_datastore=db_datastore,\n",
    "        embeddings_model=embeddings_model,\n",
    "        chat_aoai_deployment_name=chat_aoai_deployment_name,\n",
    "        embedding_aoai_deployment_name=embedding_aoai_deployment_name,\n",
    "        embeddings_dataset_name=mlindex_dataset_name,\n",
    "        embedding_connection=aoai_connection,\n",
    "        llm_connection=aoai_connection,\n",
    "        selected_tables=selected_tables,\n",
    "        max_sampling_rows=max_sampling_rows,\n",
    "    )\n",
    "    return {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6682ea80",
   "metadata": {},
   "outputs": [],
   "source": [
    "datastore_name = \"<test_db_datastore_name>\"\n",
    "aoai_embedding_model_name = \"text-embedding-ada-002\"\n",
    "asset_name = \"llm_index_db_dataset\"\n",
    "pipeline_job = db_copilot_vector_pipeline_faiss(\n",
    "    aoai_connection=aoai_connection_id,\n",
    "    db_datastore=f\"azureml://datastores/{datastore_name}\",\n",
    "    embeddings_model=f\"azure_open_ai://deployment/{aoai_embedding_deployment_name}/model/{aoai_completion_model_name}\",\n",
    "    chat_aoai_deployment_name=aoai_completion_deployment_name,\n",
    "    embedding_aoai_deployment_name=aoai_embedding_deployment_name,\n",
    "    mlindex_dataset_name=asset_name,\n",
    "    selected_tables='[\"[dbo].[jobs]\"]',\n",
    "    max_sampling_rows=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b00fd418",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are added so that in progress index generations can be listed in UI, this tagging is done automatically by UI.\n",
    "pipeline_job.properties[\"azureml.mlIndexAssetName\"] = asset_name\n",
    "pipeline_job.properties[\"azureml.mlIndexAssetKind\"] = \"faiss\"\n",
    "pipeline_job.properties[\"azureml.mlIndexAssetSource\"] = \"Database\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52d9b2be",
   "metadata": {},
   "outputs": [],
   "source": [
    "running_pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline_job, experiment_name=\"db_copilot_pipeline\"\n",
    ")\n",
    "running_pipeline_job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71f975c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client.jobs.stream(running_pipeline_job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d098fccd",
   "metadata": {},
   "source": [
    "## Use DBCopilot with Promptflow\n",
    "After the pipeline complete, it will create a promptflow which could be used to chat with the db."
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}