{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "86fc5bb2-017f-434e-8cd6-53ab214a5604",
   "metadata": {},
   "source": [
    "# Retrieval-augmented generation (RAG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de913d6d-c57f-4927-82fe-18902a636861",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/langchain-ai/langchain/blob/master/docs/docs/use_cases/question_answering/index.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5151afed",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "### What is RAG?\n",
    "\n",
    "RAG is a technique for augmenting LLM knowledge with additional, often private or real-time, data.\n",
    "\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. 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. The process of bringing the appropriate information and inserting it into the model prompt is known as Retrieval Augmented Generation (RAG).\n",
    "\n",
    "### What's in this guide?\n",
    "\n",
    "LangChain has a number of components specifically designed to help build RAG applications. To familiarize ourselves with these, we'll build a simple question-answering application over a text data source. Specifically, we'll build a QA bot over the [LLM Powered Autonomous Agents](https://lilianweng.github.io/posts/2023-06-23-agent/) blog post by Lilian Weng. Along the way we'll go over a typical QA architecture, discuss the relevant LangChain components, and highlight additional resources for more advanced QA techniques. We'll also see how LangSmith can help us trace and understand our application. LangSmith will become increasingly helpful as our application grows in complexity.\n",
    "\n",
    "**Note**\n",
    "Here we focus on RAG for unstructured data. Two RAG use cases which we cover elsewhere are:\n",
    "- [QA over structured data](/docs/use_cases/qa_structured/sql) (e.g., SQL)\n",
    "- [QA over code](/docs/use_cases/question_answering/code_understanding) (e.g., Python)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f25cbbd-0938-4e3d-87e4-17a204a03ffb",
   "metadata": {},
   "source": [
    "## Architecture\n",
    "A typical RAG application has two main components:\n",
    "\n",
    "**Indexing**: a pipeline for ingesting data from a source and indexing it. *This usually happen offline.*\n",
    "\n",
    "**Retrieval and generation**: the actual RAG chain, which takes the user query at run time and retrieves the relevant data from the index, then passes that to the model.\n",
    "\n",
    "The most common full sequence from raw data to answer looks like:\n",
    "\n",
    "#### Indexing\n",
    "1. **Load**: First we need to load our data. We'll use [DocumentLoaders](/docs/modules/data_connection/document_loaders/) for this.\n",
    "2. **Split**: [Text splitters](/docs/modules/data_connection/document_transformers/) break large `Documents` into smaller chunks. This is useful both for indexing data and for passing it in to a model, since large chunks are harder to search over and won't in a model's finite context window.\n",
    "3. **Store**: We need somewhere to store and index our splits, so that they can later be searched over. This is often done using a [VectorStore](/docs/modules/data_connection/vectorstores/) and [Embeddings](/docs/modules/data_connection/text_embedding/) model.\n",
    "\n",
    "![index_diagram](/img/rag_indexing.png)\n",
    "\n",
    "#### Retrieval and generation\n",
    "4. **Retrieve**: Given a user input, relevant splits are retrieved from storage using a [Retriever](/docs/modules/data_connection/retrievers/).\n",
    "5. **Generate**: A [ChatModel](/docs/modules/model_io/chat_models) / [LLM](/docs/modules/model_io/llms/) produces an answer using a prompt that includes the question and the retrieved data\n",
    "\n",
    "![retrieval_diagram](/img/rag_retrieval_generation.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "487d8d79-5ee9-4aa4-9fdf-cd5f4303e099",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "### Dependencies\n",
    "\n",
    "We'll use an OpenAI chat model and embeddings and a Chroma vector store in this walkthrough, but everything shown here works with any [ChatModel](/docs/integrations/chat/) or [LLM](/docs/integrations/llms/), [Embeddings](/docs/integrations/text_embedding/), and [VectorStore](/docs/integrations/vectorstores/) or [Retriever](/docs/integrations/retrievers). \n",
    "\n",
    "We'll use the following packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28d272cd-4e31-40aa-bbb4-0be0a1f49a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U langchain openai chromadb langchainhub bs4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51ef48de-70b6-4f43-8e0b-ab9b84c9c02a",
   "metadata": {},
   "source": [
    "We need to set environment variable `OPENAI_API_KEY`, which can be done directly or loaded from a `.env` file like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "143787ca-d8e6-4dc9-8281-4374f4d71720",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass()\n",
    "\n",
    "# import dotenv\n",
    "\n",
    "# dotenv.load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1665e740-ce01-4f09-b9ed-516db0bd326f",
   "metadata": {},
   "source": [
    "### LangSmith\n",
    "\n",
    "Many of the applications you build with LangChain will contain multiple steps with multiple invocations of LLM calls. As these applications get more and more complex, it becomes crucial to be able to inspect what exactly is going on inside your chain or agent. The best way to do this is with [LangSmith](https://smith.langchain.com).\n",
    "\n",
    "Note that LangSmith is not needed, but it is helpful. If you do want to use LangSmith, after you sign up at the link above, make sure to set your environment variables to start logging traces:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07411adb-3722-4f65-ab7f-8f6f57663d11",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = getpass.getpass()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa6ba684-26cf-4860-904e-a4d51380c134",
   "metadata": {},
   "source": [
    "## Quickstart\n",
    "\n",
    "Suppose we want to build a QA app over the [LLM Powered Autonomous Agents](https://lilianweng.github.io/posts/2023-06-23-agent/) blog post by Lilian Weng. We can create a simple pipeline for this in ~20 lines of code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8a913b1-0eea-442a-8a64-ec73333f104b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import bs4\n",
    "from langchain import hub\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.document_loaders import WebBaseLoader\n",
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain.vectorstores import Chroma"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "820244ae-74b4-4593-b392-822979dd91b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = WebBaseLoader(\n",
    "    web_paths=(\"https://lilianweng.github.io/posts/2023-06-23-agent/\",),\n",
    "    bs_kwargs=dict(\n",
    "        parse_only=bs4.SoupStrainer(\n",
    "            class_=(\"post-content\", \"post-title\", \"post-header\")\n",
    "        )\n",
    "    ),\n",
    ")\n",
    "docs = loader.load()\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "splits = text_splitter.split_documents(docs)\n",
    "\n",
    "vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())\n",
    "retriever = vectorstore.as_retriever()\n",
    "\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)\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": "code",
   "execution_count": 5,
   "id": "0d3b0f36-7b56-49c0-8e40-a1aa9ebcbf24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Task decomposition is a technique used to break down complex tasks into smaller and simpler steps. It can be done through prompting techniques like Chain of Thought or Tree of Thoughts, or by using task-specific instructions or human inputs. Task decomposition helps agents plan ahead and manage complicated tasks more effectively.'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag_chain.invoke(\"What is Task Decomposition?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7cb344e0-c423-400c-a079-964c08e07e32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup\n",
    "vectorstore.delete_collection()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "639dc31a-7f16-40f6-ba2a-20e7c2ecfe60",
   "metadata": {},
   "source": [
    ":::tip Check out the [LangSmith trace](https://smith.langchain.com/public/1c6ca97e-445b-4d00-84b4-c7befcbc59fe/r) \n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "842cf72d-abbc-468e-a2eb-022470347727",
   "metadata": {},
   "source": [
    "## Detailed walkthrough\n",
    "\n",
    "Let's go through the above code step-by-step to really understand what's going on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba5daed6",
   "metadata": {},
   "source": [
    "## Step 1. Load\n",
    "\n",
    "We need to first load the blog post contents. We can use `DocumentLoader`s for this, which are objects that load in data from a source as `Documents`.  A `Document` is an object with `page_content` (str) and `metadata` (dict) attributes. \n",
    "\n",
    "In this case we'll use the `WebBaseLoader`, which uses `urllib` and `BeautifulSoup` to load and parse the passed in web urls, returning one `Document` per url. We can customize the html -> text parsing by passing in parameters to the `BeautifulSoup` parser via `bs_kwargs` (see [BeautifulSoup docs](https://beautiful-soup-4.readthedocs.io/en/latest/#beautifulsoup)). In this case only HTML tags with class \"post-content\", \"post-title\", or \"post-header\" are relevant, so we'll remove all others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cf4d5c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import WebBaseLoader\n",
    "\n",
    "loader = WebBaseLoader(\n",
    "    web_paths=(\"https://lilianweng.github.io/posts/2023-06-23-agent/\",),\n",
    "    bs_kwargs={\n",
    "        \"parse_only\": bs4.SoupStrainer(\n",
    "            class_=(\"post-content\", \"post-title\", \"post-header\")\n",
    "        )\n",
    "    },\n",
    ")\n",
    "docs = loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "207f87a3-effa-4457-b013-6d233bc7a088",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "42824"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(docs[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "52469796-5ce4-4c12-bd2a-a903872dac33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "      LLM Powered Autonomous Agents\n",
      "    \n",
      "Date: June 23, 2023  |  Estimated Reading Time: 31 min  |  Author: Lilian Weng\n",
      "\n",
      "\n",
      "Building agents with LLM (large language model) as its core controller is a cool concept. Several proof-of-concepts demos, such as AutoGPT, GPT-Engineer and BabyAGI, serve as inspiring examples. The potentiality of LLM extends beyond generating well-written copies, stories, essays and programs; it can be framed as a powerful general problem solver.\n",
      "Agent System Overview#\n",
      "In\n"
     ]
    }
   ],
   "source": [
    "print(docs[0].page_content[:500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee5c6556-56be-4067-adbc-98b5aa19ef6e",
   "metadata": {},
   "source": [
    "### Go deeper\n",
    "`DocumentLoader`: Object that load data from a source as `Documents`.\n",
    "- [Docs](/docs/modules/data_connection/document_loaders/): Further documentation on how to use `DocumentLoader`s.\n",
    "- [Integrations](/docs/integrations/document_loaders/): Find the relevant `DocumentLoader` integration (of the > 160 of them) for your use case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2cc9a7",
   "metadata": {},
   "source": [
    "## Step 2. Split\n",
    "\n",
    "Our loaded document is over 42k characters long. This is too long to fit in the context window of many models. And even for those models that could fit the full post in their context window, empirically models struggle to find the relevant context in very long prompts. \n",
    "\n",
    "So we'll split the `Document` into chunks for embedding and vector storage. This should help us retrieve only the most relevant bits of the blog post at run time.\n",
    "\n",
    "In this case we'll split our documents into chunks of 1000 characters with 200 characters of overlap between chunks. The overlap helps mitigate the possibility of separating a statement from important context related to it. We use the `RecursiveCharacterTextSplitter`, which will (recursively) split the document using common separators (like new lines) until each chunk is the appropriate size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4b11c01d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=1000, chunk_overlap=200, add_start_index=True\n",
    ")\n",
    "all_splits = text_splitter.split_documents(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "3741eb67-9caf-40f2-a001-62f49349bff5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "66"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(all_splits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f868d0e5-5670-4d54-b562-f50265e907f4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "969"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(all_splits[0].page_content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5c9e5f27-c8e3-4ca7-8a8e-45c5de2901cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       " 'start_index': 7056}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_splits[10].metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a33bd4d",
   "metadata": {},
   "source": [
    "### Go deeper\n",
    "\n",
    "`DocumentSplitter`: Object that splits a list of `Document`s into smaller chunks. Subclass of `DocumentTransformer`s.\n",
    "- Explore `Context-aware splitters`, which keep the location (\"context\") of each split in the original `Document`:\n",
    "    - [Markdown files](/docs/use_cases/question_answering/document-context-aware-QA)\n",
    "    - [Code (py or js)](docs/integrations/document_loaders/source_code)\n",
    "    - [Scientific papers](/docs/integrations/document_loaders/grobid)\n",
    "\n",
    "`DocumentTransformer`: Object that performs a transformation on a list of `Document`s.\n",
    "- [Docs](/docs/modules/data_connection/document_transformers/): Further documentation on how to use `DocumentTransformer`s\n",
    "- [Integrations](/docs/integrations/document_transformers/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46547031-2352-4321-9970-d6ea27285c2e",
   "metadata": {},
   "source": [
    "## Step 3. Store\n",
    "\n",
    "Now that we've got 66 text chunks in memory, we need to store and index them so that we can search them later in our RAG app. The most common way to do this is to embed the contents of each document split and upload those embeddings to a vector store. \n",
    "\n",
    "Then, when we want to search over our splits, we take the search query, embed it as well, and perform some sort of \"similarity\" search to identify the stored splits with the most similar embeddings to our query embedding. The simplest similarity measure is cosine similarity — we measure the cosine of the angle between each pair of embeddings (which are just very high dimensional vectors).\n",
    "\n",
    "We can embed and store all of our document splits in a single command using the `Chroma` vector store and `OpenAIEmbeddings` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e9c302c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings import OpenAIEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "vectorstore = Chroma.from_documents(documents=all_splits, embedding=OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc6f22b0",
   "metadata": {},
   "source": [
    "### Go deeper\n",
    "`Embeddings`: Wrapper around a text embedding model, used for converting text to embeddings.\n",
    "- [Docs](/docs/modules/data_connection/text_embedding): Further documentation on the interface.\n",
    "- [Integrations](/docs/integrations/text_embedding/): Browse the > 30 text embedding integrations\n",
    "\n",
    "`VectorStore`: Wrapper around a vector database, used for storing and querying embeddings.\n",
    "- [Docs](/docs/modules/data_connection/vectorstores/): Further documentation on the interface.\n",
    "- [Integrations](/docs/integrations/vectorstores/): Browse the > 40 `VectorStore` integrations.\n",
    "\n",
    "This completes the **Indexing** portion of the pipeline. At this point we have an query-able vector store containing the chunked contents of our blog post. Given a user question, we should ideally be able to return the snippets of the blog post that answer the question:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70d64d40-e475-43d9-b64c-925922bb5ef7",
   "metadata": {},
   "source": [
    "## Step 4. Retrieve\n",
    "\n",
    "Now let's write the actual application logic. We want to create a simple application that let's the user ask a question, searches for documents relevant to that question, passes the retrieved documents and initial question to a model, and finally returns an answer.\n",
    "\n",
    "LangChain defines a `Retriever` interface which wraps an index that can return relevant documents given a string query. All retrievers implement a common method `get_relevant_documents()` (and its asynchronous variant `aget_relevant_documents()`).\n",
    "\n",
    "The most common type of `Retriever` is the `VectorStoreRetriever`, which uses the similarity search capabilities of a vector store to facillitate retrieval. Any `VectorStore` can easily be turned into a `Retriever`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4414df0d-5d43-46d0-85a9-5f47be0dd099",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = vectorstore.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 6})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e2c26b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieved_docs = retriever.get_relevant_documents(\n",
    "    \"What are the approaches to Task Decomposition?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8684291d-0f5e-453a-8d3e-ff9feea765d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(retrieved_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9a5dc074-816d-409a-b005-ab4eddfd76af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tree of Thoughts (Yao et al. 2023) extends CoT by exploring multiple reasoning possibilities at each step. It first decomposes the problem into multiple thought steps and generates multiple thoughts per step, creating a tree structure. The search process can be BFS (breadth-first search) or DFS (depth-first search) with each state evaluated by a classifier (via a prompt) or majority vote.\n",
      "Task decomposition can be done (1) by LLM with simple prompting like \"Steps for XYZ.\\n1.\", \"What are the subgoals for achieving XYZ?\", (2) by using task-specific instructions; e.g. \"Write a story outline.\" for writing a novel, or (3) with human inputs.\n"
     ]
    }
   ],
   "source": [
    "print(retrieved_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d5a113b",
   "metadata": {},
   "source": [
    "### Go deeper\n",
    "Vector stores are commonly used for retrieval, but there are plenty of other ways to do retrieval. \n",
    "\n",
    "`Retriever`: An object that returns `Document`s given a text query\n",
    "- [Docs](/docs/modules/data_connection/retrievers/): Further documentation on the interface and built-in retrieval techniques. Some of which include:\n",
    "    - `MultiQueryRetriever` [generates variants of the input question](/docs/modules/data_connection/retrievers/MultiQueryRetriever) to improve retrieval hit rate.\n",
    "    - `MultiVectorRetriever` (diagram below) instead generates [variants of the embeddings](/docs/modules/data_connection/retrievers/multi_vector), also in order to improve retrieval hit rate.\n",
    "    - `Max marginal relevance` selects for [relevance and diversity](https://www.cs.cmu.edu/~jgc/publication/The_Use_MMR_Diversity_Based_LTMIR_1998.pdf) among the retrieved documents to avoid passing in duplicate context.\n",
    "    - Documents can be filtered during vector store retrieval using [`metadata` filters](/docs/use_cases/question_answering/document-context-aware-QA).\n",
    "- [Integrations](/docs/integrations/retrievers/): Integrations with retrieval services."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "415d6824",
   "metadata": {},
   "source": [
    "## Step 5. Generate\n",
    "\n",
    "Let's put it all together into a chain that takes a question, retrieves relevant documents, constructs a prompt, passes that to a model, and parses the output.\n",
    "\n",
    "We'll use the gpt-3.5-turbo OpenAI chat model, but any LangChain `LLM` or `ChatModel` could be substituted in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d34d998c-9abf-4e01-a4ad-06dadfcf131c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc826723-36fc-45d1-a3ef-df8c2c8471a8",
   "metadata": {},
   "source": [
    "We'll use a prompt for RAG that is checked into the LangChain prompt hub ([here](https://smith.langchain.com/hub/rlm/rag-prompt))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "bede955b-9aeb-4fd3-964d-8e43f214ce70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain import hub\n",
    "\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "11c35354-f275-47ec-9f72-ebd5c23731eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.\n",
      "Question: filler question \n",
      "Context: filler context \n",
      "Answer:\n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    prompt.invoke(\n",
    "        {\"context\": \"filler context\", \"question\": \"filler question\"}\n",
    "    ).to_string()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51f9a210-1eee-4054-99d7-9d9ddf7e3593",
   "metadata": {},
   "source": [
    "We'll use the [LCEL Runnable](https://python.langchain.com/docs/expression_language/) protocol to define the chain, allowing us to \n",
    "- pipe together components and functions in a transparent way\n",
    "- automatically trace our chain in LangSmith\n",
    "- get streaming, async, and batched calling out of the box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "99fa1aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema import StrOutputParser\n",
    "from langchain.schema.runnable import RunnablePassthrough\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": "code",
   "execution_count": 23,
   "id": "8655a152-d7cf-466f-b1bc-fbff9ae2b889",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task decomposition is a technique used to break down complex tasks into smaller and simpler steps. It can be done through methods like Chain of Thought (CoT) or Tree of Thoughts, which involve dividing the task into manageable subtasks and exploring multiple reasoning possibilities at each step. Task decomposition can be performed by AI models with prompting, task-specific instructions, or human inputs."
     ]
    }
   ],
   "source": [
    "for chunk in rag_chain.stream(\"What is Task Decomposition?\"):\n",
    "    print(chunk, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c000e5f-2b7f-4eb9-8876-9f4b186b4a08",
   "metadata": {},
   "source": [
    ":::tip Check out the [LangSmith trace](https://smith.langchain.com/public/1799e8db-8a6d-4eb2-84d5-46e8d7d5a99b/r) \n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7d52c84",
   "metadata": {},
   "source": [
    "### Go deeper\n",
    "\n",
    "#### Choosing LLMs\n",
    "`ChatModel`: An LLM-backed chat model wrapper. Takes in a sequence of messages and returns a message.\n",
    "- [Docs](/docs/modules/model_io/chat/)\n",
    "- [Integrations](/docs/integrations/chat/): Explore over 25 `ChatModel` integrations.\n",
    "\n",
    "`LLM`: A text-in-text-out LLM. Takes in a string and returns a string.\n",
    "- [Docs](/docs/modules/model_io/llms)\n",
    "- [Integrations](/docs/integrations/llms): Explore over 75 `LLM` integrations.\n",
    "\n",
    "See a guide on RAG with locally-running models [here](/docs/modules/use_cases/question_answering/local_retrieval_qa)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa82f437",
   "metadata": {},
   "source": [
    "#### Customizing the prompt\n",
    "\n",
    "As shown above, we can load prompts (e.g., [this RAG prompt](https://smith.langchain.com/hub/rlm/rag-prompt)) from the prompt hub. The prompt can also be easily customized:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e4fee704",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Task decomposition is the process of breaking down a complex task into smaller and simpler steps. It can be done through techniques like Chain of Thought (CoT) or Tree of Thoughts, which involve dividing the problem into multiple thought steps and generating multiple thoughts per step. Task decomposition helps in enhancing model performance and understanding the thinking process of the model. Thanks for asking!'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "template = \"\"\"Use the following pieces of context to answer the question at the end. \n",
    "If you don't know the answer, just say that you don't know, don't try to make up an answer. \n",
    "Use three sentences maximum and keep the answer as concise as possible. \n",
    "Always say \"thanks for asking!\" at the end of the answer. \n",
    "{context}\n",
    "Question: {question}\n",
    "Helpful Answer:\"\"\"\n",
    "rag_prompt_custom = PromptTemplate.from_template(template)\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()}\n",
    "    | rag_prompt_custom\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "rag_chain.invoke(\"What is Task Decomposition?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94b952e6-dc4b-415b-9cf3-1ad333e48366",
   "metadata": {},
   "source": [
    ":::tip Check out the [LangSmith trace](https://smith.langchain.com/public/da23c4d8-3b33-47fd-84df-a3a582eedf84/r) \n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2f99b5-80b4-4178-bf30-c1c0a152638f",
   "metadata": {},
   "source": [
    "### Adding sources\n",
    "\n",
    "With LCEL it's easy to return the retrieved documents or certain source metadata from the documents:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ded41680-b749-4e2a-9daa-b1165d74783b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'documents': [{'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 1585},\n",
       "  {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 2192},\n",
       "  {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 17804},\n",
       "  {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 17414},\n",
       "  {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 29630},\n",
       "  {'source': 'https://lilianweng.github.io/posts/2023-06-23-agent/',\n",
       "   'start_index': 19373}],\n",
       " 'answer': 'Task decomposition is a technique used to break down complex tasks into smaller and simpler steps. It involves transforming big tasks into multiple manageable tasks, allowing for a more systematic and organized approach to problem-solving. Thanks for asking!'}"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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",
    "    | rag_prompt_custom\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(\"What is Task Decomposition\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b437da5d-ca09-4d15-9be2-c35e5a1ace77",
   "metadata": {},
   "source": [
    ":::tip Check out the [LangSmith trace](https://smith.langchain.com/public/007d7e01-cb62-4a84-8b71-b24767f953ee/r)\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "776ae958-cbdc-4471-8669-c6087436f0b5",
   "metadata": {},
   "source": [
    "### Adding memory\n",
    "\n",
    "Suppose we want to create a stateful application that remembers past user inputs. There are two main things we need to do to support this.\n",
    "1. Add a messages placeholder to our chain which allows us to pass in historical messages\n",
    "2. Add a chain that takes the latest user query and reformulates it in the context of the chat history into a standalone question that can be passed to our retriever.\n",
    "\n",
    "Let's start with 2. We can build a \"condense question\" chain that looks something like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2b685428-8b82-4af1-be4f-7232c5d55b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "condense_q_system_prompt = \"\"\"Given a chat history and the latest user question \\\n",
    "which might reference the chat history, formulate a standalone question \\\n",
    "which can be understood without the chat history. Do NOT answer the question, \\\n",
    "just reformulate it if needed and otherwise return it as is.\"\"\"\n",
    "condense_q_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", condense_q_system_prompt),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "condense_q_chain = condense_q_prompt | llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "46ee9aa1-16f1-4509-8dae-f8c71f4ad47d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'What is the definition of \"large\" in the context of a language model?'"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.schema.messages import AIMessage, HumanMessage\n",
    "\n",
    "condense_q_chain.invoke(\n",
    "    {\n",
    "        \"chat_history\": [\n",
    "            HumanMessage(content=\"What does LLM stand for?\"),\n",
    "            AIMessage(content=\"Large language model\"),\n",
    "        ],\n",
    "        \"question\": \"What is meant by large\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "31ee8481-ce37-41ae-8ca5-62196619d4b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'How do transformer models function?'"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "condense_q_chain.invoke(\n",
    "    {\n",
    "        \"chat_history\": [\n",
    "            HumanMessage(content=\"What does LLM stand for?\"),\n",
    "            AIMessage(content=\"Large language model\"),\n",
    "        ],\n",
    "        \"question\": \"How do transformers work\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42a47168-4a1f-4e39-bd2d-d5b03609a243",
   "metadata": {},
   "source": [
    "And now we can build our full QA chain. Notice we add some routing functionality to only run the \"condense question chain\" when our chat history isn't empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "66f275f3-ddef-4678-b90d-ee64576878f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_system_prompt = \"\"\"You are an assistant for question-answering tasks. \\\n",
    "Use the following pieces of retrieved context to answer the question. \\\n",
    "If you don't know the answer, just say that you don't know. \\\n",
    "Use three sentences maximum and keep the answer concise.\\\n",
    "\n",
    "{context}\"\"\"\n",
    "qa_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", qa_system_prompt),\n",
    "        MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "def condense_question(input: dict):\n",
    "    if input.get(\"chat_history\"):\n",
    "        return condense_q_chain\n",
    "    else:\n",
    "        return input[\"question\"]\n",
    "\n",
    "\n",
    "rag_chain = (\n",
    "    RunnablePassthrough.assign(context=condense_question | retriever | format_docs)\n",
    "    | qa_prompt\n",
    "    | llm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "51fd0e54-5bb4-4a9a-b012-87a18ebe2bef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Common ways of task decomposition include:\\n\\n1. Using Chain of Thought (CoT): CoT is a prompting technique that instructs the model to \"think step by step\" and decompose complex tasks into smaller and simpler steps. It utilizes more test-time computation and sheds light on the model\\'s thinking process.\\n\\n2. Prompting with LLM: Language Model (LLM) can be used to prompt the model with simple instructions like \"Steps for XYZ\" or \"What are the subgoals for achieving XYZ?\" This allows the model to generate a sequence of subtasks or thought steps.\\n\\n3. Task-specific instructions: For certain tasks, task-specific instructions can be provided to guide the model in decomposing the task. For example, for writing a novel, the instruction \"Write a story outline\" can be given to break down the task into manageable steps.\\n\\n4. Human inputs: In some cases, human inputs can be used to assist in task decomposition. Humans can provide their expertise and knowledge to identify and break down complex tasks into smaller subtasks.')"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat_history = []\n",
    "\n",
    "question = \"What is Task Decomposition?\"\n",
    "ai_msg = rag_chain.invoke({\"question\": question, \"chat_history\": chat_history})\n",
    "chat_history.extend([HumanMessage(content=question), ai_msg])\n",
    "\n",
    "second_question = \"What are common ways of doing it?\"\n",
    "rag_chain.invoke({\"question\": second_question, \"chat_history\": chat_history})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53263a65-4de2-4dd8-9291-6a8169ab6f1d",
   "metadata": {},
   "source": [
    ":::tip Check out the [LangSmith trace](https://smith.langchain.com/public/b3001782-bb30-476a-886b-12da17ec258f/r) \n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdf6c7e0-84f8-4747-b2ae-e84315152bd9",
   "metadata": {},
   "source": [
    "Here we've gone over how to add chain logic for incorporating historical outputs. But how do we actually store and retrieve historical outputs for different sessions? For that check out the LCEL [How to add message history (memory)](/docs/expression_language/how_to/message_history) page."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "580e18de-132d-4009-ba67-4aaf2c7717a2",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "That's a lot of content we've covered in a short amount of time. There's plenty of nuances, features, integrations, etc to explore in each of the above sections. Aside from the sources mentioned above, good next steps include:\n",
    "\n",
    "- Reading up on more advanced retrieval techniques in the [Retrievers](/docs/modules/data_connection/retrievers/) section.\n",
    "- Learning about the LangChain [Indexing API](/docs/modules/data_connection/indexing), which helps repeatedly sync data sources and vector stores without redundant computation or storage.\n",
    "- Exploring RAG [LangChain Templates](/docs/templates/#-advanced-retrieval), which are reference applications that can easily be deployed with [LangServe](/docs/langserve).\n",
    "- Learning about [evaluating RAG applications with LangSmith](https://github.com/langchain-ai/langsmith-cookbook/blob/main/testing-examples/qa-correctness/qa-correctness.ipynb)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv",
   "language": "python",
   "name": "poetry-venv"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
