Merge pull request #62 from gsa9989/users/garagundi/cosmosdbnosql

Adding LlamaIndex RAG Chatbot Sample for Azure CosmosDB NoSql

Author: jcodella

Python/CosmosDB-NoSQL-Integrations/LlamaIndex-CosmosDBNoSql-RAG-Chatbot-Sample.ipynb

@@ -0,0 +1,311 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "936cff50-b7ff-434a-b6ae-558f6288a5fb",
+   "metadata": {},
+   "source": [
+    "# Building a RAG Chatbot Using LlamaIndex\n",
+    "\n",
+    "LlamaIndex provides users with a simple way of creating a chatbot that works both with an LLM and data from a database. This combination is called Retrieval-Augmented Generation (RAG) and is used to give LLM's the ability to answer queries using data it was not trained on. This notebook will cover each step necessary to create a RAG chatbot using the Python SDK for Azure Cosmos DB for NoSQL. At the end, we create a UX using gradio to allow users to type in questions and see the response displayed in a chatbot style.\n",
+    "\n",
+    "Important Note: This sample requires you to have Azure Cosmos DB for NoSQL and Azure OpenAI accounts setup. To get started, visit:\n",
+    "-  [Azure Cosmos DB for NoSQL Python Quickstart](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/quickstart-python?pivots=devcontainer-codespace)\n",
+    "-  [Azure Cosmos DB for NoSQL Vector Search](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/vector-search)\n",
+    "-  [Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "25524f64-e833-419f-9762-1b07e426e55b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install llama-index-embeddings-openai\n",
+    "%pip install llama-index-llms-azure-openai"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43ab1fab-06be-4b81-aa12-0f0aa84d5780",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install llama-index"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "593a3558-fba4-406c-bf1d-c9c3d577a165",
+   "metadata": {},
+   "source": [
+    "## Setup Azure OpenAI\n",
+    "Prior to beginning we need to set up the llm and embedding model that will be used in the RAG chatbot."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "92116904-efcd-41f1-8abf-afd63863062d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.llms.azure_openai import AzureOpenAI\n",
+    "from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding\n",
+    "import os\n",
+    "from dotenv import load_dotenv"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f45c2108-ac32-42f6-9e12-1f58c65cd453",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = AzureOpenAI(\n",
+    "    model = \"gpt-35-turbo\",\n",
+    "    deployment_name = \"gpt-35-turbo\",\n",
+    "    azure_endpoint = os.getenv('AZURE_OPENAI_API_ENDPOINT'),\n",
+    "    api_key = os.getenv('AZURE_OPENAI_API_KEY'),\n",
+    "    api_version = \"2023-05-15\"\n",
+    ")\n",
+    "\n",
+    "embed_model = AzureOpenAIEmbedding(\n",
+    "    model = \"text-embedding-3-large\",\n",
+    "    deployment_name = \"text-embedding-3-large\",\n",
+    "    azure_endpoint = os.getenv('AZURE_OPENAI_API_ENDPOINT'),\n",
+    "    api_key = os.getenv('AZURE_OPENAI_API_KEY'),\n",
+    "    api_version = \"2023-05-15\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a9c75260-482f-4405-a831-b912d0c782ce",
+   "metadata": {},
+   "source": [
+    "## Loading the data\n",
+    "The first step is to load the data using the LlamaIndex function SimpleDirectoryReader."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1453d69b-d409-4eae-8048-e102ca46ca0e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import nest_asyncio\n",
+    "from llama_index.core import SimpleDirectoryReader\n",
+    "from llama_index.core.readers.base import BaseReader\n",
+    "from llama_index.core import Document"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d7be5ce9-b30b-4511-94c4-3c15f8011d4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "documents = SimpleDirectoryReader(input_files = [r\"DataSet/CVPR2019/abstracts_pdf\"]).load_data()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2aa21a24-eac7-49c2-a53a-20beb6d21aaa",
+   "metadata": {},
+   "source": [
+    "## Create the Index\n",
+    "The next step is to index the data loaded, this is done through vector embeddings. Prior to indexing it is important to initialize a Cosmos DB NoSql vector store where the embeddings will be stored."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c401450f-187a-41b7-8e77-fb7367081e13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from azure.cosmos import CosmosClient, PartitionKey\n",
+    "from llama_index.vector_stores.azurecosmosnosql import AzureCosmosDBNoSqlVectorSearch\n",
+    "from llama_index.core import StorageContext\n",
+    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d28f8d88-bc16-47fa-8088-a091c0d33a1a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from llama_index.core import Settings\n",
+    "\n",
+    "Settings.llm = llm\n",
+    "Settings.embed_model = embed_model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c35af3eb-b733-4e09-a98f-8557e14a7522",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#create cosmos client\n",
+    "URI = os.getenv('COSMOS_DB_URI')\n",
+    "KEY = os.getenv('COSMOS_DB_API_KEY')\n",
+    "client = CosmosClient(URI, credential=KEY)\n",
+    "\n",
+    "#specify vector store properties\n",
+    "indexing_policy = {\n",
+    "    \"indexingMode\": \"consistent\",\n",
+    "    \"includedPaths\": [{\"path\": \"/*\"}],\n",
+    "    \"excludedPaths\": [{\"path\": '/\"_etag\"/?'}],\n",
+    "    \"vectorIndexes\": [{\"path\": \"/embedding\", \"type\": \"quantizedFlat\"}],\n",
+    "}\n",
+    "\n",
+    "vector_embedding_policy = {\n",
+    "    \"vectorEmbeddings\": [\n",
+    "        {\n",
+    "            \"path\": \"/embedding\",\n",
+    "            \"dataType\": \"float32\",\n",
+    "            \"distanceFunction\": \"cosine\",\n",
+    "            \"dimensions\": 3072,\n",
+    "        }\n",
+    "    ]\n",
+    "}\n",
+    "\n",
+    "partition_key = PartitionKey(path=\"/id\")\n",
+    "cosmos_container_properties_test = {\"partition_key\": partition_key}\n",
+    "cosmos_database_properties_test = {}\n",
+    "\n",
+    "#create vector store\n",
+    "store = AzureCosmosDBNoSqlVectorSearch(cosmos_client=client,\n",
+    "                                       vector_embedding_policy=vector_embedding_policy,\n",
+    "                                       indexing_policy=indexing_policy,\n",
+    "                                       cosmos_container_properties=cosmos_container_properties_test,\n",
+    "                                       cosmos_database_properties=cosmos_database_properties_test,\n",
+    "                                       create_container=True,\n",
+    "                                       database_name = \"rag_chatbot_example\")\n",
+    "\n",
+    "storage_context = StorageContext.from_defaults(vector_store=store)\n",
+    "\n",
+    "#index the data\n",
+    "index = VectorStoreIndex.from_documents(\n",
+    "        documents, storage_context=storage_context\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5018f41f-a8e3-4a6b-9733-e7a4359b2f9d",
+   "metadata": {},
+   "source": [
+    "## Query the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "956c80f1-5b95-4833-b82d-bc2848bf01c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install gradio"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2ed68388-1910-401c-89d7-f957bbce2831",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import gradio as gr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "38e1b8d4-ef27-47e9-af8f-fb487000719a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "query_engine = index.as_query_engine()\n",
+    "def user_query(user_prompt, chat_history):\n",
+    "    # Create a timer to measure the time it takes to complete the request\n",
+    "    start_time = time.time()\n",
+    "    # Get LLM completion\n",
+    "    response = query_engine.query(user_prompt)    \n",
+    "    # Stop the timer\n",
+    "    end_time = time.time()\n",
+    "    elapsed_time = round((end_time - start_time) * 1000, 2)\n",
+    "    print(response)\n",
+    "    # Append user message and response to chat history\n",
+    "    details = f\"\\n (Time: {elapsed_time}ms)\"\n",
+    "    chat_history.append([user_prompt, str(response) + details])\n",
+    "        \n",
+    "    return gr.update(value=\"\"), chat_history"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f643488-442b-41e4-8dcb-d279c9310183",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chat_history = []\n",
+    "with gr.Blocks() as demo:\n",
+    "    chatbot = gr.Chatbot(label=\"RAG Chatbot\")\n",
+    "    \n",
+    "    msg = gr.Textbox(label=\"Ask me anything about the document!\")\n",
+    "    clear = gr.Button(\"Clear\")\n",
+    "    \n",
+    "    msg.submit(user_query, [msg, chatbot], [msg, chatbot], queue=False)\n",
+    "\n",
+    "    clear.click(lambda: None, None, chatbot, queue=False)\n",
+    "\n",
+    "# Launch the Gradio interface\n",
+    "demo.launch(debug=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1874d1c5-7ed1-43c2-a1f2-604f043fa449",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "demo.close()"
+   ]
+  }
+ ],
+ "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}