Retrievers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Explore Various Retrieval Strategies with Retrievers in LangChain"
      ],
      "metadata": {
        "id": "jbw4wHV4zlKj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install OpenAI, and LangChain dependencies"
      ],
      "metadata": {
        "id": "L1KvMtf54l0d"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install langchain==0.2.0\n",
        "!pip install langchain-openai==0.1.7\n",
        "!pip install langchain-community==0.2.0"
      ],
      "metadata": {
        "id": "2evPp14fy258"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install Chroma Vector DB and LangChain wrapper"
      ],
      "metadata": {
        "id": "TlfidBdQZRGj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install langchain-chroma"
      ],
      "metadata": {
        "id": "uZKQDgQURhmF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Enter Open AI API Key"
      ],
      "metadata": {
        "id": "H9c37cLnSrbg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from getpass import getpass\n",
        "\n",
        "OPENAI_KEY = getpass('Enter Open AI API Key: ')"
      ],
      "metadata": {
        "id": "cv3JzCEx_PAd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup Environment Variables"
      ],
      "metadata": {
        "id": "1T0s0um5Svfa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.environ['OPENAI_API_KEY'] = OPENAI_KEY"
      ],
      "metadata": {
        "id": "x1YSuHNF_lbh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Open AI Embedding Models\n",
        "\n",
        "LangChain enables us to access Open AI embedding models which include the newest models: a smaller and highly efficient `text-embedding-3-small` model, and a larger and more powerful `text-embedding-3-large` model."
      ],
      "metadata": {
        "id": "M8nHAP7XOGOj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_openai import OpenAIEmbeddings\n",
        "\n",
        "# details here: https://openai.com/blog/new-embedding-models-and-api-updates\n",
        "openai_embed_model = OpenAIEmbeddings(model='text-embedding-3-small')"
      ],
      "metadata": {
        "id": "jzrIVI2NAHC1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Vector Databases\n",
        "\n",
        "One of the most common ways to store and search over unstructured data is to embed it and store the resulting embedding vectors, and then at query time to embed the unstructured query and retrieve the embedding vectors that are 'most similar' to the embedded query. A vector database takes care of storing embedded data and performing vector search for you.\n",
        "\n",
        "### Chroma Vector DB\n",
        "\n",
        "[Chroma](https://docs.trychroma.com/getting-started) is a AI-native open-source vector database focused on developer productivity and happiness. Chroma is licensed under Apache 2.0."
      ],
      "metadata": {
        "id": "IO3_BzTbTBcJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Create a Vector DB and persist on disk\n",
        "\n",
        "Here we initialize a connection to a Chroma vector DB client, and also we want to save to disk, so we simply initialize the Chroma client and pass the directory where we want the data to be saved to."
      ],
      "metadata": {
        "id": "-PnV9lAXZw9a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "docs = [\n",
        " 'Quantum mechanics describes the behavior of very small particles.',\n",
        " 'Photosynthesis is the process by which green plants make food using sunlight.',\n",
        " 'Artificial Intelligence aims to create machines that can think and learn.',\n",
        " 'The pyramids of Egypt are historical monuments that have stood for thousands of years.',\n",
        " 'New Delhi is the capital of India and the seat of all three branches of the Government of India.',\n",
        " 'Biology is the study of living organisms and their interactions with the environment.',\n",
        " 'Music therapy can aid in the mental well-being of individuals.',\n",
        " 'Mumbai is the financial capital and the most populous city of India. It is the financial, commercial, and entertainment capital of South Asia.',\n",
        " 'The Milky Way is just one of billions of galaxies in the universe.',\n",
        " 'Economic theories help understand the distribution of resources in society.',\n",
        " 'Kolkata is the de facto cultural capital of India and a historically and culturally significant city. Calcutta served as the de facto capital of India until 1911.',\n",
        " 'Yoga is an ancient practice that involves physical postures and meditation.'\n",
        "]"
      ],
      "metadata": {
        "id": "oABVR8m-6ZiY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_chroma import Chroma\n",
        "\n",
        "# create vector DB of docs and embeddings - takes 1 min on Colab\n",
        "chroma_db = Chroma.from_texts(texts=docs, collection_name='db_docs',\n",
        "                              # need to set the distance function to cosine else it uses euclidean by default\n",
        "                              # check https://docs.trychroma.com/guides#changing-the-distance-function\n",
        "                              collection_metadata={\"hnsw:space\": \"cosine\"},\n",
        "                              embedding=openai_embed_model)"
      ],
      "metadata": {
        "id": "kRYfcrsHUxyZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Vector Database Retrievers\n",
        "\n",
        "Here we will explore the following retrieval strategies on our Vector Database:\n",
        "\n",
        "- Similarity or Ranking based Retrieval\n",
        "- Similarity with Threshold Retrieval\n",
        "- Custom Retriever with Similarity Scores + Thresholding\n",
        "- Multi Query Retrieval\n",
        "- Contextual Compression Retrieval\n",
        "- Ensemble Retrieval"
      ],
      "metadata": {
        "id": "bprZC4S6TLfj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Similarity or Ranking based Retrieval\n",
        "\n",
        "We use cosine similarity here and retrieve the top 3 similar documents based on the user input query"
      ],
      "metadata": {
        "id": "269fFuebBr5P"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "similarity_retriever = chroma_db.as_retriever(search_type=\"similarity\",\n",
        "                                              search_kwargs={\"k\": 3})"
      ],
      "metadata": {
        "id": "ST7PDWG2b0ln"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the capital of India?\"\n",
        "top3_docs = similarity_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "4czRBjWLe8JH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the old capital of India?\"\n",
        "top3_docs = similarity_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "CYGmZdV--qz_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"how do plants make food?\"\n",
        "top3_docs = similarity_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "ZJmJRjlV_kvS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We use maximum marginal relevance ranking here and retrieve the top 3 similar documents based on the user input query"
      ],
      "metadata": {
        "id": "Wk2cdAjpB5TE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mmr_retriever = chroma_db.as_retriever(search_type=\"mmr\",\n",
        "                                       search_kwargs={\"k\": 3,\n",
        "                                                      'fetch_k': 10})"
      ],
      "metadata": {
        "id": "TLlTPW5vfGIo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the capital of India?\"\n",
        "top3_docs = mmr_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "mHztMoMDfU2D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the old capital of India?\"\n",
        "top3_docs = mmr_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "I3wDFqdW_tA8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"how do plants make food?\"\n",
        "top3_docs = mmr_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "XrY8wGD__xR8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Similarity with Threshold Retrieval\n",
        "\n",
        "We use cosine similarity here and retrieve the top 3 similar documents based on the user input query and also introduce a cutoff to not return any documents which are below a certain similarity threshold"
      ],
      "metadata": {
        "id": "8foBD2xmCDYc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "similarity_threshold_retriever = chroma_db.as_retriever(search_type=\"similarity_score_threshold\",\n",
        "                                                        search_kwargs={\"k\": 3,\n",
        "                                                                       \"score_threshold\": 0.3})"
      ],
      "metadata": {
        "id": "6ROSNwqeCMRS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the capital of India?\"\n",
        "top3_docs = similarity_threshold_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "Nv93k_QpCZv7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the old capital of India?\"\n",
        "top3_docs = similarity_threshold_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "fpu7Wzo6Cre9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"how do plants make food?\"\n",
        "top3_docs = similarity_threshold_retriever.invoke(query)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "nmALobtYCy0-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Custom Retriever with Similarity Scores + Thresholding\n",
        "\n",
        "Here we will create a custom retriever which will:\n",
        "\n",
        "- Retrieve documents with cosine distance\n",
        "- Convert to similarity score and apply thresholding\n",
        "- Return topk documents above a similarity threshold"
      ],
      "metadata": {
        "id": "05zZy9jCLEbC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "query = 'how do plants make food?'\n",
        "chroma_db.similarity_search_with_score(query, k=3)"
      ],
      "metadata": {
        "id": "hwU_kMxIKQBG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "chroma_db._select_relevance_score_fn()"
      ],
      "metadata": {
        "id": "ArhLcC32J-OC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# converts cosine distance to similarity; cosine_similarity = 1 - cosine_distance\n",
        "cosine_sim = chroma_db._select_relevance_score_fn()\n",
        "cosine_sim(0.35375)"
      ],
      "metadata": {
        "id": "3WlBCDeTKLM2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import List\n",
        "from langchain_core.documents import Document\n",
        "from langchain_core.runnables import chain\n",
        "\n",
        "@chain\n",
        "def custom_retriever(query: str, topk=3, threshold_score=0.3) -> List[Document]:\n",
        "    # get similarity conversion function (converts cosine distance to similarity)\n",
        "    cosine_sim = chroma_db._select_relevance_score_fn()\n",
        "    # get topk documents with lowest cosine distance\n",
        "    docs, scores = zip(*chroma_db.similarity_search_with_score(query, k=topk))\n",
        "    final_docs = []\n",
        "    for doc, score in zip(docs, scores):\n",
        "        # convert cosine distance to similarity\n",
        "        score = cosine_sim(score)\n",
        "        doc.metadata[\"score\"] = round(score, 3)\n",
        "        # check if score is above threshold\n",
        "        if score > threshold_score:\n",
        "            final_docs.append(doc)\n",
        "\n",
        "    return final_docs"
      ],
      "metadata": {
        "id": "YgOMpLEnDnoi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the financial capital of India?\"\n",
        "top3_docs = custom_retriever.invoke(query, topk=3, threshold_score=0.51)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "Tt2ihp3yDFaE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = 'how do plants make food?'\n",
        "top3_docs = custom_retriever.invoke(query, topk=3, threshold_score=0.5)\n",
        "top3_docs"
      ],
      "metadata": {
        "id": "kLk3PAu5EOGJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Multi Query Retrieval\n",
        "\n",
        "Retrieval may produce different results with subtle changes in query wording, or if the embeddings do not capture the semantics of the data well. Prompt engineering / tuning is sometimes done to manually address these problems, but can be tedious.\n",
        "\n",
        "The [`MultiQueryRetriever`](https://api.python.langchain.com/en/latest/retrievers/langchain.retrievers.multi_query.MultiQueryRetriever.html) automates the process of prompt tuning by using an LLM to generate multiple queries from different perspectives for a given user input query. For each query, it retrieves a set of relevant documents and takes the unique union across all queries to get a larger set of potentially relevant documents."
      ],
      "metadata": {
        "id": "6HKdIe2_wIBR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain_openai import ChatOpenAI\n",
        "\n",
        "chatgpt = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)"
      ],
      "metadata": {
        "id": "qmfVRrRujxrA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
        "# Set logging for the queries\n",
        "import logging\n",
        "\n",
        "similarity_retriever = chroma_db.as_retriever(search_type=\"similarity\",\n",
        "                                              search_kwargs={\"k\": 3})\n",
        "\n",
        "mq_retriever = MultiQueryRetriever.from_llm(\n",
        "    retriever=similarity_retriever, llm=chatgpt\n",
        ")\n",
        "\n",
        "logging.basicConfig()\n",
        "# so we can see what queries are generated by the LLM\n",
        "logging.getLogger(\"langchain.retrievers.multi_query\").setLevel(logging.INFO)"
      ],
      "metadata": {
        "id": "s3gU4HPvfjs7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"financial capital of India?\"\n",
        "docs = mq_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "KZJBIsl8xC_N"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"old capital of India?\"\n",
        "docs = mq_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "TVBeusuSyVo3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Contextual Compression Retrieval\n",
        "\n",
        "The information most relevant to a query may be buried in a document with a lot of irrelevant text. Passing that full document through your application can lead to more expensive LLM calls and poorer responses.\n",
        "\n",
        "Contextual compression is meant to fix this. The idea is simple: instead of immediately returning retrieved documents as-is, you can compress them using the context of the given query, so that only the relevant information is returned.\n",
        "\n",
        "This compression can happen in the form of:\n",
        "\n",
        "- Remove parts of the content of retrieved documents which are not relevant to the query. This is done by extracting only relevant parts of the document to the given query\n",
        "\n",
        "- Filter out documents which are not relevant to the given query but do not remove content from the document"
      ],
      "metadata": {
        "id": "qKkMcQpf3lBs"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Here we wrap our base cosine distance retriever with a `ContextualCompressionRetriever`. Then we'll add an `LLMChainExtractor`, which will iterate over the initially returned documents and extract from each only the content that is relevant to the query."
      ],
      "metadata": {
        "id": "pNMpbsi64QcM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.retrievers import ContextualCompressionRetriever\n",
        "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
        "\n",
        "# simple cosine distance based retriever\n",
        "similarity_retriever = chroma_db.as_retriever(search_type=\"similarity\",\n",
        "                                              search_kwargs={\"k\": 3})\n",
        "\n",
        "# extracts from each document only the content that is relevant to the query\n",
        "compressor = LLMChainExtractor.from_llm(llm=chatgpt)\n",
        "\n",
        "# retrieves the documents similar to query and then applies the compressor\n",
        "compression_retriever = ContextualCompressionRetriever(\n",
        "    base_compressor=compressor, base_retriever=similarity_retriever\n",
        ")"
      ],
      "metadata": {
        "id": "b14n_2AHykTf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the financial capital of India?\"\n",
        "docs = compression_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "UxoVG1dt0GQb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the old capital of India?\"\n",
        "docs = compression_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "zW64q3FF0M-b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `LLMChainFilter` is slightly simpler but more robust compressor that uses an LLM chain to decide which of the initially retrieved documents to filter out and which ones to return, without manipulating the document contents."
      ],
      "metadata": {
        "id": "URTONl5c4aLe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.retrievers.document_compressors import LLMChainFilter\n",
        "\n",
        "#  decides which of the initially retrieved documents to filter out and which ones to return\n",
        "_filter = LLMChainFilter.from_llm(llm=chatgpt)\n",
        "\n",
        "# retrieves the documents similar to query and then applies the filter\n",
        "compression_retriever = ContextualCompressionRetriever(\n",
        "    base_compressor=_filter, base_retriever=similarity_retriever\n",
        ")"
      ],
      "metadata": {
        "id": "fs-INms_0VTX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the financial capital of India?\"\n",
        "docs = compression_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "bgKB3_lV17Hh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the old capital of India?\"\n",
        "docs = compression_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "AyhPhAld1-H3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"how do plants live?\"\n",
        "docs = compression_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "zGdq0G7v2FAc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Ensemble Retrieval\n",
        "\n",
        "The `EnsembleRetriever` takes a list of retrievers as input and ensemble the results of each of their retrievals and rerank the results based on the [Reciprocal Rank Fusion](https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf) algorithm."
      ],
      "metadata": {
        "id": "9YK8BBMBAU9o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from langchain.retrievers import EnsembleRetriever\n",
        "\n",
        "# simple cosine distance based retriever\n",
        "similarity_retriever = chroma_db.as_retriever(search_type=\"similarity\",\n",
        "                                              search_kwargs={\"k\": 3})\n",
        "\n",
        "# retrieves the documents similar to query and then applies the filter\n",
        "compression_retriever = ContextualCompressionRetriever(\n",
        "    base_compressor=_filter, base_retriever=similarity_retriever\n",
        ")\n",
        "\n",
        "\n",
        "ensemble_retriever = EnsembleRetriever(\n",
        "    retrievers=[compression_retriever, similarity_retriever],\n",
        "    weights=[0.7, 0.3]\n",
        ")"
      ],
      "metadata": {
        "id": "UgZ4eMFc2GgI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"what is the financial capital of India?\"\n",
        "docs = ensemble_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "8dVz4JGVBQTO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "query = \"how do plants live?\"\n",
        "docs = ensemble_retriever.invoke(query)\n",
        "docs"
      ],
      "metadata": {
        "id": "dFtoRwndBZuc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will cover more retrieval strategies in the Advanced Course on RAG and LangChain in the future including:\n",
        "\n",
        "- Self Query Retrieval\n",
        "- Hybrid Search Retrieval\n",
        "- Parent Document Retrieval\n",
        "- Reranker Retrieval\n",
        "\n",
        "and more..."
      ],
      "metadata": {
        "id": "OfQj71WHCKc8"
      }
    }
  ]
}