02_02_Azure-AI-Search-RAG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a RAG solution using Azure AI Search\n",
    "\n",
    "Steps in this notebook:\n",
    "1. Create an index \n",
    "2. Create a data source\n",
    "3. Create a skillset\n",
    "4. Create an indexer \n",
    "5. Send a query to the search engine to check results\n",
    "6. Send query results to a language model to generate response\n",
    "\n",
    "Note: Steps 1-4: Done during initial setup of Search Index only"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install python-dotenv\n",
    "%pip install azure-core\n",
    "%pip install azure-search-documents\n",
    "%pip install azure-storage-blob\n",
    "%pip install azure-identity\n",
    "%pip install openai\n",
    "%pip install aiohttp\n",
    "%pip install ipywidgets\n",
    "%pip install ipykernel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Azure configurations\n",
    "\n",
    "You always need to run this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "load_dotenv() # take environment variables from .env.\n",
    "\n",
    "azure_openai_endpoint = os.getenv(\"AZURE_OPENAI_ENDPOINT\")\n",
    "azure_openai_key = os.getenv(\"AZURE_OPENAI_KEY\")\n",
    "azure_openai_deployment = os.getenv(\"AZURE_OPENAI_DEPLOYMENT\")\n",
    "azure_openai_embeddings_deployment = os.getenv(\"AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT\")\n",
    "azure_openai_api_version = \"2024-10-01-preview\"\n",
    "azure_search_service_endpoint = os.getenv(\"AZURE_SEARCH_SERVICE_ENDPOINT\")\n",
    "azure_search_service_admin_key = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "azure_search_service_index_name = \"az-search-index-001\"\n",
    "azure_storage_connection_string = os.getenv(\"AZURE_STORAGE_CONNECTION_STRING\")\n",
    "azure_ai_services_key = os.getenv(\"AZURE_AI_MULTISERVICE_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from azure.search.documents.indexes.models import (\n",
    "    SearchField,\n",
    "    SearchFieldDataType,\n",
    "    VectorSearch,\n",
    "    HnswAlgorithmConfiguration,\n",
    "    VectorSearchProfile,\n",
    "    AzureOpenAIVectorizer,\n",
    "    AzureOpenAIVectorizerParameters,\n",
    "    SearchIndex,\n",
    "    SemanticConfiguration,\n",
    "    SemanticPrioritizedFields,\n",
    "    SemanticField,\n",
    "    SemanticSearch,\n",
    "    ScoringProfile,\n",
    "    TagScoringFunction,\n",
    "    TagScoringParameters\n",
    ")\n",
    "\n",
    "# Get credential from Azure AI Search Admin key\n",
    "credential = AzureKeyCredential(azure_search_service_admin_key)\n",
    "\n",
    "# Search index name  \n",
    "index_name = azure_search_service_index_name\n",
    "\n",
    "# Create a Search Index Client\n",
    "index_client = SearchIndexClient(endpoint=azure_search_service_endpoint, credential=credential)\n",
    "\n",
    "# Define the fields collection\n",
    "fields = [\n",
    "    SearchField(name=\"parent_id\", type=SearchFieldDataType.String),  \n",
    "    SearchField(name=\"title\", type=SearchFieldDataType.String),\n",
    "    SearchField(name=\"locations\", type=SearchFieldDataType.Collection(SearchFieldDataType.String), filterable=True),\n",
    "    SearchField(name=\"chunk_id\", type=SearchFieldDataType.String, key=True, sortable=True, filterable=True, facetable=True, analyzer_name=\"keyword\"),  \n",
    "    SearchField(name=\"chunk\", type=SearchFieldDataType.String, sortable=False, filterable=False, facetable=False),  \n",
    "    SearchField(name=\"text_vector\", type=SearchFieldDataType.Collection(SearchFieldDataType.Single), vector_search_dimensions=1024, vector_search_profile_name=\"myHnswProfile\")\n",
    "    ]  \n",
    "  \n",
    "# Configure the vector search configuration  \n",
    "vector_search = VectorSearch(  \n",
    "    algorithms=[  \n",
    "        HnswAlgorithmConfiguration(name=\"myHnsw\"),\n",
    "    ],  \n",
    "    profiles=[  \n",
    "        VectorSearchProfile(  \n",
    "            name=\"myHnswProfile\",  \n",
    "            algorithm_configuration_name=\"myHnsw\",  \n",
    "            vectorizer_name=\"myOpenAI\",  \n",
    "        )\n",
    "    ],  \n",
    "    vectorizers=[   # a vectorizer is software that performs vectorization\n",
    "        AzureOpenAIVectorizer(  \n",
    "            vectorizer_name=\"myOpenAI\",  \n",
    "            kind=\"azureOpenAI\",  \n",
    "            parameters=AzureOpenAIVectorizerParameters(  \n",
    "                resource_url=azure_openai_endpoint,  \n",
    "                deployment_name=azure_openai_embeddings_deployment,\n",
    "                model_name=azure_openai_embeddings_deployment\n",
    "            ),\n",
    "        ),  \n",
    "    ], \n",
    ")  \n",
    "\n",
    "# New semantic configuration\n",
    "semantic_config = SemanticConfiguration(\n",
    "    name=\"my-semantic-config\",\n",
    "    prioritized_fields=SemanticPrioritizedFields(\n",
    "        title_field=SemanticField(field_name=\"title\"),\n",
    "        keywords_fields=[SemanticField(field_name=\"locations\")],\n",
    "        content_fields=[SemanticField(field_name=\"chunk\")]\n",
    "    )\n",
    ")\n",
    "\n",
    "# Create the semantic settings with the configuration\n",
    "semantic_search = SemanticSearch(configurations=[semantic_config])\n",
    "\n",
    "# New scoring profile\n",
    "scoring_profiles = [  \n",
    "    ScoringProfile(  \n",
    "        name=\"my-scoring-profile\",\n",
    "        functions=[\n",
    "            TagScoringFunction(  \n",
    "                field_name=\"locations\",  \n",
    "                boost=5.0,  \n",
    "                parameters=TagScoringParameters(  \n",
    "                    tags_parameter=\"tags\",  \n",
    "                ),  \n",
    "            ) \n",
    "        ]\n",
    "    )\n",
    "]\n",
    "\n",
    "# Create the search index\n",
    "index = SearchIndex(name=index_name, \n",
    "                    fields=fields, \n",
    "                    vector_search=vector_search,\n",
    "                    semantic_search=semantic_search,\n",
    "                    scoring_profiles=scoring_profiles)  \n",
    "result = index_client.create_or_update_index(index)  \n",
    "print(f\"{result.name} created\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check Azure Storage access\n",
    "Verify access to storage and print out the documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.storage.blob import BlobServiceClient\n",
    "\n",
    "# Initialize the BlobServiceClient with the connection string\n",
    "blob_service_client = BlobServiceClient.from_connection_string(azure_storage_connection_string)\n",
    "\n",
    "# Get the container client\n",
    "container_client = blob_service_client.get_container_client(\"nasa-ebooks-pdfs-all\")\n",
    "\n",
    "# List blobs in the container\n",
    "try:\n",
    "    blobs_list = container_client.list_blobs()\n",
    "    print(\"Blobs in the container:\")\n",
    "    for blob in blobs_list:\n",
    "        print(blob.name)\n",
    "    print(\"Access to the blob storage was granted.\")\n",
    "except Exception as e:\n",
    "    print(f\"Failed to access the blob storage: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Data Source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents.indexes import SearchIndexerClient\n",
    "from azure.search.documents.indexes.models import (\n",
    "    SearchIndexerDataContainer,\n",
    "    SearchIndexerDataSourceConnection\n",
    ")\n",
    "\n",
    "# Create a data source \n",
    "indexer_client = SearchIndexerClient(endpoint=azure_search_service_endpoint, credential=credential)\n",
    "container = SearchIndexerDataContainer(name=\"nasa-ebooks-pdfs-all\")\n",
    "data_source_connection = SearchIndexerDataSourceConnection(\n",
    "    name=\"ai-search-ds\",\n",
    "    type=\"azureblob\",\n",
    "    connection_string=azure_storage_connection_string,\n",
    "    container=container\n",
    ")\n",
    "data_source = indexer_client.create_or_update_data_source_connection(data_source_connection)\n",
    "\n",
    "print(f\"Data source '{data_source.name}' created or updated\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Skillset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents.indexes.models import (\n",
    "    SplitSkill,\n",
    "    InputFieldMappingEntry,\n",
    "    OutputFieldMappingEntry,\n",
    "    AzureOpenAIEmbeddingSkill,\n",
    "    EntityRecognitionSkill,\n",
    "    SearchIndexerIndexProjection,\n",
    "    SearchIndexerIndexProjectionSelector,\n",
    "    SearchIndexerIndexProjectionsParameters,\n",
    "    IndexProjectionMode,\n",
    "    SearchIndexerSkillset,\n",
    "    CognitiveServicesAccountKey\n",
    ")\n",
    "\n",
    "# Create a skillset  \n",
    "skillset_name = \"ai-search-ss\"\n",
    "\n",
    "split_skill = SplitSkill(  \n",
    "    description=\"Split skill to chunk documents\",  \n",
    "    text_split_mode=\"pages\",  \n",
    "    context=\"/document\",  \n",
    "    maximum_page_length=2000,  \n",
    "    page_overlap_length=500,  \n",
    "    inputs=[  \n",
    "        InputFieldMappingEntry(name=\"text\", source=\"/document/content\"),  \n",
    "    ],  \n",
    "    outputs=[  \n",
    "        OutputFieldMappingEntry(name=\"textItems\", target_name=\"pages\")  \n",
    "    ],  \n",
    ")  \n",
    "  \n",
    "embedding_skill = AzureOpenAIEmbeddingSkill(  \n",
    "    description=\"Skill to generate embeddings via Azure OpenAI\",  \n",
    "    context=\"/document/pages/*\",  \n",
    "    resource_url=azure_openai_endpoint,  \n",
    "    deployment_name=azure_openai_embeddings_deployment,  \n",
    "    model_name=azure_openai_embeddings_deployment,\n",
    "    dimensions=1024,\n",
    "    inputs=[  \n",
    "        InputFieldMappingEntry(name=\"text\", source=\"/document/pages/*\"),  \n",
    "    ],  \n",
    "    outputs=[  \n",
    "        OutputFieldMappingEntry(name=\"embedding\", target_name=\"text_vector\")  \n",
    "    ],  \n",
    ")\n",
    "\n",
    "entity_skill = EntityRecognitionSkill(\n",
    "    description=\"Skill to recognize entities in text\",\n",
    "    context=\"/document/pages/*\",\n",
    "    categories=[\"Location\"],\n",
    "    default_language_code=\"en\",\n",
    "    inputs=[\n",
    "        InputFieldMappingEntry(name=\"text\", source=\"/document/pages/*\")\n",
    "    ],\n",
    "    outputs=[\n",
    "        OutputFieldMappingEntry(name=\"locations\", target_name=\"locations\")\n",
    "    ]\n",
    ")\n",
    "  \n",
    "index_projections = SearchIndexerIndexProjection(  \n",
    "    selectors=[  \n",
    "        SearchIndexerIndexProjectionSelector(  \n",
    "            target_index_name=azure_search_service_index_name,  \n",
    "            parent_key_field_name=\"parent_id\",  \n",
    "            source_context=\"/document/pages/*\",  \n",
    "            mappings=[  \n",
    "                InputFieldMappingEntry(name=\"chunk\", source=\"/document/pages/*\"),  \n",
    "                InputFieldMappingEntry(name=\"text_vector\", source=\"/document/pages/*/text_vector\"),\n",
    "                InputFieldMappingEntry(name=\"locations\", source=\"/document/pages/*/locations\"),  \n",
    "                InputFieldMappingEntry(name=\"title\", source=\"/document/metadata_storage_name\"),  \n",
    "            ],  \n",
    "        ),  \n",
    "    ],  \n",
    "    parameters=SearchIndexerIndexProjectionsParameters(  \n",
    "        projection_mode=IndexProjectionMode.SKIP_INDEXING_PARENT_DOCUMENTS  \n",
    "    ),  \n",
    ") \n",
    "\n",
    "cognitive_services_account = CognitiveServicesAccountKey(key=azure_ai_services_key)\n",
    "\n",
    "skills = [split_skill, embedding_skill, entity_skill]\n",
    "\n",
    "skillset = SearchIndexerSkillset(  \n",
    "    name=skillset_name,  \n",
    "    description=\"Skillset to chunk documents, generate embeddings, and extract location entities\",  \n",
    "    skills=skills,  \n",
    "    index_projection=index_projections,\n",
    "    cognitive_services_account=cognitive_services_account\n",
    ")\n",
    "  \n",
    "client = SearchIndexerClient(endpoint=azure_search_service_endpoint, credential=credential)  \n",
    "client.create_or_update_skillset(skillset)  \n",
    "print(f\"{skillset.name} created\")  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an Indexer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents.indexes.models import (\n",
    "    SearchIndexer,\n",
    "    IndexingSchedule\n",
    ")\n",
    "\n",
    "# Create an indexer  \n",
    "indexer_name = \"ai-search-idxr\" \n",
    "\n",
    "# Schedule to run every 24 hours\n",
    "schedule = IndexingSchedule(interval=\"P1D\")\n",
    "\n",
    "indexer_parameters = None\n",
    "\n",
    "indexer = SearchIndexer(  \n",
    "    name=indexer_name,  \n",
    "    description=\"Indexer to index documents, generate embeddings, and extract entities\",  \n",
    "    skillset_name=skillset_name,  \n",
    "    target_index_name=index_name,  \n",
    "    data_source_name=data_source.name,\n",
    "    parameters=indexer_parameters,\n",
    "    schedule=schedule\n",
    ")  \n",
    "\n",
    "# Create and run the indexer  \n",
    "indexer_client = SearchIndexerClient(endpoint=azure_search_service_endpoint, credential=credential)  \n",
    "indexer_result = indexer_client.create_or_update_indexer(indexer)  \n",
    "\n",
    "print(f' {indexer_name} is created and running. Give the indexer a few minutes before running a query.')  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Send a query to the search engine to check results\n",
    "\n",
    "Executed every time a user makes a query"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform a Hybrid Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.models import VectorizableTextQuery\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "\n",
    "# Get credential from Azure AI Search Admin key\n",
    "credential = AzureKeyCredential(azure_search_service_admin_key)\n",
    "search_client = SearchClient(endpoint=azure_search_service_endpoint, \n",
    "                             credential=credential, \n",
    "                             index_name=azure_search_service_index_name)\n",
    "\n",
    "# User Query\n",
    "query = \"What can I see in the United States?\"  \n",
    "\n",
    "# Convert query into vector form\n",
    "vector_query = VectorizableTextQuery(text=query, \n",
    "                                     k_nearest_neighbors=50, \n",
    "                                     fields=\"text_vector\",\n",
    "                                     weight=1)\n",
    "\n",
    "results = search_client.search(  \n",
    "    search_text=query,  #This performs a full-text search using the query\n",
    "    vector_queries= [vector_query], #This adds a vector search component\n",
    "    select=[\"title\",\"chunk\",\"locations\"], #Specify fields to be return in the result\n",
    "    top=3\n",
    ") \n",
    "\n",
    "for result in results:  \n",
    "    print(f\"Score: {result['@search.score']} \\n\")\n",
    "    print(f\"Title: {result['title']} \\n\")\n",
    "    print(f\"Chunk: {result['chunk']} \\n\")\n",
    "    print(f\"Locations: {result['locations']}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform a Semantic Hybrid Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.models import VectorizableTextQuery\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "\n",
    "# Get credential from Azure AI Search Admin key\n",
    "credential = AzureKeyCredential(azure_search_service_admin_key)\n",
    "search_client = SearchClient(endpoint=azure_search_service_endpoint, \n",
    "                             credential=credential, \n",
    "                             index_name=azure_search_service_index_name)\n",
    "\n",
    "# User Query\n",
    "query = \"What can I see in the United States?\"  \n",
    "\n",
    "# Convert query into vector form\n",
    "vector_query = VectorizableTextQuery(text=query, \n",
    "                                     k_nearest_neighbors=50, \n",
    "                                     fields=\"text_vector\",\n",
    "                                     weight=1)\n",
    "\n",
    "results = search_client.search(\n",
    "    query_type=\"semantic\", \n",
    "    semantic_configuration_name='my-semantic-config',\n",
    "    scoring_profile=\"my-scoring-profile\",\n",
    "    scoring_parameters=[\"tags-beach, 'United States'\"], \n",
    "    search_text=query,\n",
    "    vector_queries= [vector_query],\n",
    "    select=[\"title\",\"chunk\",\"locations\"],\n",
    "    top=3,\n",
    ")\n",
    "\n",
    "# Sort the results by score in descending order\n",
    "sorted_results = sorted(results, key=lambda x: x['@search.score'], reverse=True)\n",
    "\n",
    "for result in sorted_results:  \n",
    "    print(f\"Score: {result['@search.score']} \\n\")\n",
    "    print(f\"Title: {result['title']} \\n\")\n",
    "    print(f\"Chunk: {result['chunk']} \\n\")\n",
    "    print(f\"Locations: {result['locations']}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Send query results to a language model to generate response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "from azure.search.documents.models import VectorizableTextQuery\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "# Get credential from Azure AI Search Admin key\n",
    "credential = AzureKeyCredential(azure_search_service_admin_key)\n",
    "search_client = SearchClient(endpoint=azure_search_service_endpoint, \n",
    "                             credential=credential, \n",
    "                             index_name=azure_search_service_index_name)\n",
    "\n",
    "# Azure OpenAI client\n",
    "openai_client = AzureOpenAI(\n",
    "    # to get version: https://learn.microsoft.com/en-us/azure/ai-services/openai/api-version-deprecation\n",
    "    api_version=azure_openai_api_version,\n",
    "    azure_endpoint=azure_openai_endpoint,\n",
    "    api_key=azure_openai_key)\n",
    "\n",
    "# Provide instructions to the model\n",
    "SYSTEM_PROMPT=\"\"\"\n",
    "You are an AI assistant that helps users learn from the information found in the source material.\n",
    "Answer the query using only the sources provided below.\n",
    "Use bullets if the answer has multiple points.\n",
    "If the answer is longer than 3 sentences, provide a summary.\n",
    "Answer ONLY with the facts listed in the list of sources below. Cite your source when you answer the question\n",
    "If there isn't enough information below, say you don't know.\n",
    "Do not generate answers that don't use the sources below.\n",
    "Query: {query}\n",
    "Sources:\\n{sources}\n",
    "\"\"\"\n",
    "\n",
    "# User Query\n",
    "query = \"What can I see in the United States?\"  \n",
    "\n",
    "# Convert query into vector form\n",
    "vector_query = VectorizableTextQuery(text=query, \n",
    "                                     k_nearest_neighbors=50, \n",
    "                                     fields=\"text_vector\",\n",
    "                                     weight=1)\n",
    "\n",
    "results = search_client.search(\n",
    "    query_type=\"semantic\", \n",
    "    semantic_configuration_name='my-semantic-config',\n",
    "    search_text=query,\n",
    "    vector_queries= [vector_query],\n",
    "    select=[\"title\",\"chunk\",\"locations\"],\n",
    "    top=5,\n",
    ")\n",
    "\n",
    "# Use a unique separator to make the sources distinct. \n",
    "# We chose repeated equal signs (=) followed by a newline because it's unlikely the source documents contain this sequence.\n",
    "sources_formatted = \"=================\\n\".join([f'TITLE: {document[\"title\"]}, CONTENT: {document[\"chunk\"]}, LOCATIONS: {document[\"locations\"]}' for document in results])\n",
    "\n",
    "response = openai_client.chat.completions.create(\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": SYSTEM_PROMPT.format(query=query, sources=sources_formatted)\n",
    "        }\n",
    "    ],\n",
    "    model=azure_openai_deployment\n",
    ")\n",
    "\n",
    "print(response.choices[0].message.content)\n"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}