{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c5fb0b9e-f9cd-404f-bd8d-0273e94ac1fe",
   "metadata": {},
   "source": [
    "# RAG Example Using NVIDIA API Catalog and LangChain"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2969cdab-82fc-4ce5-bde1-b4f629691f27",
   "metadata": {},
   "source": [
    "This notebook introduces how to use LangChain to interact with NVIDIA hosted NIM microservices like chat, embedding, and reranking models to build a simple retrieval-augmented generation (RAG) application."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4253bd0-4313-4056-95f5-899a180879c2",
   "metadata": {},
   "source": [
    "## Terminology"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5a084a00-b65d-483a-a7c6-b4c12e4272dd",
   "metadata": {},
   "source": [
    "#### RAG\n",
    "\n",
    "- RAG is a technique for augmenting LLM knowledge with additional data.\n",
    "- LLMs can reason about wide-ranging topics, but their knowledge is limited to the public data up to a specific point in time that they were trained on.\n",
    "- If you want to build AI applications that can reason about private data or data introduced after a model's cutoff date, you need to augment the knowledge of the model with the specific information it needs.\n",
    "- The process of bringing the appropriate information and inserting it into the model prompt is known as retrieval augmented generation (RAG).\n",
    "\n",
    "The preceding summary of RAG originates in the LangChain v0.2 tutorial [Build a RAG App](https://python.langchain.com/v0.2/docs/tutorials/rag/) tutorial in the LangChain v0.2 documentation.\n",
    "\n",
    "#### NIM\n",
    "\n",
    "- [NIM microservices](https://developer.nvidia.com/blog/nvidia-nim-offers-optimized-inference-microservices-for-deploying-ai-models-at-scale/) are containerized microservices that simplify the deployment of generative AI models like LLMs and are optimized to run on NVIDIA GPUs. \n",
    "- NIM microservices support models across domains like chat, embedding, reranking, and more from both the community and NVIDIA.\n",
    "\n",
    "#### NVIDIA API Catalog\n",
    "\n",
    "- [NVIDIA API Catalog](https://build.nvidia.com/explore/discover) is a hosted platform for accessing a wide range of microservices online.\n",
    "- You can test models on the catalog and then export them with an NVIDIA AI Enterprise license for on-premises or cloud deployment\n",
    "\n",
    "#### langchain-nvidia-ai-endpoints\n",
    "\n",
    "- The [`langchain-nvidia-ai-endpoints`](https://pypi.org/project/langchain-nvidia-ai-endpoints/) Python package contains LangChain integrations for building applications that communicate with NVIDIA NIM microservices."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca300278-5ff4-47c4-ab70-c6584ef73c9f",
   "metadata": {},
   "source": [
    "## Installation and Requirements\n",
    "\n",
    "Create a Python environment (preferably with Conda) using Python version 3.10.14. \n",
    "To install Jupyter Lab, refer to the [installation](https://jupyter.org/install) page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5de8e53a-7940-4e72-87f3-e1c014128806",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Requirements\n",
    "!pip install langchain==0.2.5\n",
    "!pip install langchain_community==0.2.5\n",
    "!pip install faiss-cpu==1.8.0 # replace with faiss-gpu if you are using GPU\n",
    "!pip install langchain-nvidia-ai-endpoints==0.1.2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b7a52a0-7e5e-4064-9665-cb947d600f84",
   "metadata": {},
   "source": [
    "## Getting Started!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04495732-c2db-4c97-91d0-96708814334d",
   "metadata": {},
   "source": [
    "To get started you need an `NVIDIA_API_KEY` to use the NVIDIA API Catalog:\n",
    "\n",
    "1) Create a free account with [NVIDIA](https://build.nvidia.com/explore/discover).\n",
    "2) Click on your model of choice.\n",
    "3) Under Input select the Python tab, and click **Get API Key** and then click **Generate Key**.\n",
    "4) Copy and save the generated key as NVIDIA_API_KEY. From there, you should have access to the endpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbb51115-79f8-48c3-b3ee-d434916945f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter your NVIDIA API key:  ········\n"
     ]
    }
   ],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if not os.environ.get(\"NVIDIA_API_KEY\", \"\").startswith(\"nvapi-\"):\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    assert nvidia_api_key.startswith(\"nvapi-\"), f\"{nvidia_api_key[:5]}... is not a valid key\"\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25656ab5-0046-4e27-be65-b3d3d547b4c6",
   "metadata": {},
   "source": [
    "## RAG Example using LLM & Embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54e86bc0-e9c5-4a2b-be0e-7fca0331e886",
   "metadata": {},
   "source": [
    "### 1) Initialize the LLM\n",
    "\n",
    "The ChatNVIDIA class is part of LangChain's integration (langchain_nvidia_ai_endpoints) with NVIDIA NIM microservices. \n",
    "It allows access to NVIDIA NIM for chat applications, connecting to hosted or locally-deployed microservices.\n",
    "\n",
    "Here we will use **mixtral-8x7b-instruct-v0.1** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "88c2fafe-5ded-4238-82de-f094232bf6fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "llm = ChatNVIDIA(model=\"mistralai/mixtral-8x7b-instruct-v0.1\", max_tokens=1024)\n",
    "\n",
    "# Here we are using mixtral-8x7b-instruct-v0.1 model\n",
    "# But you are free to choose any model hosted at Nvidia API Catalog\n",
    "# Uncomment the below code to list the availabe models\n",
    "# ChatNVIDIA.get_available_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35cc87a6-2f83-4652-95f1-cf349db8bad6",
   "metadata": {},
   "source": [
    "### 2) Intiatlize the embedding\n",
    "NVIDIAEmbeddings is a client to NVIDIA embeddings models that provides access to a NVIDIA NIM for embedding. It can connect to a hosted NIM or a local NIM using a base URL\n",
    "\n",
    "We selected **NV-Embed-QA** as the embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d88f7838-b9f9-4fc5-8779-84df6cb26017",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings\n",
    "\n",
    "embedder = NVIDIAEmbeddings(model=\"NV-Embed-QA\", truncate=\"END\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9862f2e-5055-4fe4-818d-708091243d74",
   "metadata": {},
   "source": [
    "### 3) Obtain some toy text dataset\n",
    "Here we are loading a toy data from a text documents and in real-time data can be loaded from various sources. \n",
    "Read [here](https://python.langchain.com/v0.2/docs/tutorials/rag/#go-deeper) for loading data from different sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9c615b9c-527e-4e3b-86b7-49ef258e2d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# For this example we load a toy data set (it's a simple text file with some information about Sweden)\n",
    "TOY_DATA_PATH = \"./data/\"\n",
    "# We read in the text data and prepare them into vectorstore\n",
    "ps = os.listdir(TOY_DATA_PATH)\n",
    "data = []\n",
    "sources = []\n",
    "for p in ps:\n",
    "    if p.endswith('.txt'):\n",
    "        path2file=TOY_DATA_PATH+p\n",
    "        with open(path2file,encoding=\"utf-8\") as f:\n",
    "            lines=f.readlines()\n",
    "            for line in lines:\n",
    "                if len(line)>=1:\n",
    "                    data.append(line)\n",
    "                    sources.append(path2file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d0a1447d-444a-4ae9-9484-4546424c047d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(400,\n",
       " 230,\n",
       " 'Sweden, formally the Kingdom of Sweden, is a Nordic country located on the Scandinavian Peninsula in Northern Europe. It borders Norway to the west and north, Finland to the east, and is connected to Denmark in the southwest by a bridge–tunnel across the Öresund. At 447,425 square kilometres (172,752 sq mi), Sweden is the largest Nordic country, the third-largest country in the European Union, and the fifth-largest country in Europe. The capital and largest city is Stockholm. Sweden has a total population of 10.5 million, and a low population density of 25.5 inhabitants per square kilometre (66/sq mi), with around 87% of Swedes residing in urban areas, which cover 1.5% of the entire land area, in the central and southern half of the country.\\n')"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Do some basic cleaning and remove empty lines\n",
    "documents=[d for d in data if d != '\\n']\n",
    "len(data), len(documents), data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "851b16b3-43ac-4269-9f37-05a33efe24fb",
   "metadata": {},
   "source": [
    "### 4) Process the documents into vectorstore and save it to disk\n",
    "\n",
    "Real world documents can be very long, this makes it hard to fit in the context window of many models. Even for those models that could fit the full post in their context window, models can struggle to find information in very long inputs.\n",
    "\n",
    "To handle this we’ll split the Document into chunks for embedding and vector storage. More on text splitting [here](https://python.langchain.com/v0.2/docs/concepts/#text-splitters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "804c85f6-181b-4291-a685-d6b378015544",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we create a faiss vector store from the documents and save it to disk.\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "\n",
    "text_splitter = CharacterTextSplitter(chunk_size=400, separator=\" \", chunk_overlap=80)\n",
    "docs = []\n",
    "metadatas = []\n",
    "\n",
    "for i, d in enumerate(documents):\n",
    "    splits = text_splitter.split_text(d)\n",
    "    docs.extend(splits)\n",
    "    metadatas.extend([{\"source\": sources[i]}] * len(splits))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f867df18-11c8-45ea-b81c-1603459431f9",
   "metadata": {},
   "source": [
    "To enable runtime search, we index text chunks by embedding each document split and storing these embeddings in a vector database. Later to search, we embed the query and perform a similarity search to find the stored splits with embeddings most similar to the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2d7b2fbd-8cb1-4d68-9659-2426b9ecffe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# you will only need to do this once, later on we will restore the already saved vectorstore\n",
    "store = FAISS.from_texts(docs, embedder , metadatas=metadatas)\n",
    "VECTOR_STORE = './data/nv_embedding'\n",
    "store.save_local(VECTOR_STORE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fe85dad-12bb-47d2-a407-9b89b5270d4e",
   "metadata": {},
   "source": [
    "### 5) Read the previously processed & saved vectore store back"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5de3e07d-5fbe-4fe7-8f23-ed0b082f2413",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the FAISS vectorestore back.\n",
    "store = FAISS.load_local(VECTOR_STORE, embedder, allow_dangerous_deserialization=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a41ff63-6adc-4055-8bc4-e7ecaad0fb4d",
   "metadata": {},
   "source": [
    "### 6) Wrap the restored vectorsore into a retriever and ask our question "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5aa362c9-48ab-4646-bc29-bc2aca92505d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\" Sweden is the 55th-largest country in the world, the fifth-largest country in Europe, and the largest country in Northern Europe, with a total area of 449,964 km2 (173,732 sq mi). In terms of elevation, the lowest point in Sweden is in the bay of Lake Hammarsjön, near Kristianstad, at -2.41 m (-7.91 ft) below sea level, while the highest point is Kebnekaise, which is 2,111 m (6,926 ft) above sea level.\\n\\nSweden has a Nordic social welfare system that provides universal health care and tertiary education for its citizens. The country has a high standard of living and ranks very highly in various international metrics, including quality of life, health, education, protection of civil liberties, economic competitiveness, income equality, gender equality, and prosperity. Sweden's GDP per capita is the world's 14th highest.\\n\\nHistorically, Sweden has been both a kingdom and an empire. Currently, it is a constitutional monarchy and a parliamentary democracy, with a popularly elected parliament and a monarch who serves a ceremonial role. Sweden is a member of the European Union but has opted to remain outside the Eurozone.\""
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever = store.as_retriever()\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"Answer solely based on the following context:\\n<Documents>\\n{context}\\n</Documents>\",\n",
    "        ),\n",
    "        (\"user\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Langchain's LCEL(LangChain Expression Language) Runnable protocol is used to define the chain\n",
    "# LCEL allows pipe together components and functions\n",
    "chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "chain.invoke(\"Tell me about Sweden.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29478b0-0fb1-4678-93cd-b159dc9884a7",
   "metadata": {},
   "source": [
    "## RAG Example with LLM, Embedding & Reranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "972c310b-5333-4b41-a6dd-ce83e739e6dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\" The documents provided do not include information about Gustav's grandson ascending the throne. Gustav had several grandchildren, and the documents do not specify which one you are referring to. Moreover, the documents do not provide enough information about the timeline of Gustav's grandson's ascension to the throne. Therefore, it is not possible to answer this question without additional context.\""
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's test a more complex query using the above LLM Embedding chain and see if the reranker can help.\n",
    "chain.invoke(\"In which year Gustav's grandson ascended the throne?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d3854c7-68a3-45b4-9e69-2c4e583d651f",
   "metadata": {},
   "source": [
    "### Enhancing accuracy for single data sources\n",
    "\n",
    "This example demonstrates how a re-ranking model can be used to combine retrieval results and improve accuracy during retrieval of documents.\n",
    "\n",
    "Typically, reranking is a critical piece of high-accuracy, efficient retrieval pipelines. Generally, there are two important use cases:\n",
    "\n",
    "- Combining results from multiple data sources\n",
    "- Enhancing accuracy for single data sources\n",
    "\n",
    "Here, we focus on demonstrating only the second use case. If you want to know more, check [here](https://github.com/langchain-ai/langchain-nvidia/blob/main/libs/ai-endpoints/docs/retrievers/nvidia_rerank.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b7e8677e-a37f-42e2-8fea-4c4413f7d682",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\" Gustav's grandson, Sigismund, ascended the throne in 1592.\""
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_nvidia_ai_endpoints import NVIDIARerank\n",
    "from langchain_core.runnables import RunnableParallel\n",
    "\n",
    "# We will narrow the collection to 100 results and further narrow it to 10 with the reranker.\n",
    "retriever = store.as_retriever(search_kwargs={'k':100}) # typically k will be 1000 for real world use-cases\n",
    "ranker = NVIDIARerank(model='nv-rerank-qa-mistral-4b:1', top_n=10)\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"Answer solely based on the following context:\\n<Documents>\\n{context}\\n</Documents>\",\n",
    "        ),\n",
    "        (\"user\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "reranker = lambda input: ranker.compress_documents(query=input['question'], documents=input['context'])\n",
    "\n",
    "chain_with_ranker = (\n",
    "    RunnableParallel({\"context\": retriever, \"question\": RunnablePassthrough()})\n",
    "    | {\"context\": reranker, \"question\": lambda input: input['question']}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "chain_with_ranker.invoke(\"In which year Gustav's grandson ascended the throne?\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79d5186f-12c0-47c9-9e85-5987fedf7b97",
   "metadata": {},
   "source": [
    "#### Note:\n",
    "- In this notebook, we have used NVIDIA NIM microservices from the NVIDIA API Catalog.\n",
    "- The above APIs, ChatNVIDIA, NVIDIAEmbedding, and NVIDIARerank, also support self-hosted NIM microservices.\n",
    "- Change the `base_url` to your deployed NIM URL.\n",
    "- Example: `llm = ChatNVIDIA(base_url=\"http://localhost:8000/v1\", model=\"meta/llama3-8b-instruct\")`\n",
    "- NIM can be hosted locally using Docker, following the [NVIDIA NIM for LLMs](https://docs.nvidia.com/nim/large-language-models/latest/getting-started.html) documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61236a22-922f-403f-89d1-1172251aeb0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example Code snippet if you want to use a self-hosted NIM\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA\n",
    "\n",
    "# connect to an LLM NIM running at localhost:8000, specifying a specific model\n",
    "llm = ChatNVIDIA(base_url=\"http://localhost:8000/v1\", model=\"meta/llama3-8b-instruct\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (rag_notebooks)",
   "language": "python",
   "name": "rag_notebooks"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}