ai-rag-lab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Lab Documentation and Solutions](https://img.shields.io/badge/Lab%20Documentation%20and%20Solutions-purple)](https://mongodb-developer.github.io/ai-rag-lab/)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Setup prerequisites\n",
    "\n",
    "Replace `<MONGODB_URI>` with your **MongoDB connection string**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pymongo import MongoClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retain the quotes (\"\") when pasting the URI\n",
    "MONGODB_URI = \"<MONGODB_URI>\"\n",
    "# Initialize a MongoDB Python client\n",
    "mongodb_client = MongoClient(MONGODB_URI, appname=\"devrel.workshop.rag\")\n",
    "# Check the connection to the server\n",
    "mongodb_client.admin.command(\"ping\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SERVERLESS_URL = os.environ.get(\"SERVERLESS_URL\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Download the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the `mongodb-docs` dataset from Hugging Face\n",
    "data = load_dataset(\"mongodb/mongodb-docs\", split=\"train\")\n",
    "# Convert the dataset into a dataframe first, then into a list of Python objects/dictionaries\n",
    "docs = pd.DataFrame(data).to_dict(\"records\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note the number of documents in the dataset\n",
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preview a document to understand its structure\n",
    "docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Chunk up the data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from typing import Dict, List"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Common list of separators for text data\n",
    "separators = [\"\\n\\n\", \"\\n\", \" \", \"\", \"#\", \"##\", \"###\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the `RecursiveCharacterTextSplitter` from LangChain to first split a piece of text on the list of `separators` above.\n",
    "# Then recursively merge them into tokens until the specified chunk size is reached.\n",
    "# For text data, you typically want to keep 1-2 paragraphs (~200 tokens) in a single chunk.\n",
    "# Chunk overlap of 15-20% of the chunk size is recommended to maintain context between chunks.\n",
    "# Pass the `separators` list above as an argument called `separators`\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    model_name=\"gpt-4\", separators=separators, chunk_size=200, chunk_overlap=30\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://api.python.langchain.com/en/latest/character/langchain_text_splitters.character.RecursiveCharacterTextSplitter.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_chunks(doc: Dict, text_field: str) -> List[Dict]:\n",
    "    \"\"\"\n",
    "    Chunk up a document.\n",
    "\n",
    "    Args:\n",
    "        doc (Dict): Parent document to generate chunks from.\n",
    "        text_field (str): Text field to chunk.\n",
    "\n",
    "    Returns:\n",
    "        List[Dict]: List of chunked documents.\n",
    "    \"\"\"\n",
    "    # Extract the field to chunk from `doc`\n",
    "    text = doc[text_field]\n",
    "    # Split `text` using the appropriate method of the `RecursiveCharacterTextSplitter` class\n",
    "    # NOTE: `text` is a string\n",
    "    chunks = <CODE_BLOCK_1>\n",
    "\n",
    "    # Iterate through `chunks` and for each chunk:\n",
    "    # 1. Create a shallow copy of `doc`, call it `temp`\n",
    "    # 2. Set the `text_field` field in `temp` to the content of the chunk\n",
    "    # 3. Append `temp` to `chunked_data`\n",
    "    chunked_data = []\n",
    "    for chunk in chunks:\n",
    "        temp = doc.copy()\n",
    "        temp[text_field]=chunk\n",
    "        chunked_data.append(temp)\n",
    "\n",
    "    return chunked_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split_docs = []\n",
    "# Iterate through `docs`, use the `get_chunks` function to chunk up the documents based on the \"body\" field, and add the list of chunked documents to `split_docs` initialized above.\n",
    "for doc in docs:\n",
    "    chunks = <CODE_BLOCK_2>\n",
    "    split_docs.extend(chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notice that the length of `split_docs` is greater than the length of `docs` from Step 2 above\n",
    "# This is because `split_docs` contains smaller chunks of the documents in `docs`\n",
    "len(split_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Preview a chunked document to understand its structure\n",
    "# Note that the structure looks similar to the original docs, except the `body` field now contains smaller chunks of text\n",
    "split_docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Generate embeddings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You may see a warning upon running this cell. You can ignore it.\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the `gte-small` model using the Sentence Transformers library\n",
    "embedding_model = SentenceTransformer(\"thenlper/gte-small\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://huggingface.co/thenlper/gte-small#usage (See \"Use with sentence-transformers\" under Usage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function that takes a piece of text (`text`) as input, embeds it using the `embedding_model` instantiated above and returns the embedding as a list\n",
    "# An array can be converted to a list using the `tolist()` method\n",
    "def get_embedding(text: str) -> List[float]:\n",
    "    \"\"\"\n",
    "    Generate the embedding for a piece of text.\n",
    "\n",
    "    Args:\n",
    "        text (str): Text to embed.\n",
    "\n",
    "    Returns:\n",
    "        List[float]: Embedding of the text as a list.\n",
    "    \"\"\"\n",
    "    embedding = <CODE_BLOCK_3>\n",
    "    return embedding.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded_docs = []\n",
    "# Add an `embedding` field to each dictionary in `split_docs`\n",
    "# The `embedding` field should correspond to the embedding of the value of the `body` field\n",
    "# Use the `get_embedding` function defined above to generate the embedding\n",
    "# Append the updated dictionaries to `embedded_docs` initialized above.\n",
    "for doc in tqdm(split_docs):\n",
    "    <CODE_BLOCK_4>\n",
    "    embedded_docs.append(doc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check that the length of `embedded_docs` is the same as that of `split_docs`\n",
    "len(embedded_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Ingest data into MongoDB\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Name of the database -- Change if needed or leave as is\n",
    "DB_NAME = \"mongodb_genai_devday\"\n",
    "# Name of the collection -- Change if needed or leave as is\n",
    "COLLECTION_NAME = \"knowledge_base\"\n",
    "# Name of the vector search index -- Change if needed or leave as is\n",
    "ATLAS_VECTOR_SEARCH_INDEX_NAME = \"vector_index\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://pymongo.readthedocs.io/en/stable/tutorial.html#getting-a-database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to the `DB_NAME` database.\n",
    "# Use the `mongodb_client` instantiated in Step 1.\n",
    "db = <CODE_BLOCK_5>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://pymongo.readthedocs.io/en/stable/tutorial.html#getting-a-collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Connect to the `COLLECTION_NAME` collection.\n",
    "# Use the `db` and collection name defined above.\n",
    "collection = <CODE_BLOCK_6>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bulk delete all existing records from the collection defined above\n",
    "collection.delete_many({})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://pymongo.readthedocs.io/en/stable/examples/bulk.html#bulk-insert\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bulk insert `embedded_docs` into the collection defined above -- should be a one-liner\n",
    "<CODE_BLOCK_7>\n",
    "\n",
    "print(f\"Ingested {collection.count_documents({})} documents into the {COLLECTION_NAME} collection.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Create a vector search index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create vector index definition specifying:\n",
    "# path: Path to the embeddings field\n",
    "# numDimensions: Number of embedding dimensions- depends on the embedding model used\n",
    "# similarity: Similarity metric. One of cosine, euclidean, dotProduct.\n",
    "model = {\n",
    "    \"name\": ATLAS_VECTOR_SEARCH_INDEX_NAME,\n",
    "    \"type\": \"vectorSearch\",\n",
    "    \"definition\": {\n",
    "        \"fields\": [\n",
    "            {\n",
    "                \"type\": \"vector\",\n",
    "                \"path\": \"embedding\",\n",
    "                \"numDimensions\": 384,\n",
    "                \"similarity\": \"cosine\",\n",
    "            }\n",
    "        ]\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.create_search_index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a vector search index with the above definition for the `collection` collection\n",
    "<CODE_BLOCK_8>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 7: Perform vector search on your data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a vector search function\n",
    "\n",
    "📚 https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#ann-examples (Refer to the \"Basic Example\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to retrieve relevant documents for a user query using vector search\n",
    "def vector_search(user_query: str) -> List[Dict]:\n",
    "    \"\"\"\n",
    "    Retrieve relevant documents for a user query using vector search.\n",
    "\n",
    "    Args:\n",
    "    user_query (str): The user's query string.\n",
    "\n",
    "    Returns:\n",
    "    list: A list of matching documents.\n",
    "    \"\"\"\n",
    "\n",
    "    # Generate embedding for the `user_query` using the `get_embedding` function defined in Step 4\n",
    "    query_embedding = <CODE_BLOCK_9>\n",
    "\n",
    "    # Define an aggregation pipeline consisting of a $vectorSearch stage, followed by a $project stage\n",
    "    # Set the number of candidates to 150 and only return the top 5 documents from the vector search\n",
    "    # In the $project stage, exclude the `_id` field and include only the `body` field and `vectorSearchScore`\n",
    "    # NOTE: Use variables defined previously for the `index`, `queryVector` and `path` fields in the $vectorSearch stage\n",
    "    pipeline = <CODE_BLOCK_10>\n",
    "\n",
    "    # Execute the aggregation `pipeline` and store the results in `results`\n",
    "    results = <CODE_BLOCK_11>\n",
    "    return list(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run vector search queries\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_search(\"What are some best practices for data backups in MongoDB?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_search(\"How to resolve alerts in MongoDB?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🦹‍♀️ Combine pre-filtering with vector search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter for documents where the product name is `MongoDB Atlas`\n",
    "\n",
    "📚 https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#about-the-filter-type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the vector search index `model` from Step 6 to include the `metadata.productName` field as a `filter` field\n",
    "model = <CODE_BLOCK_12>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update vector search index\n",
    "collection.update_search_index(\n",
    "    name=ATLAS_VECTOR_SEARCH_INDEX_NAME, definition=model[\"definition\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE:** Check that the index update is complete before proceeding. To do so, ensure that the status of the index shows \"Ready\" in the Atlas UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed the user query\n",
    "query_embedding = get_embedding(\n",
    "    \"What are some best practices for data backups in MongoDB?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-stage/#ann-examples (Refer to the \"Filter Example\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the $vectorSearch stage of the aggregation pipeline defined previously to include a filter for documents where the `metadata.productName` field has the value \"MongoDB Atlas\"\n",
    "pipeline = <CODE_BLOCK_13>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute the aggregation pipeline and view the results\n",
    "results = collection.aggregate(pipeline)\n",
    "list(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Filter on documents which have been updated on or after `2024-05-19` and where the content type is `Tutorial`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the vector search index `model` from Step 6 to include `metadata.contentType` and `updated` as `filter` fields\n",
    "model = <CODE_BLOCK_14>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update vector search index\n",
    "collection.update_search_index(\n",
    "    name=ATLAS_VECTOR_SEARCH_INDEX_NAME, definition=model[\"definition\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE:** Check that the index update is complete before proceeding. To do so, ensure that the status of the index shows \"Ready\" in the Atlas UI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed the user query\n",
    "query_embedding = get_embedding(\n",
    "    \"What are some best practices for data backups in MongoDB?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the $vectorSearch stage of the aggregation pipeline defined previously to include a filter for documents where\n",
    "# the `metadata.contentType` field has the value \"Tutorial\"\n",
    "# AND\n",
    "# the `updated` field is greater than or equal to \"2024-05-19\"\n",
    "pipeline = <CODE_BLOCK_15>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute the aggregation pipeline and view the results\n",
    "results = collection.aggregate(pipeline)\n",
    "list(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 8: Build the RAG application\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function to create the chat prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to create the user prompt for our RAG application\n",
    "def create_prompt(user_query: str) -> str:\n",
    "    \"\"\"\n",
    "    Create a chat prompt that includes the user query and retrieved context.\n",
    "\n",
    "    Args:\n",
    "        user_query (str): The user's query string.\n",
    "\n",
    "    Returns:\n",
    "        str: The chat prompt string.\n",
    "    \"\"\"\n",
    "    # Retrieve the most relevant documents for the `user_query` using the `vector_search` function defined in Step 7\n",
    "    context = <CODE_BLOCK_16>\n",
    "    # Join the retrieved documents into a single string, where each document is separated by two new lines (\"\\n\\n\")\n",
    "    context = \"\\n\\n\".join([doc.get('body') for doc in context])\n",
    "    # Prompt consisting of the question and relevant context to answer it\n",
    "    prompt = f\"Answer the question based only on the following context. If the context is empty, say I DON'T KNOW\\n\\nContext:\\n{context}\\n\\nQuestion:{user_query}\"\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function to answer user queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to answer user queries\n",
    "def generate_answer(user_query: str) -> None:\n",
    "    \"\"\"\n",
    "    Generate an answer to the user query.\n",
    "\n",
    "    Args:\n",
    "        user_query (str): The user's query string.\n",
    "    \"\"\"\n",
    "    # Use the `create_prompt` function above to create a chat prompt\n",
    "    prompt = <CODE_BLOCK_17>\n",
    "    # Format the message to the LLM in the format {\"role\": <role_value>, \"content\": <content_value>}\n",
    "    # The role value for user messages must be \"user\"\n",
    "    # Use the `prompt` created above to populate the `content` field in the chat message\n",
    "    messages = [<CODE_BLOCK_18>]\n",
    "    # Send the chat messages to a serverless function to get back an LLM response\n",
    "    response = requests.post(url=SERVERLESS_URL, json={\"task\": \"completion\", \"data\": messages})\n",
    "    # Print the final answer\n",
    "    print(response.json()[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Query the RAG application\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_answer(\"What are some best practices for data backups in MongoDB?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_answer(\"What did I just ask you?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🦹‍♀️ Re-rank retrieved results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import CrossEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rerank_model = CrossEncoder(\"mixedbread-ai/mxbai-rerank-xsmall-v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://huggingface.co/mixedbread-ai/mxbai-rerank-xsmall-v1#quickstart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add a re-ranking step to the following function\n",
    "def create_prompt(user_query: str) -> str:\n",
    "    \"\"\"\n",
    "    Create a chat prompt that includes the user query and retrieved context.\n",
    "\n",
    "    Args:\n",
    "        user_query (str): The user's query string.\n",
    "\n",
    "    Returns:\n",
    "        str: The chat prompt string.\n",
    "    \"\"\"\n",
    "    # Retrieve the most relevant documents for the `user_query` using the `vector_search` function defined in Step 7\n",
    "    context = vector_search(user_query)\n",
    "    # Extract the \"body\" field from each document in `context`\n",
    "    documents = [d.get(\"body\") for d in context]\n",
    "    # Use the `rerank_model` instantiated above to re-rank `documents`\n",
    "    # Set the `top_k` argument to 5 \n",
    "    reranked_documents = <CODE_BLOCK_19>\n",
    "    # Join the re-ranked documents into a single string, where each document is separated by two new lines (\"\\n\\n\")\n",
    "    context = \"\\n\\n\".join([d.get(\"text\", \"\") for d in reranked_documents])\n",
    "    # Prompt consisting of the question and relevant context to answer it\n",
    "    prompt = f\"Answer the question based only on the following context. If the context is empty, say I DON'T KNOW\\n\\nContext:\\n{context}\\n\\nQuestion:{user_query}\"\n",
    "    return prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note the impact of re-ranking on the generated answer\n",
    "generate_answer(\"What are some best practices for data backups in MongoDB?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 9: Add memory to the RAG application\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history_collection = mongodb_client[DB_NAME][\"chat_history\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "📚 https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.create_index\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an index on the key `session_id` for the `history_collection` collection\n",
    "<CODE_BLOCK_20>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function to store chat messages in MongoDB\n",
    "\n",
    "📚 https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.insert_one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def store_chat_message(session_id: str, role: str, content: str) -> None:\n",
    "    \"\"\"\n",
    "    Store a chat message in a MongoDB collection.\n",
    "\n",
    "    Args:\n",
    "        session_id (str): Session ID of the message.\n",
    "        role (str): Role for the message. One of `system`, `user` or `assistant`.\n",
    "        content (str): Content of the message.\n",
    "    \"\"\"\n",
    "    # Create a message object with `session_id`, `role`, `content` and `timestamp` fields\n",
    "    # `timestamp` should be set the current timestamp\n",
    "    message = {\n",
    "        \"session_id\": session_id,\n",
    "        \"role\": role,\n",
    "        \"content\": content,\n",
    "        \"timestamp\": datetime.now(),\n",
    "    }\n",
    "    # Insert the `message` into the `history_collection` collection\n",
    "    <CODE_BLOCK_21>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function to retrieve chat history from MongoDB\n",
    "\n",
    "📚 https://pymongo.readthedocs.io/en/stable/api/pymongo/cursor.html#pymongo.cursor.Cursor.sort"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def retrieve_session_history(session_id: str) -> List:\n",
    "    \"\"\"\n",
    "    Retrieve chat message history for a particular session.\n",
    "\n",
    "    Args:\n",
    "        session_id (str): Session ID to retrieve chat message history for.\n",
    "\n",
    "    Returns:\n",
    "        List: List of chat messages.\n",
    "    \"\"\"\n",
    "    # Query the `history_collection` collection for documents where the \"session_id\" field has the value of the input `session_id`\n",
    "    # Sort the results in increasing order of the values in `timestamp` field\n",
    "    cursor =  <CODE_BLOCK_22>\n",
    "\n",
    "    if cursor:\n",
    "        # Iterate through the cursor and extract the `role` and `content` field from each entry\n",
    "        # Then format each entry as: {\"role\": <role_value>, \"content\": <content_value>}\n",
    "        messages = [{\"role\": msg[\"role\"], \"content\": msg[\"content\"]} for msg in cursor]\n",
    "    else:\n",
    "        # If cursor is empty, return an empty list\n",
    "        messages = []\n",
    "\n",
    "    return messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Handle chat history in the `generate_answer` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_answer(session_id: str, user_query: str) -> None:\n",
    "    \"\"\"\n",
    "    Generate an answer to the user's query taking chat history into account.\n",
    "\n",
    "    Args:\n",
    "        session_id (str): Session ID to retrieve chat history for.\n",
    "        user_query (str): The user's query string.\n",
    "    \"\"\"\n",
    "    # Initialize list of messages to pass to the chat completion model\n",
    "    messages = []\n",
    "\n",
    "    # Retrieve documents relevant to the user query and convert them to a single string\n",
    "    context = vector_search(user_query)\n",
    "    context = \"\\n\\n\".join([d.get(\"body\", \"\") for d in context])\n",
    "    # Create a system prompt containing the retrieved context\n",
    "    system_message = {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": f\"Answer the question based only on the following context. If the context is empty, say I DON'T KNOW\\n\\nContext:\\n{context}\",\n",
    "    }\n",
    "    # Append the system prompt to the `messages` list\n",
    "    messages.append(system_message)\n",
    "\n",
    "    # Use the `retrieve_session_history` function to retrieve message history from MongoDB for the session ID `session_id` \n",
    "    # And add all messages in the message history to the `messages` list \n",
    "    message_history = <CODE_BLOCK_23>\n",
    "    messages.extend(message_history)\n",
    "\n",
    "    # Format the user query in the format {\"role\": <role_value>, \"content\": <content_value>}\n",
    "    # The role value for user messages must be \"user\"\n",
    "    # And append the user message to the `messages` list\n",
    "    user_message = <CODE_BLOCK_24>\n",
    "    messages.append(user_message)\n",
    "\n",
    "    # Send the chat messages to a serverless function to get back an LLM response\n",
    "    response = requests.post(url=SERVERLESS_URL, json={\"task\": \"completion\", \"data\": messages})\n",
    "\n",
    "    # Extract the answer from the response\n",
    "    answer = response.json()[\"text\"]\n",
    "\n",
    "    # Use the `store_chat_message` function to store the user message and also the generated answer in the message history collection\n",
    "    # The role value for user messages is \"user\", and \"assistant\" for the generated answer\n",
    "    <CODE_BLOCK_25>\n",
    "\n",
    "    print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_answer(\n",
    "    session_id=\"1\",\n",
    "    user_query=\"What are some best practices for data backups in MongoDB?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_answer(\n",
    "    session_id=\"1\",\n",
    "    user_query=\"What did I just ask you?\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}