sdk/python/generative-ai/rag/notebooks/faiss/scheduled_update_faiss_index.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "%pip install azure-ai-ml\n",
    "%pip install -U 'azureml-rag[faiss,hugging_face]>=0.1.11'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Update a FAISS based Vector Index on a Schedule\n",
    "\n",
    "We'll walk through setting up an AzureML Pipeline which runs on schedule to continuously load data, processes the data into chunks, embeds the chunks and creates updated langchain compatible FAISS Vector Index which can be loaded and used anywhere."
   ]
  },
  {
   "cell_type": "markdown",
   "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",
    "If you don't have a Workspace and want to create and Index locally see [here to create one](https://learn.microsoft.com/azure/machine-learning/quickstart-create-resources?view=azureml-api-2).\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,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "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",
   "metadata": {},
   "source": [
    "`MLClient` is how you interact with AzureML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "    # ml_client = MLClient(\n",
    "    #     credential=credential,\n",
    "    #     subscription_id=\"\",\n",
    "    #     resource_group_name=\"\",\n",
    "    #     workspace_name=\"\"\n",
    "    # )\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",
   "metadata": {},
   "source": [
    "## Which Embeddings Model to use?\n",
    "\n",
    "There are currently two supported Embedding options: OpenAI's `text-embedding-ada-002` embedding model or HuggingFace embedding models. Here are some factors that might influence your decision:\n",
    "\n",
    "### OpenAI\n",
    "\n",
    "OpenAI has [great documentation](https://platform.openai.com/docs/guides/embeddings) on their Embeddings model `text-embedding-ada-002`, it can handle up to 8191 tokens and can be accessed using [Azure OpenAI](https://learn.microsoft.com/azure/cognitive-services/openai/concepts/models#embeddings-models) or OpenAI directly.\n",
    "If you have an existing Azure OpenAI Instance you can connect it to AzureML, if you don't AzureML provisions a default one for you called `Default_AzureOpenAI`.\n",
    "The main limitation when using `text-embedding-ada-002` is cost/quota available for the model. Otherwise it provides high quality embeddings across a wide array of text domains while being simple to use.\n",
    "\n",
    "### HuggingFace\n",
    "\n",
    "HuggingFace hosts many different models capable of embedding text into single-dimensional vectors. The [MTEB Leaderboard](https://huggingface.co/spaces/mteb/leaderboard) ranks the performance of embeddings models on a few axis, not all models ranked can be run locally (e.g. `text-embedding-ada-002` is on the list), though many can and there is a range of larger and smaller models. When embedding with HuggingFace the model is loaded locally for inference, this will potentially impact your choice of compute resources.\n",
    "\n",
    "**NOTE:** The default PromptFlow Runtime does not come with HuggingFace model dependencies installed, Indexes created using HuggingFace embeddings will not work in PromptFlow by default. **Pick OpenAI if you want to use PromptFlow**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run the cells under _either_ heading (OpenAI or HuggingFace) to use the respective embedding model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### OpenAI\n",
    "\n",
    "We can use the automatically created `Default_AzureOpenAI` connection.\n",
    "\n",
    "If you would rather use an existing Azure OpenAI connection then change `aoai_connection_name` below.\n",
    "If you would rather use an existing Azure OpenAI resource, but don't have a connection created, modify `aoai_connection_name` and the details under the `# Create New Connection` code comment, or navigate the PromptFlow section in your AzureML Workspace and use the Connections create UI flow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aoai_connection_name = \"Default_AzureOpenAI\"\n",
    "aoai_connection_id = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we get the connection from the Workspace and save its `id` so we can reference it later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "metadata": {},
   "source": [
    "Now that your Workspace has a connection to Azure OpenAI we will make sure the `text-embedding-ada-002` 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,
   "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": "markdown",
   "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,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_model_uri = f\"azure_open_ai://deployment/{aoai_embedding_deployment_name}/model/{aoai_embedding_model_name}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### HuggingFace\n",
    "\n",
    "AzureMLs default model from HuggingFace is `all-mpnet-base-v2`, it can be run by most laptops. Any `sentence-transformer` models should be supported, you can learn more about `sentence-transformers` [here](https://huggingface.co/sentence-transformers)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings_model_uri = \"hugging_face://model/sentence-transformers/all-mpnet-base-v2\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Pipeline to process data into Index\n",
    "\n",
    "AzureML [Pipelines](https://learn.microsoft.com/azure/machine-learning/concept-ml-pipelines?view=azureml-api-2) connect together multiple [Components](https://learn.microsoft.com/azure/machine-learning/concept-component?view=azureml-api-2). Each Component defines inputs, code that consumes the inputs and outputs produced from the code. Pipelines themselves can have inputs, and outputs produced by connecting together individual sub Components.\n",
    "To process your data for embedding and indexing we will chain together multiple components each performing their own step of the workflow.\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,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_registry = MLClient(credential=credential, registry_name=\"azureml\")\n",
    "\n",
    "# Clones git repository to output folder of pipeline, by default this will be on the default Workspace Datastore `workspaceblobstore`\n",
    "git_clone_component = ml_registry.components.get(\"llm_rag_git_clone\", label=\"latest\")\n",
    "# Walks input folder according to provided glob pattern (all files by default: '**/*') and attempts to open them, extract text chunks and further chunk if necessary to fir within provided `chunk_size`.\n",
    "crack_and_chunk_component = ml_registry.components.get(\n",
    "    \"llm_rag_crack_and_chunk\", label=\"latest\"\n",
    ")\n",
    "# Reads input folder of files containing chunks and their metadata as batches, in parallel, and generates embeddings for each chunk. Output format is produced and loaded by `azureml.rag.embeddings.EmbeddingContainer`.\n",
    "generate_embeddings_component = ml_registry.components.get(\n",
    "    \"llm_rag_generate_embeddings\", label=\"latest\"\n",
    ")\n",
    "# Reads input folder produced by `azureml.rag.embeddings.EmbeddingsContainer.save()` and inserts all documents (chunk, metadata, embedding_vector) int a Faiss index and in-memory document store. Writes an MLIndex yaml detailing the index and embeddings model information.\n",
    "create_faiss_index_component = ml_registry.components.get(\n",
    "    \"llm_rag_create_faiss_index\", label=\"latest\"\n",
    ")\n",
    "# Takes a uri to a storage location where an MLIndex yaml is stored and registers it as an MLIndex Data asset in the AzureML Workspace.\n",
    "register_mlindex_component = ml_registry.components.get(\n",
    "    \"llm_rag_register_mlindex_asset\", label=\"latest\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each Component has documentation which provides an overall description of the Components purpose and each of the inputs/outputs.\n",
    "For example we can see understand what `crack_and_chunk` does by inspecting the Component definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(crack_and_chunk_component)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below a Pipeline is built by defining a python function which chains together the above components inputs and outputs. Arguments to the function are inputs to the Pipeline itself and the return value is a dictionary defining the outputs of the Pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import Input, Output\n",
    "from azure.ai.ml.dsl import pipeline\n",
    "from azure.ai.ml.entities._job.pipeline._io import PipelineInput\n",
    "from typing import Optional\n",
    "\n",
    "\n",
    "def use_automatic_compute(component, instance_count=1, instance_type=\"Standard_E8s_v3\"):\n",
    "    \"\"\"Configure input `component` to use automatic compute with `instance_count` and `instance_type`.\n",
    "\n",
    "    This avoids the need to provision a compute cluster to run the component.\n",
    "    \"\"\"\n",
    "    component.set_resources(\n",
    "        instance_count=instance_count,\n",
    "        instance_type=instance_type,\n",
    "        properties={\"compute_specification\": {\"automatic\": True}},\n",
    "    )\n",
    "    return component\n",
    "\n",
    "\n",
    "def optional_pipeline_input_provided(input: Optional[PipelineInput]):\n",
    "    \"\"\"Checks if optional pipeline inputs are provided.\"\"\"\n",
    "    return input is not None and input._data is not None\n",
    "\n",
    "\n",
    "# If you have an existing compute cluster you want to use instead of automatic compute, uncomment the following line, replace `dedicated_cpu_compute` with the name of your cluster.\n",
    "# Also comment out the `component.set_resources` line in `use_automatic_compute` above and the `default_compute='serverless'` line below.\n",
    "# @pipeline(compute=dedicated_cpu_compute)\n",
    "@pipeline(default_compute=\"serverless\")\n",
    "def git_to_faiss(\n",
    "    git_url: str,\n",
    "    embeddings_model: str,\n",
    "    asset_name: str,\n",
    "    branch_name: str = None,\n",
    "    chunk_size: int = 1024,\n",
    "    data_source_glob: str = None,\n",
    "    data_source_url: str = None,\n",
    "    document_path_replacement_regex: str = None,\n",
    "    git_connection_id=None,\n",
    "    aoai_connection_id=None,\n",
    "    embeddings_container=None,\n",
    "):\n",
    "    \"\"\"Pipeline to generate embeddings for a `input_data` source and create a Faiss index.\"\"\"\n",
    "\n",
    "    git_clone = git_clone_component(git_repository=git_url, branch_name=branch_name)\n",
    "    use_automatic_compute(git_clone)\n",
    "    if optional_pipeline_input_provided(git_connection_id):\n",
    "        git_clone.environment_variables[\n",
    "            \"AZUREML_WORKSPACE_CONNECTION_ID_GIT\"\n",
    "        ] = git_connection_id\n",
    "\n",
    "    crack_and_chunk = crack_and_chunk_component(\n",
    "        input_data=git_clone.outputs.output_data,\n",
    "        input_glob=data_source_glob,\n",
    "        chunk_size=chunk_size,\n",
    "        data_source_url=data_source_url,\n",
    "        document_path_replacement_regex=document_path_replacement_regex,\n",
    "    )\n",
    "    use_automatic_compute(crack_and_chunk)\n",
    "\n",
    "    generate_embeddings = generate_embeddings_component(\n",
    "        chunks_source=crack_and_chunk.outputs.output_chunks,\n",
    "        embeddings_container=embeddings_container,\n",
    "        embeddings_model=embeddings_model,\n",
    "    )\n",
    "    use_automatic_compute(generate_embeddings)\n",
    "    if optional_pipeline_input_provided(aoai_connection_id):\n",
    "        generate_embeddings.environment_variables[\n",
    "            \"AZUREML_WORKSPACE_CONNECTION_ID_AOAI\"\n",
    "        ] = aoai_connection_id\n",
    "    if optional_pipeline_input_provided(embeddings_container):\n",
    "        # If provided, `embeddings_container` is expected to be a URI to folder, the folder can be empty.\n",
    "        # Each sub-folder is generated by a `create_embeddings_component` run and can be reused for subsequent embeddings runs.\n",
    "        generate_embeddings.outputs.embeddings = Output(\n",
    "            type=\"uri_folder\", path=f\"{embeddings_container.path}/{{name}}\"\n",
    "        )\n",
    "\n",
    "    create_faiss_index = create_faiss_index_component(\n",
    "        embeddings=generate_embeddings.outputs.embeddings,\n",
    "    )\n",
    "    use_automatic_compute(create_faiss_index)\n",
    "\n",
    "    register_mlindex = register_mlindex_component(\n",
    "        storage_uri=create_faiss_index.outputs.index, asset_name=asset_name\n",
    "    )\n",
    "    use_automatic_compute(register_mlindex)\n",
    "    return {\n",
    "        \"mlindex_asset_uri\": create_faiss_index.outputs.index,\n",
    "        \"mlindex_asset_id\": register_mlindex.outputs.asset_id,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create the Pipeline Job by calling the `@pipeline` annotated function and providing input arguments.\n",
    "`asset_name` will be used when registering the MLIndex Data Asset produced by the `register_mlindex` component in the pipeline. This is how you can refer to the MLIndex within AzureML.\n",
    "\n",
    "`crack_and_chunk` will infer the data_source_url if the input_data is a git repository. If you want to have source links your index that link to somewhere else, refer to how `data_source_url` is handled in [ACS MLIndex Notebook](../azure_cognitive_search/acs_mlindex_with_langchain.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "git_url = \"https://github.com/rust-lang/book.git\"\n",
    "data_source_glob = \"src/**/*\"\n",
    "asset_name = \"rustdocs_aoai_faiss_mlindex\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import Input\n",
    "\n",
    "pipeline_job = git_to_faiss(\n",
    "    git_url=git_url,\n",
    "    data_source_glob=data_source_glob,\n",
    "    # Each run will save latest Embeddings to subfolder under this path, runs will load latest embeddings from container and reuse any unchanged chunk embeddings.\n",
    "    embeddings_model=embeddings_model_uri,\n",
    "    aoai_connection_id=aoai_connection_id,\n",
    "    embeddings_container=Input(\n",
    "        type=\"uri_folder\",\n",
    "        path=f\"azureml://datastores/workspaceblobstore/paths/embeddings/{asset_name}\",\n",
    "    ),\n",
    "    # Name of asset to register MLIndex under\n",
    "    asset_name=asset_name,\n",
    ")\n",
    "\n",
    "# By default AzureML Pipelines will reuse the output of previous component Runs when inputs have not changed.\n",
    "# If you want to rerun the Pipeline every time each time so that any changes to upstream data sources are processed uncomment the below line.\n",
    "# pipeline_job.settings.force_rerun = True # Rerun each time so that git_clone isn't cached, if intent is to ingest latest data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we add some properties to `pipeline_job` which ensure the Index generation progress and final Artifact appear in the PromptFlow Vector Index UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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\"] = \"Git Repository\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Schedule Pipeline\n",
    "\n",
    "The scheduled Pipeline will run once daily at the `scheduled_start_time`, at the end of each run a new version of the MLIndex data asset will be published which will contain any updated documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.constants import TimeZone\n",
    "from azure.ai.ml.entities import JobSchedule, RecurrenceTrigger, RecurrencePattern\n",
    "from datetime import datetime, timedelta\n",
    "\n",
    "\n",
    "schedule_name = \"rustbook_aoai_faiss_mlindex_daily\"\n",
    "\n",
    "# Make sure pipeline runs git_clone every trigger\n",
    "pipeline_job.settings.force_rerun = True\n",
    "\n",
    "schedule_start_time = datetime.utcnow() + timedelta(minutes=1)\n",
    "recurrence_trigger = RecurrenceTrigger(\n",
    "    frequency=\"day\",\n",
    "    interval=1,\n",
    "    # schedule=RecurrencePattern(hours=16, minutes=[15]),\n",
    "    start_time=schedule_start_time,\n",
    "    time_zone=TimeZone.UTC,\n",
    ")\n",
    "\n",
    "job_schedule = JobSchedule(\n",
    "    name=schedule_name,\n",
    "    trigger=recurrence_trigger,\n",
    "    create_job=pipeline_job,\n",
    "    properties={\n",
    "        \"azureml.mlIndexAssetName\": asset_name,\n",
    "        \"azureml.mlIndexAssetKind\": \"faiss\",\n",
    "        \"azureml.mlIndexAssetSource\": \"Git Repository\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enable Schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_schedule_res = ml_client.schedules.begin_create_or_update(\n",
    "    schedule=job_schedule\n",
    ").result()\n",
    "job_schedule_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disable Schedule\n",
    "job_schedule_res = ml_client.schedules.begin_disable(name=schedule_name).result()\n",
    "job_schedule_res.is_enabled"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use Index with langchain\n",
    "\n",
    "The Data Asset produced by the AzureML Pipeline above contains a yaml file named 'MLIndex' which contains all the information needed to use the Azure Cognitive Search index.\n",
    "For instance if an AOAI deployment was used to embed the documents the details of that deployment and a reference to the secret are there.\n",
    "This allows easy loading of the MLIndex into a langchain retriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.mlindex import MLIndex\n",
    "\n",
    "question = \"What makes Rust Enums so awesome?\"\n",
    "\n",
    "retriever = MLIndex(\n",
    "    ml_client.data.get(asset_name, label=\"latest\")\n",
    ").as_langchain_retriever()\n",
    "retriever.get_relevant_documents(question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have not deployed `gpt-35-turbo` on your Azure OpenAI resource the below cell will fail indicated the `API deployment for this resource does not exist`. Follow the previous instructions for deploying `text-embedding-ada-002` to deploy `gpt-35-turbo`, note the chosen deployment name below and use the same or update it if you choose different one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains import RetrievalQA\n",
    "from azureml.rag.models import init_llm, parse_model_uri\n",
    "\n",
    "model_config = parse_model_uri(\n",
    "    \"azure_open_ai://deployment/gpt-35-turbo/model/gpt-35-turbo\"\n",
    ")\n",
    "model_config[\"api_base\"] = aoai_connection[\"properties\"][\"target\"]\n",
    "model_config[\"key\"] = aoai_connection[\"properties\"][\"credentials\"][\"key\"]\n",
    "model_config[\"temperature\"] = 0.3\n",
    "model_config[\"max_retries\"] = 3\n",
    "\n",
    "qa = RetrievalQA.from_chain_type(\n",
    "    llm=init_llm(model_config), chain_type=\"stuff\", retriever=retriever\n",
    ")\n",
    "\n",
    "qa.run(question)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: To control the number of documents returned when searching try getting the the MLIndex `as_langchain_vectorstore()` instead, this implements the `VectorStore` interface which has more parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use MLIndex with PromptFlow\n",
    "\n",
    "To use the MLindex in PromptFlow the asset_id can be used with the `Vector Index Lookup​` Tool. Replace `versions/2` with `versions/latest` to use the latest version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use MLIndex with PromptFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asset_id = f\"azureml:/{ml_client.data.get(asset_name, label='latest').id}\"\n",
    "asset_id.replace(\"resourceGroups\", \"resourcegroups\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the MLindex in PromptFlow the asset_id can be used with the `Vector Index Lookup​` Tool. Replace `versions/2` with `versions/latest` to use the latest version."
   ]
  }
 ],
 "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}