SampleCode/Python/sample_rag_langchain.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Retrieval Augmented Generation (RAG)\n",
    "\n",
    "This notebook demonstrates an example of using [LangChain](https://www.langchain.com/) to delvelop a Retrieval Augmented Generation (RAG) pattern. It uses Azure AI Document Intelligence as document loader, which can extracts tables, paragraphs, and layout information from pdf, image, office and html files. The output markdown can be used in LangChain's markdown header splitter, which enables semantic chunking of the documents. Then the chunked documents are indexed into Azure AI Search vectore store. Given a user query, it will use Azure AI Search to get the relevant chunks, then feed the context into the prompt with the query to generate the answer.\n",
    "\n",
    "![Semantic chunking in RAG](../media/semantic-chunking-rag.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "- An Azure AI Document Intelligence resource in one of the 3 preview regions: **East US**, **West US2**, **West Europe** - follow [this document](https://learn.microsoft.com/azure/ai-services/document-intelligence/create-document-intelligence-resource?view=doc-intel-4.0.0) to create one if you don't have.\n",
    "- An Azure AI Search resource - follow [this document](https://learn.microsoft.com/azure/search/search-create-service-portal) to create one if you don't have.\n",
    "- An Azure OpenAI resource and deployments for embeddings model and chat model - follow [this document](https://learn.microsoft.com/azure/ai-services/openai/how-to/create-resource?pivots=web-portal) to create one if you don't have."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install python-dotenv langchain langchain-community langchain-openai langchainhub openai tiktoken azure-ai-documentintelligence azure-identity azure-search-documents==11.6.0b3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "This code loads environment variables using the `dotenv` library and sets the necessary environment variables for Azure services.\n",
    "The environment variables are loaded from the `.env` file in the same directory as this notebook.\n",
    "\"\"\"\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "os.environ[\"AZURE_OPENAI_ENDPOINT\"] = os.getenv(\"AZURE_OPENAI_ENDPOINT\")\n",
    "os.environ[\"AZURE_OPENAI_API_KEY\"] = os.getenv(\"AZURE_OPENAI_API_KEY\")\n",
    "doc_intelligence_endpoint = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT\")\n",
    "doc_intelligence_key = os.getenv(\"AZURE_DOCUMENT_INTELLIGENCE_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "from langchain_openai import AzureChatOpenAI\n",
    "from langchain_community.document_loaders import AzureAIDocumentIntelligenceLoader\n",
    "from langchain_openai import AzureOpenAIEmbeddings\n",
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.text_splitter import MarkdownHeaderTextSplitter\n",
    "from langchain.vectorstores.azuresearch import AzureSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load a document and split it into semantic chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initiate Azure AI Document Intelligence to load the document. You can either specify file_path or url_path to load the document.\n",
    "loader = AzureAIDocumentIntelligenceLoader(file_path=\"<path to your file>\", api_key = doc_intelligence_key, api_endpoint = doc_intelligence_endpoint, api_model=\"prebuilt-layout\")\n",
    "docs = loader.load()\n",
    "\n",
    "# Split the document into chunks base on markdown headers.\n",
    "headers_to_split_on = [\n",
    "    (\"#\", \"Header 1\"),\n",
    "    (\"##\", \"Header 2\"),\n",
    "    (\"###\", \"Header 3\"),\n",
    "]\n",
    "text_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)\n",
    "\n",
    "docs_string = docs[0].page_content\n",
    "splits = text_splitter.split_text(docs_string)\n",
    "\n",
    "print(\"Length of splits: \" + str(len(splits)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embed and index the chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed the splitted documents and insert into Azure Search vector store\n",
    "\n",
    "aoai_embeddings = AzureOpenAIEmbeddings(\n",
    "    azure_deployment=\"<Azure OpenAI embeddings model>\",\n",
    "    openai_api_version=\"<Azure OpenAI API version>\",  # e.g., \"2023-12-01-preview\"\n",
    ")\n",
    "\n",
    "vector_store_address: str = os.getenv(\"AZURE_SEARCH_ENDPOINT\")\n",
    "vector_store_password: str = os.getenv(\"AZURE_SEARCH_ADMIN_KEY\")\n",
    "\n",
    "index_name: str = \"<your index name>\"\n",
    "vector_store: AzureSearch = AzureSearch(\n",
    "    azure_search_endpoint=vector_store_address,\n",
    "    azure_search_key=vector_store_password,\n",
    "    index_name=index_name,\n",
    "    embedding_function=aoai_embeddings.embed_query,\n",
    ")\n",
    "\n",
    "vector_store.add_documents(documents=splits)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retrive relevant chunks based on a question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Retrieve relevant chunks based on the question\n",
    "\n",
    "retriever = vector_store.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 3})\n",
    "\n",
    "retrieved_docs = retriever.get_relevant_documents(\n",
    "    \"<your question>\"\n",
    ")\n",
    "\n",
    "print(retrieved_docs[0].page_content)\n",
    "\n",
    "# Use a prompt for RAG that is checked into the LangChain prompt hub (https://smith.langchain.com/hub/rlm/rag-prompt?organizationId=989ad331-949f-4bac-9694-660074a208a7)\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "llm = AzureChatOpenAI(\n",
    "    openai_api_version=\"<Azure OpenAI API version>\",  # e.g., \"2023-12-01-preview\"\n",
    "    azure_deployment=\"<your chat model deployment name>\",\n",
    "    temperature=0,\n",
    ")\n",
    "\n",
    "\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Document Q&A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ask a question about the document\n",
    "\n",
    "rag_chain.invoke(\"<your question>\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Doucment Q&A with references"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Return the retrieved documents or certain source metadata from the documents\n",
    "\n",
    "from operator import itemgetter\n",
    "\n",
    "from langchain.schema.runnable import RunnableMap\n",
    "\n",
    "rag_chain_from_docs = (\n",
    "    {\n",
    "        \"context\": lambda input: format_docs(input[\"documents\"]),\n",
    "        \"question\": itemgetter(\"question\"),\n",
    "    }\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "rag_chain_with_source = RunnableMap(\n",
    "    {\"documents\": retriever, \"question\": RunnablePassthrough()}\n",
    ") | {\n",
    "    \"documents\": lambda input: [doc.metadata for doc in input[\"documents\"]],\n",
    "    \"answer\": rag_chain_from_docs,\n",
    "}\n",
    "\n",
    "rag_chain_with_source.invoke(\"<your question>\")"
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}