Cohere-MongoDB-Atlas-Vector-Search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dc56c0b7-1f5a-4294-82a0-4f621a21f517",
   "metadata": {},
   "source": [
    "# How to use Cohere embeddings and rerank module with MongoDB Atlas Search / Vector Search.\n",
    "Cohere embeddings are a way of representing text as high-dimensional vectors that capture the semantic meaning and context of the text. Rerank module is a feature of Cohere that allows you to customize the ranking of search results(retrieved from MongoDB Atlas) based on reranker module. Working with Cohere embeddings and rerank module, you can create powerful and flexible vector search applications/Information retrieval systems that can handle natural language queries and also address complex scenarios.\n",
    "\n",
    "The powerful combination of technologies you can further enhance existing operational data store(ODS) functionality of MongoDB with features such as semantic search, question-answering as in RAG framework, real time threat detection, and plethora of other applications that rely on NLP and Full text search capabilities over Real time and transactions data sources that generate text data a scale."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b1ce77-5ad1-46d9-b2ed-01cf5f274cc8",
   "metadata": {},
   "source": [
    "# What we will do in this tutorial?\n",
    "\n",
    "### Store Embeddings and Prepare the Index\n",
    "1. Use the Cohere Embed Jobs to generate vector embeddings for the first time on large datasets in an asynchronous and scheduled manner using the Cohere platform.\n",
    "2. Add vector embeddings into MongoDB Atlas, which can store and index these vector embeddings alongside your other operational/metadata. \n",
    "3. Finally, prepare the indexes for both vector embeddings as well as full-text search on our private dataset.\n",
    "\n",
    "### Search with vector embeddings\n",
    "1. Write a simple python function to accept search terms/phrases and pass it through the Cohere embed api again to get a query vector.\n",
    "2. Take these resultant query vector embedding and perform a vector search query using the `$vectorsearch` operator in the MongoDB Aggregation Pipeline.\n",
    "3. Pre-filter documents using meta information to narrow the search across your dataset; thereby speeding up the performance of vector search results while retaining accuracy.\n",
    "4. The retrieved semantically similar documents can be post-filtered(relevancy score) to demonstrate a higher degree of control over the semantic search behavior.\n",
    "\n",
    "### Search with text and rerank with Cohere\n",
    "1. Write a simple python function to accept search terms / phrases and prepare a search query using $search operator and MongoDB Aggregation Pipeline.\n",
    "2. Take these resultant documents perform reranking operation of the retrieved documents to achieve higher accuracy with full text search results using Cohere rerank module.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![MDB_vector_search_Cohere.png](MDB_vector_search_Cohere.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b08e9796-6d3f-45d5-bccc-3e6fcc79a6a9",
   "metadata": {},
   "source": [
    "## Install all dependecies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4d999aca-90dd-4f41-bbcd-f2b3ba28ef1b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: cohere in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (4.44)\n",
      "Requirement already satisfied: pymongo in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (4.6.1)\n",
      "Requirement already satisfied: pandas in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (1.4.4)\n",
      "Requirement already satisfied: aiohttp<4.0,>=3.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (3.8.3)\n",
      "Requirement already satisfied: backoff<3.0,>=2.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (2.2.1)\n",
      "Requirement already satisfied: fastavro<2.0,>=1.8 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (1.9.0)\n",
      "Requirement already satisfied: importlib_metadata<7.0,>=6.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (6.0.0)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.25.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (2.31.0)\n",
      "Requirement already satisfied: urllib3<3,>=1.26 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from cohere) (1.26.18)\n",
      "Requirement already satisfied: dnspython<3.0.0,>=1.16.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from pymongo) (2.3.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from pandas) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from pandas) (2022.7)\n",
      "Requirement already satisfied: numpy>=1.20.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from pandas) (1.24.3)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (22.1.0)\n",
      "Requirement already satisfied: charset-normalizer<3.0,>=2.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (2.0.4)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (6.0.2)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0.0a3 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (4.0.2)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (1.8.1)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (1.3.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from aiohttp<4.0,>=3.0->cohere) (1.2.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from importlib_metadata<7.0,>=6.0->cohere) (3.11.0)\n",
      "Requirement already satisfied: six>=1.5 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from python-dateutil>=2.8.1->pandas) (1.16.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from requests<3.0.0,>=2.25.0->cohere) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/ashwin.gangadhar/anaconda3/lib/python3.9/site-packages (from requests<3.0.0,>=2.25.0->cohere) (2023.5.7)\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "pip install cohere pymongo pandas s3fs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad4ad9e-a09d-413d-bdc1-74f136a71539",
   "metadata": {},
   "source": [
    "# Intitialize the packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21779e56-2654-417b-8135-4217d510ffb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import getpass\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c355cd43-ebb3-4b9e-975a-28d3e8a9d2db",
   "metadata": {},
   "source": [
    "# Initialize Cohere API Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dce2ab16-6c4f-467e-9e2a-6d58f3614770",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    cohere_api_key = os.environ[\"COHERE_API_KEY\"]\n",
    "except KeyError:\n",
    "    cohere_api_key = getpass.getpass(\"Please enter your COHERE API KEY (hit enter): \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1fa9dab-983f-4ca8-8d73-fff6a24f765a",
   "metadata": {},
   "source": [
    "# Initialize MongoDB Connection String"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5c75b317-3f06-4129-9e48-c2c5bb0dca3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    MONGO_CONN_STR = os.environ[\"MONGO_CONNECTION_STR\"]\n",
    "except KeyError:\n",
    "    MONGO_CONN_STR = getpass.getpass(\"Please enter your MongoDB Atlas Connection String (hit enter): \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76949067-52d4-4715-9cf1-71d867446616",
   "metadata": {},
   "source": [
    "# Load dataset from S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "86fcff30-e5cb-458d-8121-12d2b6e351ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>overview</th>\n",
       "      <th>title</th>\n",
       "      <th>release_date</th>\n",
       "      <th>vote_average</th>\n",
       "      <th>vote_count</th>\n",
       "      <th>adult</th>\n",
       "      <th>year</th>\n",
       "      <th>month</th>\n",
       "      <th>day</th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Led by Woody, Andy's toys live happily in his ...</td>\n",
       "      <td>Toy Story</td>\n",
       "      <td>1995-10-30</td>\n",
       "      <td>7.7</td>\n",
       "      <td>5415</td>\n",
       "      <td>False</td>\n",
       "      <td>1995</td>\n",
       "      <td>10</td>\n",
       "      <td>30</td>\n",
       "      <td>Title: Toy Story  Genres: Animation,Comedy,Fam...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>When siblings Judy and Peter discover an encha...</td>\n",
       "      <td>Jumanji</td>\n",
       "      <td>1995-12-15</td>\n",
       "      <td>6.9</td>\n",
       "      <td>2413</td>\n",
       "      <td>False</td>\n",
       "      <td>1995</td>\n",
       "      <td>12</td>\n",
       "      <td>15</td>\n",
       "      <td>Title: Jumanji  Genres: Animation,Comedy,Famil...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A family wedding reignites the ancient feud be...</td>\n",
       "      <td>Grumpier Old Men</td>\n",
       "      <td>1995-12-22</td>\n",
       "      <td>6.5</td>\n",
       "      <td>92</td>\n",
       "      <td>False</td>\n",
       "      <td>1995</td>\n",
       "      <td>12</td>\n",
       "      <td>22</td>\n",
       "      <td>Title: Grumpier Old Men  Genres: Animation,Com...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            overview             title  \\\n",
       "0  Led by Woody, Andy's toys live happily in his ...         Toy Story   \n",
       "1  When siblings Judy and Peter discover an encha...           Jumanji   \n",
       "2  A family wedding reignites the ancient feud be...  Grumpier Old Men   \n",
       "\n",
       "  release_date  vote_average  vote_count  adult  year  month  day  \\\n",
       "0   1995-10-30           7.7        5415  False  1995     10   30   \n",
       "1   1995-12-15           6.9        2413  False  1995     12   15   \n",
       "2   1995-12-22           6.5          92  False  1995     12   22   \n",
       "\n",
       "                                                text  \n",
       "0  Title: Toy Story  Genres: Animation,Comedy,Fam...  \n",
       "1  Title: Jumanji  Genres: Animation,Comedy,Famil...  \n",
       "2  Title: Grumpier Old Men  Genres: Animation,Com...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import s3fs\n",
    "df = pd.read_json(\"s3://ashwin-partner-bucket/cohere/movies_sample_dataset.jsonl\", orient=\"records\", lines=True)\n",
    "df.to_json(\"./movies_sample_dataset.jsonl\", orient='records', lines=True)\n",
    "df[:3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efc92716-cec1-4556-a16d-b5c0aa430894",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "713e18d5-a0e9-4a96-a69a-b5d6f1f91510",
   "metadata": {},
   "source": [
    "# Initialize and Schedule the Cohere Embeddings Job to embed the `sample_movies` dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6847c873-ab5f-49a2-843e-54fb46654188",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cohere  \n",
    "co_client = cohere.Client(cohere_api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a3380fb0-2977-413f-9f74-d676bfde0c6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "uploading file, starting validation...\n",
      "movies-wz7adk was uploaded\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "dataset = co_client.create_dataset(name='movies',\n",
    "                                   data=open(\"./movies_sample_dataset.jsonl\",'r'),\n",
    "                                   keep_fields=[\"overview\",\"title\",\"year\"],\n",
    "                                   dataset_type=\"embed-input\").wait()\n",
    "dataset.wait()\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "633a57ca-cb5b-482e-bb34-4a0014f396ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_job = co_client.create_embed_job(dataset_id=dataset.id, \n",
    "    input_type='search_document',\n",
    "    model='embed-english-v3.0', \n",
    "    truncate='END')\n",
    "embed_job.wait()\n",
    "output_dataset = co_client.get_dataset(embed_job.output.id)\n",
    "results = list(output_dataset)\n",
    "len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "94cdf891",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cohere.Dataset {\n",
       "\tid: embeded-movies-p4kcg2\n",
       "\tname: embeded-movies\n",
       "\tdataset_type: embed-result\n",
       "\tvalidation_status: validated\n",
       "\tcreated_at: 2024-01-23 10:54:50.794112\n",
       "\tupdated_at: 2024-01-23 10:54:50.794112\n",
       "\tdownload_urls: ['https://storage.googleapis.com/cohere-user/dataset-api-temp/8268d35d-90bc-4d70-88d1-3851ec01c1c1/e89b6118-1480-4765-b6a8-0819432a0932/embeded-movies-p4kcg2/001_embeded-movies.avro?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=dataset%40cohere-production.iam.gserviceaccount.com%2F20240123%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20240123T105749Z&X-Goog-Expires=259199&X-Goog-Signature=428f0194c4ea7bbcbf42e05dd37f0278b21452f8f797467834fa223c8c215e5ed2726525c4dd1779533b1e3865aa323adb92449ded1ed48e4502ff864542e47c23764ddc861e6834616f1acff52d46f1cdef16478aaa76a5d8f3f1bcc015cff6d3e4e3003c94ee621f61c04482e9ea0624f4ea317d05405221b6344136ef390bae9486bd30e02969577ecc123bc4acaa50ec6c2db7effe60a3ab8788fd66b60df0b5b8568e2a4d0c5eba822a6c14e29473201780b0a373b2b7025cf0b986076cb78165967527b2e24e585fcacef70191accf71667a0c64d6edf301c4b7628863e3d226b3544bb7cbe63df0eab7bb8400346c273a79d9e32f4d440e9cc005ca51&X-Goog-SignedHeaders=host', 'https://storage.googleapis.com/cohere-user/dataset-api-temp/8268d35d-90bc-4d70-88d1-3851ec01c1c1/e89b6118-1480-4765-b6a8-0819432a0932/embeded-movies-p4kcg2/002_embeded-movies.avro?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=dataset%40cohere-production.iam.gserviceaccount.com%2F20240123%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20240123T105749Z&X-Goog-Expires=259199&X-Goog-Signature=a15587271ec304febdfb7e06b9a98137be0c89351c0c00af6251b6386ec769d38842880d6b164dcb24d12ae4684c339e04b70077e9082738a2e73751c336b4317714b6f6e66bd090588509adf61dc2ff3f3bb0b57039ab3a9317264d4710288dfd44df766c0e41c0b7fe55d325626711faba35dc8fccc9392f6b5ac347bae994a1406251f5be32f011b29fe77810e5ec382d0e82bfd78ecd621814d4fa306e9875932655535ec46d796acbc388feeeadf0b4505d67e03fa4999947a1c269f8e847d308f91fad583f3f04b70c2981f9948996fee266a318da8b73b028f7af69ba789b202db485c47afce5a71b80dd5b7b6c00be0dac8dc1abad42f03f804a351d&X-Goog-SignedHeaders=host', 'https://storage.googleapis.com/cohere-user/dataset-api-temp/8268d35d-90bc-4d70-88d1-3851ec01c1c1/e89b6118-1480-4765-b6a8-0819432a0932/embeded-movies-p4kcg2/003_embeded-movies.avro?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=dataset%40cohere-production.iam.gserviceaccount.com%2F20240123%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20240123T105749Z&X-Goog-Expires=259199&X-Goog-Signature=932a32103c9ece7dd33913b5f281049ae771355551db80bcb34703059cad682fb44c58724d88b3a75cde5b24521e39bb0595876a343fe655388f2d89c8993840081e253919de6ea63e4a59f878f9bad27e2766bf5961a2a81febe035b52592e00ece09cbd9774f6622006ac1087caf91a528300226ec8d3b595a5afa6f0702742005668c1c7b921737483c56951248d66de2d15f3c491fba2f005587ed5c064e7a1198285e487aa0761d40907e87bab135cff293475e93aa18076ae445e9567738c779e771272398427f4e26be69f75cacfdfb35e3240dbebbf422fb672df3a7a8f19bb520becd88dffce04737dafd5d2f8c24d0a17abe1bb47c21289456ac5e&X-Goog-SignedHeaders=host', 'https://storage.googleapis.com/cohere-user/dataset-api-temp/8268d35d-90bc-4d70-88d1-3851ec01c1c1/e89b6118-1480-4765-b6a8-0819432a0932/embeded-movies-p4kcg2/004_embeded-movies.avro?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=dataset%40cohere-production.iam.gserviceaccount.com%2F20240123%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20240123T105749Z&X-Goog-Expires=259199&X-Goog-Signature=39add535aaba94b76558bb8c0cdfca529cd321de686ff8959b21d157c9d2ab98a7a90d1ce3cf0ab40931733acd67bc926eb173350a9eed8877c5493782678fdd69b63b35297c7ff2d2f9d2f9098910d4cbfee9471981d1cd2f80f884a7a3ec736a14ad3e7287356c41cd3c71f400b444ea2afe20f178e137c5c8fc10d692e4547c8dfd3827b2dcf622ddb86a4e8fc90cc919814de8095e6f06e7e980a2de036f98232b5086fba2f91d9fa81066f0885b3fc9f9088ed55bf8fef7d04c8f08fcfe6be21c1b3002b99b4fc3e838b7869adb759ad5c98fd0e2caac489baa003c9d03f4ff4ff846ca0ecee8d365bb5888283e6f4ce9e868b71538a426d2a9aa41b24f&X-Goog-SignedHeaders=host', 'https://storage.googleapis.com/cohere-user/dataset-api-temp/8268d35d-90bc-4d70-88d1-3851ec01c1c1/e89b6118-1480-4765-b6a8-0819432a0932/embeded-movies-p4kcg2/005_embeded-movies.avro?X-Goog-Algorithm=GOOG4-RSA-SHA256&X-Goog-Credential=dataset%40cohere-production.iam.gserviceaccount.com%2F20240123%2Fauto%2Fstorage%2Fgoog4_request&X-Goog-Date=20240123T105749Z&X-Goog-Expires=259199&X-Goog-Signature=b5427b56f491c315bedebc933f5fe98162962678f4656c61972df3d6419e4d8b576e95d4c5ec2c316045b5abe77d3ae2e735b5ea1b611c315426ed23e6482689b05705a0f18b42d925d7b43bc6a69d965c8601afa7a0aa3b32045b25237f9fefd8bdd2f80b47dceeaf8ee9b9fcf90d8f9bc8a3da30e699b3af963427fe948fcc9a5f3279fce6a1252fcce44a3fa9ad75ec70665713c84f2cf2dc8cb08cc5029e266aa31ec64b348b26516177cfac60c216edc8c2af2dfd081c7faba107fd9f10b9ff1d23fe47da64dfbde92f1f110f22270d13fb46778926ded3568cdbd082af857ea98258a7f8b7841469f7fd67f293971c4b52729f0b9ea6a02325e5b72ef8&X-Goog-SignedHeaders=host']\n",
       "\tvalidation_error: None\n",
       "\tvalidation_warnings: []\n",
       "}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a30964-c359-41ed-83d6-045aec5d1cc2",
   "metadata": {},
   "source": [
    "# How to Initialize MongoDB Atlas and Insert data to MongoDB collection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b41283ce-cc5f-4dd5-a44c-310c21e19b6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymongo import MongoClient\n",
    "mongo_client = MongoClient(MONGO_CONN_STR)\n",
    "\n",
    "output_collection = mongo_client[\"sample_mflix\"][\"cohere_embed_movies\"]\n",
    "if output_collection.count_documents({})>0:\n",
    "    output_collection.delete_many({})\n",
    "e = output_collection.insert_many(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "310065a6-2bb9-44b5-a04f-0339c97f23e1",
   "metadata": {},
   "source": [
    "# Programatically create Vector Search and Full text search index in MongoDB Atlas for the respective collection used in the tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "8731d226-9006-42fb-b992-3039dfbd20bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'default'"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_collection.create_search_index({\"definition\":\n",
    "        {\"mappings\":\n",
    "         {\"dynamic\": True,\n",
    "          \"fields\": {\n",
    "            \"embedding\" : {\n",
    "                \"dimensions\": 1024,\n",
    "                \"similarity\": \"cosine\",\n",
    "                \"type\": \"vector\"\n",
    "                }\n",
    "            }}},\n",
    "     \"name\": \"default\"\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ab0cfa4-450c-44ee-92d0-25053cf08f42",
   "metadata": {},
   "source": [
    "# Python function to embed search term/phrase and perform vectors search query using `$vectorsearch` operator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "45008d86-4fd9-4675-9b6e-5361d7b9d30f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_vector_search(q, prefilter = {}, postfilter = {},path=\"embedding\",topK=2):\n",
    "    ele = co_client.embed(model=\"embed-english-v3.0\",input_type=\"search_document\",texts=[q])\n",
    "    query_embedding = ele.embeddings[0]\n",
    "    vs_query = {\n",
    "                \"index\": \"default\",\n",
    "                \"path\": path,\n",
    "                \"queryVector\": query_embedding,\n",
    "                \"numCandidates\": 10,\n",
    "                \"limit\": topK,\n",
    "            }\n",
    "    if len(prefilter)>0:\n",
    "        vs_query[\"filter\"] = prefilter\n",
    "    new_search_query = {\"$vectorSearch\": vs_query}\n",
    "    project = {\"$project\": {\"score\": {\"$meta\": \"vectorSearchScore\"},\"_id\": 0,\"title\": 1, \"release_date\": 1, \"overview\": 1,\"year\": 1}}\n",
    "    if len(postfilter.keys())>0:\n",
    "        postFilter = {\"$match\":postfilter}\n",
    "        res = list(output_collection.aggregate([new_search_query, project, postFilter]))\n",
    "    else:\n",
    "        res = list(output_collection.aggregate([new_search_query, project]))\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a7cb915-2d33-4e04-959e-22573addf3bf",
   "metadata": {},
   "source": [
    "# Vector search query example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d7b1af8d-eb2f-47d4-b5fe-e7cbcc37fa61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'overview': 'German Comedy',\n",
       "  'title': 'Honeymoon',\n",
       "  'year': 1996,\n",
       "  'score': 0.7810907959938049},\n",
       " {'overview': 'Seven close friends reunite for the wedding of two of their friends. Problems arise because the bride and the maid of honor have had a long rivalry over the groom.',\n",
       "  'title': 'The Romantics',\n",
       "  'year': 2010,\n",
       "  'score': 0.7723205089569092},\n",
       " {'overview': 'Spoof of romantic comedies which focuses on a man (Campbell), his crush (Hannigan), his parents (Coolidge, Willard), and her father (Griffin).',\n",
       "  'title': 'Date Movie',\n",
       "  'year': 2006,\n",
       "  'score': 0.7711508274078369},\n",
       " {'overview': 'A film editor breaks up with his girlfriend, unsure if he is in love.',\n",
       "  'title': 'Modern Romance',\n",
       "  'year': 1981,\n",
       "  'score': 0.7707068920135498},\n",
       " {'overview': 'A pair of lovers have no idea that their affair is overseen and orchestrated by an angel sent by Cupid to see that they remain together in this romantic comedy.',\n",
       "  'title': \"Lover's Knot\",\n",
       "  'year': 1996,\n",
       "  'score': 0.7698203921318054}]"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_vector_search(\"romantic comedy movies\", topK=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb0052c7-79f0-453e-bab6-1f489db42a06",
   "metadata": {},
   "source": [
    "# Vector search query example with prefilter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "88d94702-b948-4560-9c0c-efa3f6c5101d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'overview': 'A film editor breaks up with his girlfriend, unsure if he is in love.',\n",
       "  'title': 'Modern Romance',\n",
       "  'year': 1981,\n",
       "  'score': 0.7707068920135498},\n",
       " {'overview': 'A 1938 romantic comedy film which tells the story of a teenage boy who becomes entangled with three different girls all at the same time.',\n",
       "  'title': 'Love Finds Andy Hardy',\n",
       "  'year': 1938,\n",
       "  'score': 0.7665838003158569},\n",
       " {'overview': 'Bickering husband-and-wife stage stars are manipulated into a break-up for publicity purposes.',\n",
       "  'title': 'Sweethearts',\n",
       "  'year': 1938,\n",
       "  'score': 0.7654489278793335},\n",
       " {'overview': 'Two couples go to a mutual friends wedding, and end up swapping partners.',\n",
       "  'title': 'Cousins',\n",
       "  'year': 1989,\n",
       "  'score': 0.7594757080078125},\n",
       " {'overview': 'A couple struggle to find happiness after a whirlwind courtship.',\n",
       "  'title': 'Made for Each Other',\n",
       "  'year': 1939,\n",
       "  'score': 0.7590867280960083}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Romantic comedy movies released before 1990\n",
    "query_vector_search(\"romantic comedy movies\", prefilter={\"year\":{\"$lt\": 1990}}, topK=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40189e33-4207-48b5-9b78-fd6bc90c46d3",
   "metadata": {},
   "source": [
    "# Vector search query example with prefilter and postfilter to control the semantic search relavancy and behaviour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "81b00224-3e68-47dc-8577-f60ccec73563",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'overview': 'A film editor breaks up with his girlfriend, unsure if he is in love.',\n",
       "  'title': 'Modern Romance',\n",
       "  'year': 1981,\n",
       "  'score': 0.7707068920135498},\n",
       " {'overview': 'A 1938 romantic comedy film which tells the story of a teenage boy who becomes entangled with three different girls all at the same time.',\n",
       "  'title': 'Love Finds Andy Hardy',\n",
       "  'year': 1938,\n",
       "  'score': 0.7665838003158569},\n",
       " {'overview': 'Bickering husband-and-wife stage stars are manipulated into a break-up for publicity purposes.',\n",
       "  'title': 'Sweethearts',\n",
       "  'year': 1938,\n",
       "  'score': 0.7654489278793335}]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_vector_search(\"romantic comedy movies\", prefilter={\"year\":{\"$lt\": 1990}}, postfilter={\"score\": {\"$gt\":0.76}},topK=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38375156-5261-4835-b2e2-98c89743a5a0",
   "metadata": {},
   "source": [
    "# Leverage MongoDB full text and Cohere Rerank Module to rank retrieved results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "1448ecd2-696d-41a4-a968-e3866c7188b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_fulltext_search(q,topK=25):\n",
    "    v = {\"$search\": {\n",
    "      \"text\": {\n",
    "        \"query\": q,\n",
    "        \"path\":\"overview\"\n",
    "      }\n",
    "    }}\n",
    "    project = {\"$project\": {\"score\": {\"$meta\": \"searchScore\"},\"_id\": 0,\"title\": 1, \"release-date\": 1, \"overview\": 1}}\n",
    "    docs = list(output_collection.aggregate([v,project, {\"$limit\":topK}]))\n",
    "    return docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "6883410d-0093-4df8-9864-64375abeb81c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'overview': 'A 1993 Bollywood Comedy starring Govinda and Chunky Pandey. This movie kick started a long line  of Comedy movies for Govinda.',\n",
       "  'title': 'Aankhen',\n",
       "  'score': 5.915504455566406},\n",
       " {'overview': \"Valentin Vaala's crazy, romantic comedy.\",\n",
       "  'title': 'Morsian yllättää',\n",
       "  'score': 5.859979629516602},\n",
       " {'overview': 'Romantic comedy about aspiring writers in NY.',\n",
       "  'title': 'Broadway Damage',\n",
       "  'score': 5.725950241088867},\n",
       " {'overview': 'Romantic comedy adapted from a Somerset Maugham play.',\n",
       "  'title': 'Too Many Husbands',\n",
       "  'score': 5.661208629608154},\n",
       " {'overview': 'A romantic comedy about the invention of the vibrator.',\n",
       "  'title': 'Hysteria',\n",
       "  'score': 5.597914695739746},\n",
       " {'overview': 'A romantic comedy set on Christmas Eve in Warsaw and centered around a series of characters.',\n",
       "  'title': 'Letters to Santa',\n",
       "  'score': 5.1916093826293945},\n",
       " {'overview': 'A romantic comedy about a neurotic girl who makes a vow of solitude after getting dumped.',\n",
       "  'title': 'Lovely Loneliness',\n",
       "  'score': 5.1916093826293945},\n",
       " {'overview': 'A \"romantic comedy\" based loosely on the suicide of the poet Henrich von Kleist in 1811.',\n",
       "  'title': 'Amour fou',\n",
       "  'score': 5.1916093826293945},\n",
       " {'overview': 'A mysterious stranger seduces three sisters in this sensual romantic comedy from director Jang Hyeon-Su.',\n",
       "  'title': 'Everybody Has Secrets',\n",
       "  'score': 5.1916093826293945},\n",
       " {'overview': \"A romantic comedy about a young girl who's coming to the kolkhoz to work as a cook.\",\n",
       "  'title': 'The Cook',\n",
       "  'score': 5.138331413269043}]"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# results before re ranking\n",
    "docs = query_fulltext_search(\"romantic comedy movies\", topK=10)\n",
    "docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "4f6dc536-8546-4b38-a302-562f747dc224",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document Rank: 1, Document Index: 22\n",
      "Document Title: Love Finds Andy Hardy\n",
      "Document: A 1938 romantic comedy film which tells the story of a teenage boy who becomes entangled with three different girls all at the same time.\n",
      "Relevance Score: 0.99\n",
      "\n",
      "\n",
      "Document Rank: 2, Document Index: 12\n",
      "Document Title: Seventh Heaven\n",
      "Document: Seventh Heaven or De zevende zemel is a 1993 Dutch romantic comedy film directed by Jean-Paul Lilienfeld.\n",
      "Relevance Score: 0.99\n",
      "\n",
      "\n",
      "Document Rank: 3, Document Index: 19\n",
      "Document Title: Shared Rooms\n",
      "Document: A new romantic comedy feature film that brings together three interrelated tales of gay men seeking family, love and sex during the holiday season.\n",
      "Relevance Score: 0.97\n",
      "\n",
      "\n",
      "Document Rank: 4, Document Index: 3\n",
      "Document Title: Too Many Husbands\n",
      "Document: Romantic comedy adapted from a Somerset Maugham play.\n",
      "Relevance Score: 0.97\n",
      "\n",
      "\n",
      "Document Rank: 5, Document Index: 20\n",
      "Document Title: Walking the Streets of Moscow\n",
      "Document: \"I Am Walking Along Moscow\" aka \"Ya Shagayu Po Moskve\" (1963) is a charming lyrical comedy directed by Georgi Daneliya in 1963 that was nominated for Golden Palm at Cannes Film Festival. Daneliya proved that it is possible to create a masterpiece in the most difficult genre of romantic comedy. Made by the team of young and incredibly talented artists that besides Daneliya included writer/poet Gennady Shpalikov, composer Andrei Petrov, and cinematographer Vadim Yusov (who had made four films with Andrei Tarkovski), and the dream cast of the talented actors even in the smaller cameos, \"I Am Walking Along Moscow\" keeps walking victoriously through the decades remaining deservingly one of the best and most beloved Russian comedies and simply one of the best Russian movies ever made. Funny and gentle, dreamy and humorous, romantic and realistic, the film is blessed with the eternal youth and will always take to the walk on the streets of Moscow new generations of the grateful viewers.\n",
      "Relevance Score: 0.96\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "q = \"romantic comedy movies\"\n",
    "docs = query_fulltext_search(q)\n",
    "results = co_client.rerank(query=q, documents=list(map(lambda x:x[\"overview\"], docs)), top_n=5, model='rerank-english-v2.0') # Change top_n to change the number of results returned. If top_n is not passed, all results will be returned.\n",
    "for idx, r in enumerate(results):\n",
    "    print(f\"Document Rank: {idx + 1}, Document Index: {r.index}\")\n",
    "    print(f\"Document Title: {docs[r.index]['title']}\")\n",
    "    print(f\"Document: {r.document['text']}\")\n",
    "    print(f\"Relevance Score: {r.relevance_score:.2f}\")\n",
    "    print(\"\\n\")"
   ]
  }
 ],
 "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.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}