Journey 2 - Build the Ultimate Retrieval System for RAG/sample/2-Build-The-Ultimate-Retrieval-System-For-RAG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Advanced Search Techniques with Azure AI Search: Keyword, Vector, and Hybrid Methods\n",
    "\n",
    "This notebook demonstrates how to perform different types of searches using Azure AI Search, including keyword search, vector search, hybrid search, semantic ranking, and query rewriting.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "Before running the notebook, ensure you have the following: \n",
    "\n",
    "- [Fork](https://github.com/microsoft/rag-time/fork) the repository and clone it to your local machine by following the script below:\n",
    "\n",
    "    ```bash\n",
    "    git clone https://github.com/your-org/rag-time.git\n",
    "    cd rag-time\n",
    "    ```\n",
    "\n",
    "- An [Azure account](https://portal.azure.com) with proper permissions to access the following services:\n",
    "    - An **Azure AI Search** service with an index that contains vectorized text data. Follow the instructions in the [Quickstart](https://learn.microsoft.com/en-us/azure/search/search-get-started-portal-import-vectors?tabs=sample-data-storage%2Cmodel-aoai%2Cconnect-data-storage) to index the [MSFT_cloud_architecture_contoso.pdf](MSFT_cloud_architecture_contoso.pdf) file. \n",
    "- Install Python 3.8 or later from [python.org](https://python.org).\n",
    "\n",
    "## Steps to Use the Notebook\n",
    "\n",
    "### 1. Install Required Libraries\n",
    "\n",
    "Run the first code cell to install the required Python libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install azure-search-documents==11.6.0b9 azure-identity python-dotenv pandas jinja2 --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Set Up Environment Variables\n",
    "\n",
    "To store credentials securely, rename `.env.sample` file to `.env` in the same directory as the notebook and update the following variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AZURE_SEARCH_SERVICE_ENDPOINT=\"<your_search_service_endpoint>\"\n",
    "AZURE_SEARCH_INDEX=\"<your_search_index_name>\"\n",
    "AZURE_SEARCH_ADMIN_KEY=\"<your_search_admin_key>\"  # Leave blank if using Managed Identity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After setting up, the notebook will automatically load these values using dotenv.\n",
    "\n",
    "### 3. Load Environment Variables\n",
    "\n",
    "Run the following command to load environment variables from the .env file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True) # take environment variables from .env.\n",
    "\n",
    "endpoint = os.environ[\"AZURE_SEARCH_SERVICE_ENDPOINT\"]\n",
    "index_name = os.environ[\"AZURE_SEARCH_INDEX\"]\n",
    "credential = AzureKeyCredential(os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")) if os.getenv(\"AZURE_SEARCH_ADMIN_KEY\") else DefaultAzureCredential()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This will ensure all necessary credentials are available before setting up the API client."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Set Up API Client and Define the Display Function\n",
    "\n",
    "Initialize the Azure AI Search Client for interacting with the Azure Search service and make the search results easier to read by defining a function that formats and displays results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents import SearchClient\n",
    "import pandas as pd\n",
    "\n",
    "search_client = SearchClient(endpoint, index_name, credential)\n",
    "\n",
    "def display_results(results):\n",
    "    df = pd.json_normalize(list(results)).dropna(axis=1, how='all')\n",
    "    df[\"chunk\"] = df[\"chunk\"].apply(lambda c: c[:300] + '...' if len(c) > 300 else c)\n",
    "    first_cols = ['title', 'chunk', '@search.score']\n",
    "    df = df[first_cols + [col for col in df.columns if col not in first_cols]]\n",
    "\n",
    "    df = df.style.set_properties(**{\n",
    "        'max-width': '500px',\n",
    "        'text-align': 'left',\n",
    "        'white-space': 'normal',\n",
    "        'word-wrap': 'break-word'\n",
    "    }).hide(axis=\"index\")\n",
    "\n",
    "\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Perform Different Search Methods\n",
    "\n",
    "#### Keyword Search\n",
    "\n",
    "Execute a traditional keyword-based search:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(search_text=\"What is Contoso\", top=5, select=[\"title\", \"chunk\"])\n",
    "\n",
    "display_results(results)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Vector Search\n",
    "\n",
    "Retrieve documents using vector similarity search:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.search.documents.models import VectorizableTextQuery\n",
    "\n",
    "results = search_client.search(vector_queries=[VectorizableTextQuery(text=\"What is Contoso\", k_nearest_neighbors=50, fields=\"text_vector\")], top=5, select=[\"title\", \"chunk\"])\n",
    "\n",
    "display_results(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hybrid Search (Keyword + Vector Search)\n",
    "\n",
    "Combine keyword and vector searches for better accuracy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(\n",
    "    search_text=\"What is Contoso\",\n",
    "    vector_queries=[VectorizableTextQuery(text=\"What is Contoso\", k_nearest_neighbors=50, fields=\"text_vector\")],\n",
    "    top=5,\n",
    "    select=[\"title\", \"chunk\"]\n",
    ")\n",
    "\n",
    "display_results(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hybrid Search + Semantic Ranker\n",
    "\n",
    "Enhance search results using a semantic ranker:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(\n",
    "    search_text=\"What is Contoso\",\n",
    "    vector_queries=[VectorizableTextQuery(text=\"What is Contoso\", k_nearest_neighbors=50, fields=\"text_vector\")],\n",
    "    top=5,\n",
    "    select=[\"title\", \"chunk\"],\n",
    "    query_type=\"semantic\",\n",
    "    semantic_configuration_name=\"ragtime2-semantic-configuration\"\n",
    ")\n",
    "\n",
    "display_results(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Hybrid Search + Semantic Ranker + Query Rewriting\n",
    "\n",
    "Use semantic ranking and query rewriting for improved relevance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = search_client.search(\n",
    "    search_text=\"What is Contoso\",\n",
    "    vector_queries=[VectorizableTextQuery(text=\"What is Contoso\", k_nearest_neighbors=50, fields=\"text_vector\")],\n",
    "    top=5,\n",
    "    select=[\"title\", \"chunk\"],\n",
    "    query_type=\"semantic\",\n",
    "    semantic_configuration_name=\"ragtime2-semantic-configuration\",\n",
    "    query_rewrites=\"generative\",\n",
    "    query_language=\"en\"\n",
    ")\n",
    "\n",
    "display_results(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Troubleshooting\n",
    "\n",
    "- **Environment Variables Not Loaded:** Ensure you have correctly set the .env file or manually export them in your terminal before running the notebook.\n",
    "- **Authentication Issues:** If using Managed Identity, make sure your Azure identity has proper role assignments.\n",
    "- **Search Results Are Empty:** Ensure your Azure AI Search index contains vectorized data.\n",
    "- **Query Rewriting Issues:** Ensure your search service supports semantic configurations and generative query rewrites.\n",
    "\n",
    "## Summary\n",
    "\n",
    "This notebook demonstrates different search techniques using Azure AI Search, including keyword search, vector search, hybrid search, semantic ranking, and query rewriting. The approach enhances search accuracy by leveraging vector embeddings and semantic understanding to retrieve the most relevant documents.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}