RAG.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "52cb90b4",
   "metadata": {},
   "source": [
    "# Installing dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "387c03a3-4a4f-40de-842f-2cc838f5c0a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install langchain pypdf chromadb rapidocr-onnxruntime lark tiktoken"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ed5ec02",
   "metadata": {},
   "source": [
    "# Loading and Preprocessing the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d5e17c5",
   "metadata": {},
   "source": [
    "## Load files from a directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbae3683-cfda-4baa-ab18-6e4d5982843f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders.pdf import PyPDFDirectoryLoader\n",
    "\n",
    "DOC_DIR = './commentary_files'\n",
    "\n",
    "# Load Documents\n",
    "loader = PyPDFDirectoryLoader(DOC_DIR, extract_images=True)\n",
    "docs = loader.load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8fa3997",
   "metadata": {},
   "source": [
    "## Check file metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e4f8b4-2689-4625-8f15-d498a6f72bbc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# print(len(docs))\n",
    "for doc in docs:\n",
    "    print(doc.metadata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b41ccb8-60ed-4e82-9c75-f8fcd26c5e00",
   "metadata": {},
   "source": [
    "## Split the documents into chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef536247-f573-4f47-ba84-f1dd8e573781",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)\n",
    "splits = text_splitter.split_documents(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc0e2c29-bd0b-404a-acd2-ad637a8eac9c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for doc in splits:\n",
    "    if 'keywords' in doc.metadata:\n",
    "        if type(doc.metadata['keywords']) == list:\n",
    "            print(\"Fixing list\", doc.metadata['keywords'])\n",
    "            doc.metadata['keywords'] = ','.join(doc.metadata['keywords'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a530a060-caaf-4a8b-9856-a621a3af5bac",
   "metadata": {},
   "source": [
    "## Generate some metadata for each chunk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f2e23c-e9ff-4c1f-b89b-66c45bbb06ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import json\n",
    "\n",
    "def extract_json_objects(text):\n",
    "    # Regex to extract potential JSON objects (not guaranteed to be valid JSON)\n",
    "    pattern = r'\\{[^{}]*\\}'\n",
    "    potential_jsons = re.findall(pattern, text, re.DOTALL)\n",
    "\n",
    "    # List to store valid JSON objects\n",
    "    valid_json_objects = []\n",
    "\n",
    "    # Validate each extracted string as JSON\n",
    "    for potential_json in potential_jsons:\n",
    "        try:\n",
    "            # Attempt to parse the JSON string\n",
    "            potential_json = re.sub(\",[ \\t\\r\\n]+}\", \"}\", potential_json)\n",
    "            potential_json = re.sub(\",[ \\t\\r\\n]+\\]\", \"]\", potential_json)\n",
    "            json_object = json.loads(potential_json)\n",
    "            # If successful, append to the list of valid JSON objects\n",
    "            valid_json_objects.append(json_object)\n",
    "        except json.JSONDecodeError:\n",
    "            # If JSON is not valid, skip it\n",
    "            continue\n",
    "    \n",
    "    return valid_json_objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99d00e24-ae94-45af-bbd8-36f222da8dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOllama\n",
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Prompt\n",
    "template = \"\"\"Given some input text, summarize it and generate a JSON containing the following keys only:\n",
    "- topic: A title that describes the content\n",
    "- summary: A concise and accurate summary of the input document\n",
    "- keywords: A string of comma-separated semantic keywords associated with the content text\n",
    "\n",
    "Return the output as a json only.\n",
    "\n",
    "Text: {text}\n",
    "\"\"\"\n",
    "\n",
    "# LLM\n",
    "llm = ChatOllama(model_name=\"llama2\", temperature=1)\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "summarizer = prompt | llm | StrOutputParser()\n",
    "\n",
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ad498bb-4aeb-429b-a6b3-e90e678e791e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "for idx, doc in enumerate(splits):\n",
    "    print(f\"-------------- Processing chunk {idx}\")\n",
    "    tags = summarizer.invoke(doc.page_content)\n",
    "    try:\n",
    "        tags = extract_json_objects(tags)[0]\n",
    "    except:\n",
    "        print(f\"Received non-json for document {doc.metadata['source']}. \\nReceived: {tags}\\n\\nPlease enter the expected JSON:\")\n",
    "        j = input()\n",
    "        tags = json.loads(j)\n",
    "        \n",
    "    if tags:\n",
    "        doc.metadata.update(tags)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51264021-062c-4315-b180-337b6ad23897",
   "metadata": {},
   "source": [
    "## Vectorize the chunks (using OpenAI due to a compatibility issue with LangChain and Ollama for Self Query retrievers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ea2a16b-b9fc-48b2-b7bd-8af3d5f8ac2a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_community.embeddings import OpenAIEmbeddings\n",
    "\n",
    "# Embed\n",
    "vectorstore = Chroma.from_documents(documents=splits, \n",
    "                                    embedding=OpenAIEmbeddings(), persist_directory=\"./chroma_db\")\n",
    "\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6addd691-c3cc-4d86-a585-fa2dbcea62c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = retriever.get_relevant_documents(\"What are mega grants?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "364bf989-5b5a-46de-bb29-158fac646a56",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(len(docs))\n",
    "for doc in docs:\n",
    "    print(\"--------\")\n",
    "    print(doc.metadata.keys())\n",
    "    print(doc.page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f56b89bd-55b6-4db3-924d-2f979bfae036",
   "metadata": {},
   "source": [
    "## Use this LLM for all below types except Self Query (that one needs OpenAI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e25a301-661b-430c-9bb0-820d8dbf8db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LLM\n",
    "from langchain.chat_models import ChatOllama\n",
    "llm = ChatOllama(model_name=\"llama2\", temperature=0.7)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9151ed14",
   "metadata": {},
   "source": [
    "# RAGs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b08408a8-6627-4355-9be8-7d9df5032819",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Basic RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32296a62-8b4f-49e9-9f11-1109aa568186",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatOllama\n",
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "\n",
    "# Prompt\n",
    "template = \"\"\"Please list the semantic keywords associated with the following user question. Return as a comma-separated list. Also, answer the question based only on the following context:\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70f48c03-c307-41a3-9a6c-f420034aada2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Chain\n",
    "chain = prompt | llm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73cfb8b0-fb71-45eb-8021-c7f4b11ea8df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run\n",
    "chain.invoke({\"context\": docs, \"question\": \"What are pending mega grants?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696a7e4e-4dcb-4616-8bff-e764219f49e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "rag_chain.invoke(\"What are pending mega grants?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ff78a07-ae74-492a-8bf7-0e1e08f791bf",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Multi Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2027b31a-306c-4a01-89e1-09ae68f93aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.chat_models import ChatOllama\n",
    "\n",
    "question = 'What are insights?'\n",
    "\n",
    "template = \"\"\"You are an AI language model assistant. Your task is to generate five \n",
    "different versions of the given user question to retrieve relevant documents from a vector \n",
    "database. By generating multiple perspectives on the user question, your goal is to help\n",
    "the user overcome some of the limitations of the distance-based similarity search. \n",
    "Provide these alternative questions separated by newlines. Original question: {question}\"\"\"\n",
    "prompt_perspectives = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "generate_queries = (\n",
    "    prompt_perspectives \n",
    "    | ChatOllama(model_name=\"llama2\", temperature=1) \n",
    "    | StrOutputParser() \n",
    "    | (lambda x: x.split(\"\\n\"))\n",
    ")\n",
    "\n",
    "generate_queries.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caeff482-07f8-44f9-81e5-707c6961dde3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.load import dumps, loads\n",
    "\n",
    "def get_unique_union(documents: list[list]):\n",
    "    \"\"\" Unique union of retrieved docs \"\"\"\n",
    "    # Flatten list of lists, and convert each Document to string\n",
    "    flattened_docs = [dumps(doc) for sublist in documents for doc in sublist]\n",
    "    # Get unique documents\n",
    "    unique_docs = list(set(flattened_docs))\n",
    "    # Return\n",
    "    return [loads(doc) for doc in unique_docs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89cf7ef7-e417-4b3b-8e3e-fce794da0140",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "from langchain.chat_models import ChatOllama\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "# RAG\n",
    "template = \"\"\"Answer the following question based on this context:\n",
    "\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "llm = ChatOllama(model_name=\"llama2\", temperature=0)\n",
    "retrieval_chain = generate_queries | retriever.map() | get_unique_union\n",
    "\n",
    "final_rag_chain = (\n",
    "    {\"context\": retrieval_chain, \n",
    "     \"question\": itemgetter(\"question\")} \n",
    "    | prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "final_rag_chain.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "460047d0-cc70-47f6-bb3a-0e2b01984aae",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## RAG-Fusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "307571e4-c01f-46e1-846b-433f801b0297",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "\n",
    "template = \"\"\"You are a helpful assistant thinks through a question based on a context. Reply only with a list of questions. \\n\n",
    "For context: {context}, generate multiple search queries related to: {question} \\n\n",
    "Output (4 queries):\"\"\"\n",
    "prompt_rag_fusion = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0efc479c-43f0-4a97-9a95-c8c76d12cd1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain.chat_models import ChatOllama\n",
    "\n",
    "generate_queries = (\n",
    "    prompt_rag_fusion \n",
    "    | llm\n",
    "    | StrOutputParser() \n",
    "    | (lambda x: x.split(\"\\n\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c124148e-787a-4ecc-ab1d-4e537cd90a7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.load import dumps, loads\n",
    "\n",
    "def reciprocal_rank_fusion(results: list[list], k=60):\n",
    "    \"\"\" Reciprocal_rank_fusion that takes multiple lists of ranked documents \n",
    "        and an optional parameter k used in the RRF formula \"\"\"\n",
    "    \n",
    "    # Initialize a dictionary to hold fused scores for each unique document\n",
    "    fused_scores = {}\n",
    "\n",
    "    # Iterate through each list of ranked documents\n",
    "    for docs in results:\n",
    "        # Iterate through each document in the list, with its rank (position in the list)\n",
    "        for rank, doc in enumerate(docs):\n",
    "            # Convert the document to a string format to use as a key (assumes documents can be serialized to JSON)\n",
    "            doc_str = dumps(doc)\n",
    "            # If the document is not yet in the fused_scores dictionary, add it with an initial score of 0\n",
    "            if doc_str not in fused_scores:\n",
    "                fused_scores[doc_str] = 0\n",
    "            # Retrieve the current score of the document, if any\n",
    "            previous_score = fused_scores[doc_str]\n",
    "            # Update the score of the document using the RRF formula: 1 / (rank + k)\n",
    "            fused_scores[doc_str] += 1 / (rank + k)\n",
    "\n",
    "    # Sort the documents based on their fused scores in descending order to get the final reranked results\n",
    "    reranked_results = [\n",
    "        (loads(doc), score)\n",
    "        for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)\n",
    "    ]\n",
    "\n",
    "    # Return the reranked results as a list of tuples, each containing the document and its fused score\n",
    "    for doc in reranked_results:\n",
    "        print(doc[0].metadata['source'], doc[1])\n",
    "    return reranked_results\n",
    "\n",
    "retrieval_chain_rag_fusion = generate_queries | retriever.map() | reciprocal_rank_fusion\n",
    "# docs = retrieval_chain_rag_fusion.invoke({\"question\": question})\n",
    "# len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc049faa-fc4e-4534-874c-f43d5cf912c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# len(docs)\n",
    "# for doc in docs:\n",
    "#     print(doc[0].metadata['source'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d682b481-8bbd-4141-be35-cf7494b26908",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from operator import itemgetter\n",
    "\n",
    "# Retrieve\n",
    "question = input(\"Q:\")\n",
    "# question = \"What are pending mega grants?\"\n",
    "\n",
    "# print(\"Thinking these questions ------------\")\n",
    "questions = generate_queries.invoke({\"question\": question, \"context\": retriever})\n",
    "for question in questions:\n",
    "    print(question)\n",
    "# docs = retrieval_chain_rag_fusion.invoke({\"question\": question})\n",
    "\n",
    "# retrieval_chain = generate_queries | retriever.map() | get_unique_union\n",
    "# docs = retrieval_chain.invoke({\"context\": retriever, \"question\": question})\n",
    "# for doc in docs:\n",
    "#     print(doc.metadata['source'])\n",
    "# print(len(docs))\n",
    "\n",
    "# RAG\n",
    "template = \"\"\"Answer the following questions ONLY from this context:\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "final_rag_chain = (\n",
    "    {\"context\": retrieval_chain_rag_fusion, \n",
    "     \"question\": itemgetter(\"question\")} \n",
    "    | prompt\n",
    "    | llm  \n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "final_rag_chain.invoke({\"question\": question, \"context\": retriever}).split('\\n\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a377ef02-21cb-425c-8cf8-7743757349f7",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Decomposition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "711e95ea-def7-44ef-9d86-dfafdbadb949",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "\n",
    "# Decomposition\n",
    "template = \"\"\"You are a helpful assistant that generates multiple sub-questions related to an input question. \\n\n",
    "The goal is to break down the input into a set of sub-problems / sub-questions that can be answers in isolation. \\n\n",
    "Generate multiple search queries related to: {question} \\n\n",
    "Output (3 queries):\"\"\"\n",
    "prompt_decomposition = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3150f0c0-1785-4f95-923c-0298e19ffa9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "question = \"What are mega grants?\"\n",
    "\n",
    "# Chain\n",
    "generate_queries_decomposition = ( prompt_decomposition | llm | StrOutputParser() | (lambda x: x.split(\"\\n\")))\n",
    "\n",
    "# Run\n",
    "questions = generate_queries_decomposition.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01896140-f965-43f2-818e-30d29c128ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7be6a3d-235b-4200-ba9c-5cc9236868ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt\n",
    "template = \"\"\"Here is the question you need to answer:\n",
    "\n",
    "\\n --- \\n {question} \\n --- \\n\n",
    "\n",
    "Here is any available background question + answer pairs:\n",
    "\n",
    "\\n --- \\n {q_a_pairs} \\n --- \\n\n",
    "\n",
    "Here is additional context relevant to the question: \n",
    "\n",
    "\\n --- \\n {context} \\n --- \\n\n",
    "\n",
    "Use the above context and any background question + answer pairs to answer the question: \\n {question}\n",
    "\"\"\"\n",
    "\n",
    "decomposition_prompt = ChatPromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "813a3b9c-81f7-4b2a-8711-213a7bb13fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "def format_qa_pair(question, answer):\n",
    "    \"\"\"Format Q and A pair\"\"\"\n",
    "    \n",
    "    formatted_string = \"\"\n",
    "    formatted_string += f\"Question: {question}\\nAnswer: {answer}\\n\\n\"\n",
    "    return formatted_string.strip()\n",
    "\n",
    "q_a_pairs = \"\"\n",
    "for q in questions:\n",
    "    \n",
    "    rag_chain = (\n",
    "    {\"context\": itemgetter(\"question\") | retriever, \n",
    "     \"question\": itemgetter(\"question\"),\n",
    "     \"q_a_pairs\": itemgetter(\"q_a_pairs\")} \n",
    "    | decomposition_prompt\n",
    "    | llm\n",
    "    | StrOutputParser())\n",
    "\n",
    "    answer = rag_chain.invoke({\"question\":q,\"q_a_pairs\":q_a_pairs})\n",
    "    q_a_pair = format_qa_pair(q,answer)\n",
    "    q_a_pairs = q_a_pairs + \"\\n---\\n\"+  q_a_pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c753820-d306-437c-8eda-0454724cdff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53163e12-48b5-401f-a0da-21126af24d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"\n",
    "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: {question} \n",
    "Context: {context} \n",
    "\"\"\"\n",
    "\n",
    "prompt_rag = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "def retrieve_and_rag(question,prompt_rag,sub_question_generator_chain):\n",
    "    \"\"\"RAG on each sub-question\"\"\"\n",
    "    \n",
    "    # Use our decomposition / \n",
    "    sub_questions = sub_question_generator_chain.invoke({\"question\":question})\n",
    "    \n",
    "    # Initialize a list to hold RAG chain results\n",
    "    rag_results = []\n",
    "    \n",
    "    for sub_question in sub_questions:\n",
    "        \n",
    "        # Retrieve documents for each sub-question\n",
    "        retrieved_docs = retriever.get_relevant_documents(sub_question)\n",
    "        \n",
    "        # Use retrieved documents and sub-question in RAG chain\n",
    "        answer = (prompt_rag | llm | StrOutputParser()).invoke({\"context\": retrieved_docs, \n",
    "                                                                \"question\": sub_question})\n",
    "        rag_results.append(answer)\n",
    "    \n",
    "    return rag_results,sub_questions\n",
    "\n",
    "# Wrap the retrieval and RAG process in a RunnableLambda for integration into a chain\n",
    "answers, questions = retrieve_and_rag(question, prompt_rag, generate_queries_decomposition)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8ffe167-d23b-4628-805a-ad14a7987311",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_qa_pairs(questions, answers):\n",
    "    \"\"\"Format Q and A pairs\"\"\"\n",
    "    \n",
    "    formatted_string = \"\"\n",
    "    for i, (question, answer) in enumerate(zip(questions, answers), start=1):\n",
    "        formatted_string += f\"Question {i}: {question}\\nAnswer {i}: {answer}\\n\\n\"\n",
    "    return formatted_string.strip()\n",
    "\n",
    "context = format_qa_pairs(questions, answers)\n",
    "\n",
    "# Prompt\n",
    "template = \"\"\"Here is a set of Q+A pairs:\n",
    "\n",
    "{context}\n",
    "\n",
    "Use these to synthesize an answer to the question: {question}\n",
    "\"\"\"\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "final_rag_chain = (\n",
    "    prompt\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "final_rag_chain.invoke({\"context\":context,\"question\":question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "465485ee-6e06-466a-832c-7f154b9c6808",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Contextual Compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b4b9efc-dc3f-4531-bb0c-3799de40db9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
    "\n",
    "compressor = LLMChainExtractor.from_llm(llm)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69748b8d-b6e0-4d9c-8729-cbc0466e9830",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pretty_print_docs(docs):\n",
    "    print(\n",
    "        f\"\\n{'-' * 100}\\n\".join(\n",
    "            [f\"Document {i+1} ({doc.metadata['source']}):\\n\\n\" + d.page_content for i, d in enumerate(docs)]\n",
    "        )\n",
    "    )\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What are mega grants?\"\n",
    ")\n",
    "\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40604638-bc44-41b4-adf6-967fe1551133",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import EmbeddingsFilter\n",
    "from langchain_community.embeddings import OllamaEmbeddings\n",
    "\n",
    "embeddings = OllamaEmbeddings()\n",
    "embeddings_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.5)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=embeddings_filter, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What are mega grants?\"\n",
    ")\n",
    "\n",
    "pretty_print_docs(compressed_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e326048-0c25-4fb4-a45e-d10330bf1519",
   "metadata": {},
   "source": [
    "## Self Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2b45f37-c936-4f05-b1a4-7d0dd9312904",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "\n",
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"topic\",\n",
    "        description=\"The topic describing the content\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"summary\",\n",
    "        description=\"A brief summary of the content\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"keywords\",\n",
    "        description=\"Semantic keywords associated with the document\",\n",
    "        type=\"string\",\n",
    "    )\n",
    "]\n",
    "\n",
    "print(metadata_field_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c1e43aa-cd63-4e19-b108-8f6f4a7a0a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "\n",
    "document_content_description = \"B3i usage guides\"\n",
    "\n",
    "llm = ChatOpenAI(temperature=0.7)\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectorstore,\n",
    "    document_content_description,\n",
    "    metadata_field_info\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b386b09-1c58-40c3-be84-836e5ca0c3ea",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "retriever.get_relevant_documents(\"What are mega grants?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7592e7bc-22e5-41cd-983d-0939b74cba6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = retriever.invoke(\"What are mega grants?\")\n",
    "for doc in docs:\n",
    "    print(doc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4722f7d4-0437-4973-aac6-1860d2691c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts.chat import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from operator import itemgetter\n",
    "\n",
    "template = \"\"\"\n",
    "You are a helpful assistant and your job is to answer the following question only from this context. Make it layman and easy to understand. Suggest steps to explore the product and explain how to interact. Do not make stuff up, if there's no data, say you don't know\n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "final_rag_chain = (\n",
    "    {\"context\": retriever, \n",
    "     \"question\": itemgetter(\"question\")} \n",
    "    | prompt\n",
    "    | llm  \n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "question = input()\n",
    "final_rag_chain.invoke({\"question\": question})"
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}