openai-cookbook/examples/vector_databases/Using_vector_databases_for_embeddings_search.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb1537e6",
   "metadata": {},
   "source": [
    "# Using Vector Databases for Embeddings Search\n",
    "\n",
    "This notebook takes you through a simple flow to download some data, embed it, and then index and search it using a selection of vector databases. This is a common requirement for customers who want to store and search our embeddings with their own data in a secure environment to support production use cases such as chatbots, topic modelling and more.\n",
    "\n",
    "### What is a Vector Database\n",
    "\n",
    "A vector database is a database made to store, manage and search embedding vectors. The use of embeddings to encode unstructured data (text, audio, video and more) as vectors for consumption by machine-learning models has exploded in recent years, due to the increasing effectiveness of AI in solving use cases involving natural language, image recognition and other unstructured forms of data. Vector databases have emerged as an effective solution for enterprises to deliver and scale these use cases.\n",
    "\n",
    "### Why use a Vector Database\n",
    "\n",
    "Vector databases enable enterprises to take many of the embeddings use cases we've shared in this repo (question and answering, chatbot and recommendation services, for example), and make use of them in a secure, scalable environment. Many of our customers make embeddings solve their problems at small scale but performance and security hold them back from going into production - we see vector databases as a key component in solving that, and in this guide we'll walk through the basics of embedding text data, storing it in a vector database and using it for semantic search.\n",
    "\n",
    "\n",
    "### Demo Flow\n",
    "The demo flow is:\n",
    "- **Setup**: Import packages and set any required variables\n",
    "- **Load data**: Load a dataset and embed it using OpenAI embeddings\n",
    "- **Pinecone**\n",
    "    - *Setup*: Here we'll set up the Python client for Pinecone. For more details go [here](https://docs.pinecone.io/docs/quickstart)\n",
    "    - *Index Data*: We'll create an index with namespaces for __titles__ and __content__\n",
    "    - *Search Data*: We'll test out both namespaces with search queries to confirm it works\n",
    "- **Weaviate**\n",
    "    - *Setup*: Here we'll set up the Python client for Weaviate. For more details go [here](https://weaviate.io/developers/weaviate/current/client-libraries/python.html)\n",
    "    - *Index Data*: We'll create an index with __title__ search vectors in it\n",
    "    - *Search Data*: We'll run a few searches to confirm it works\n",
    "- **Milvus**\n",
    "    - *Setup*: Here we'll set up the Python client for Milvus. For more details go [here](https://milvus.io/docs)\n",
    "    - *Index Data* We'll create a collection and index it for both __titles__ and __content__\n",
    "    - *Search Data*: We'll test out both collections with search queries to confirm it works\n",
    "- **Qdrant**\n",
    "    - *Setup*: Here we'll set up the Python client for Qdrant. For more details go [here](https://github.com/qdrant/qdrant_client)\n",
    "    - *Index Data*: We'll create a collection with vectors for __titles__ and __content__\n",
    "    - *Search Data*: We'll run a few searches to confirm it works\n",
    "- **Redis**\n",
    "    - *Setup*: Set up the Redis-Py client. For more details go [here](https://github.com/redis/redis-py)\n",
    "    - *Index Data*: Create the search index for vector search and hybrid search (vector + full-text search) on all available fields.\n",
    "    - *Search Data*: Run a few example queries with various goals in mind.\n",
    "\n",
    "\n",
    "Once you've run through this notebook you should have a basic understanding of how to setup and use vector databases, and can move on to more complex use cases making use of our embeddings."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b59250",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Import the required libraries and set the embedding model that we'd like to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d8810f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll need to install the clients for all vector databases\n",
    "!pip install pinecone-client\n",
    "!pip install weaviate-client\n",
    "!pip install pymilvus\n",
    "!pip install qdrant-client\n",
    "!pip install redis\n",
    "\n",
    "#Install wget to pull zip file\n",
    "!pip install wget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5be94df6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "import tiktoken\n",
    "from typing import List, Iterator\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import wget\n",
    "from ast import literal_eval\n",
    "\n",
    "# Redis client library for Python\n",
    "import redis\n",
    "\n",
    "# Pinecone's client library for Python\n",
    "import pinecone\n",
    "\n",
    "# Weaviate's client library for Python\n",
    "import weaviate\n",
    "\n",
    "# Milvus's client library for Python\n",
    "import pymilvus\n",
    "\n",
    "# Qdrant's client library for Python\n",
    "import qdrant_client\n",
    "\n",
    "# I've set this to our new embeddings model, this can be changed to the embedding model of your choice\n",
    "EMBEDDING_MODEL = \"text-embedding-ada-002\"\n",
    "\n",
    "# Ignore unclosed SSL socket warnings - optional in case you get these errors\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(action=\"ignore\", message=\"unclosed\", category=ResourceWarning)\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5d9d2e1",
   "metadata": {},
   "source": [
    "## Load data\n",
    "\n",
    "In this section we'll load embedded data that we've prepared previous to this session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5dff8b55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'vector_database_wikipedia_articles_embedded.zip'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings_url = 'https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip'\n",
    "\n",
    "# The file is ~700 MB so this will take some time\n",
    "wget.download(embeddings_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21097972",
   "metadata": {},
   "outputs": [],
   "source": [
    "import zipfile\n",
    "with zipfile.ZipFile(\"vector_database_wikipedia_articles_embedded.zip\",\"r\") as zip_ref:\n",
    "    zip_ref.extractall(\"../data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70bbd8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "article_df = pd.read_csv('../data/vector_database_wikipedia_articles_embedded.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1721e45d",
   "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>id</th>\n",
       "      <th>url</th>\n",
       "      <th>title</th>\n",
       "      <th>text</th>\n",
       "      <th>title_vector</th>\n",
       "      <th>content_vector</th>\n",
       "      <th>vector_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>https://simple.wikipedia.org/wiki/April</td>\n",
       "      <td>April</td>\n",
       "      <td>April is the fourth month of the year in the J...</td>\n",
       "      <td>[0.001009464613161981, -0.020700545981526375, ...</td>\n",
       "      <td>[-0.011253940872848034, -0.013491976074874401,...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>https://simple.wikipedia.org/wiki/August</td>\n",
       "      <td>August</td>\n",
       "      <td>August (Aug.) is the eighth month of the year ...</td>\n",
       "      <td>[0.0009286514250561595, 0.000820168002974242, ...</td>\n",
       "      <td>[0.0003609954728744924, 0.007262262050062418, ...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>6</td>\n",
       "      <td>https://simple.wikipedia.org/wiki/Art</td>\n",
       "      <td>Art</td>\n",
       "      <td>Art is a creative activity that expresses imag...</td>\n",
       "      <td>[0.003393713850528002, 0.0061537534929811954, ...</td>\n",
       "      <td>[-0.004959689453244209, 0.015772193670272827, ...</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>8</td>\n",
       "      <td>https://simple.wikipedia.org/wiki/A</td>\n",
       "      <td>A</td>\n",
       "      <td>A or a is the first letter of the English alph...</td>\n",
       "      <td>[0.0153952119871974, -0.013759135268628597, 0....</td>\n",
       "      <td>[0.024894846603274345, -0.022186409682035446, ...</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>9</td>\n",
       "      <td>https://simple.wikipedia.org/wiki/Air</td>\n",
       "      <td>Air</td>\n",
       "      <td>Air refers to the Earth's atmosphere. Air is a...</td>\n",
       "      <td>[0.02224554680287838, -0.02044147066771984, -0...</td>\n",
       "      <td>[0.021524671465158463, 0.018522677943110466, -...</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id                                       url   title  \\\n",
       "0   1   https://simple.wikipedia.org/wiki/April   April   \n",
       "1   2  https://simple.wikipedia.org/wiki/August  August   \n",
       "2   6     https://simple.wikipedia.org/wiki/Art     Art   \n",
       "3   8       https://simple.wikipedia.org/wiki/A       A   \n",
       "4   9     https://simple.wikipedia.org/wiki/Air     Air   \n",
       "\n",
       "                                                text  \\\n",
       "0  April is the fourth month of the year in the J...   \n",
       "1  August (Aug.) is the eighth month of the year ...   \n",
       "2  Art is a creative activity that expresses imag...   \n",
       "3  A or a is the first letter of the English alph...   \n",
       "4  Air refers to the Earth's atmosphere. Air is a...   \n",
       "\n",
       "                                        title_vector  \\\n",
       "0  [0.001009464613161981, -0.020700545981526375, ...   \n",
       "1  [0.0009286514250561595, 0.000820168002974242, ...   \n",
       "2  [0.003393713850528002, 0.0061537534929811954, ...   \n",
       "3  [0.0153952119871974, -0.013759135268628597, 0....   \n",
       "4  [0.02224554680287838, -0.02044147066771984, -0...   \n",
       "\n",
       "                                      content_vector  vector_id  \n",
       "0  [-0.011253940872848034, -0.013491976074874401,...          0  \n",
       "1  [0.0003609954728744924, 0.007262262050062418, ...          1  \n",
       "2  [-0.004959689453244209, 0.015772193670272827, ...          2  \n",
       "3  [0.024894846603274345, -0.022186409682035446, ...          3  \n",
       "4  [0.021524671465158463, 0.018522677943110466, -...          4  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "article_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "960b82af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read vectors from strings back into a list\n",
    "article_df['title_vector'] = article_df.title_vector.apply(literal_eval)\n",
    "article_df['content_vector'] = article_df.content_vector.apply(literal_eval)\n",
    "\n",
    "# Set vector_id to be a string\n",
    "article_df['vector_id'] = article_df['vector_id'].apply(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a334ab8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 25000 entries, 0 to 24999\n",
      "Data columns (total 7 columns):\n",
      " #   Column          Non-Null Count  Dtype \n",
      "---  ------          --------------  ----- \n",
      " 0   id              25000 non-null  int64 \n",
      " 1   url             25000 non-null  object\n",
      " 2   title           25000 non-null  object\n",
      " 3   text            25000 non-null  object\n",
      " 4   title_vector    25000 non-null  object\n",
      " 5   content_vector  25000 non-null  object\n",
      " 6   vector_id       25000 non-null  object\n",
      "dtypes: int64(1), object(6)\n",
      "memory usage: 1.3+ MB\n"
     ]
    }
   ],
   "source": [
    "article_df.info(show_counts=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed32fc87",
   "metadata": {},
   "source": [
    "## Pinecone\n",
    "\n",
    "We'll index these embedded documents in a vector database and search them. The first option we'll look at is **Pinecone**, a managed vector database which offers a cloud-native option.\n",
    "\n",
    "Before you proceed with this step you'll need to navigate to [Pinecone](pinecone.io), sign up and then save your API key as an environment variable titled ```PINECONE_API_KEY```.\n",
    "\n",
    "For section we will:\n",
    "- Create an index with multiple namespaces for article titles and content\n",
    "- Store our data in the index with separate searchable \"namespaces\" for article **titles** and **content**\n",
    "- Fire some similarity search queries to verify our setup is working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92e6152a",
   "metadata": {},
   "outputs": [],
   "source": [
    "api_key = os.getenv(\"PINECONE_API_KEY\")\n",
    "pinecone.init(api_key=api_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b28543",
   "metadata": {},
   "source": [
    "### Create Index\n",
    "\n",
    "First we will need to create an index, which we'll call `wikipedia-articles`. Once we have an index, we can create multiple namespaces, which can make a single index searchable for various use cases. For more details, consult [Pinecone documentation](https://docs.pinecone.io/docs/namespaces#:~:text=Pinecone%20allows%20you%20to%20partition,different%20subsets%20of%20your%20index.).\n",
    "\n",
    "If you want to batch insert to your index in parallel to increase insertion speed then there is a great guide in the Pinecone documentation on [batch inserts in parallel](https://docs.pinecone.io/docs/insert-data#sending-upserts-in-parallel)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a71c575",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Models a simple batch generator that make chunks out of an input DataFrame\n",
    "class BatchGenerator:\n",
    "    \n",
    "    \n",
    "    def __init__(self, batch_size: int = 10) -> None:\n",
    "        self.batch_size = batch_size\n",
    "    \n",
    "    # Makes chunks out of an input DataFrame\n",
    "    def to_batches(self, df: pd.DataFrame) -> Iterator[pd.DataFrame]:\n",
    "        splits = self.splits_num(df.shape[0])\n",
    "        if splits <= 1:\n",
    "            yield df\n",
    "        else:\n",
    "            for chunk in np.array_split(df, splits):\n",
    "                yield chunk\n",
    "\n",
    "    # Determines how many chunks DataFrame contains\n",
    "    def splits_num(self, elements: int) -> int:\n",
    "        return round(elements / self.batch_size)\n",
    "    \n",
    "    __call__ = to_batches\n",
    "\n",
    "df_batcher = BatchGenerator(300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea9ad46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pick a name for the new index\n",
    "index_name = 'wikipedia-articles'\n",
    "\n",
    "# Check whether the index with the same name already exists - if so, delete it\n",
    "if index_name in pinecone.list_indexes():\n",
    "    pinecone.delete_index(index_name)\n",
    "    \n",
    "# Creates new index\n",
    "pinecone.create_index(name=index_name, dimension=len(article_df['content_vector'][0]))\n",
    "index = pinecone.Index(index_name=index_name)\n",
    "\n",
    "# Confirm our index was created\n",
    "pinecone.list_indexes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5daeba00",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upsert content vectors in content namespace - this can take a few minutes\n",
    "print(\"Uploading vectors to content namespace..\")\n",
    "for batch_df in df_batcher(article_df):\n",
    "    index.upsert(vectors=zip(batch_df.vector_id, batch_df.content_vector), namespace='content')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fc1b083",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upsert title vectors in title namespace - this can also take a few minutes\n",
    "print(\"Uploading vectors to title namespace..\")\n",
    "for batch_df in df_batcher(article_df):\n",
    "    index.upsert(vectors=zip(batch_df.vector_id, batch_df.title_vector), namespace='title')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90c7fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check index size for each namespace to confirm all of our docs have loaded\n",
    "index.describe_index_stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2da40a69",
   "metadata": {},
   "source": [
    "### Search data\n",
    "\n",
    "Now we'll enter some dummy searches and check we get decent results back"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8280363",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First we'll create dictionaries mapping vector IDs to their outputs so we can retrieve the text for our search results\n",
    "titles_mapped = dict(zip(article_df.vector_id,article_df.title))\n",
    "content_mapped = dict(zip(article_df.vector_id,article_df.text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c8c2aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_article(query, namespace, top_k=5):\n",
    "    '''Queries an article using its title in the specified\n",
    "     namespace and prints results.'''\n",
    "\n",
    "    # Create vector embeddings based on the title column\n",
    "    embedded_query = openai.Embedding.create(\n",
    "                                            input=query,\n",
    "                                            model=EMBEDDING_MODEL,\n",
    "                                            )[\"data\"][0]['embedding']\n",
    "\n",
    "    # Query namespace passed as parameter using title vector\n",
    "    query_result = index.query(embedded_query, \n",
    "                                      namespace=namespace, \n",
    "                                      top_k=top_k)\n",
    "\n",
    "    # Print query results \n",
    "    print(f'\\nMost similar results to {query} in \"{namespace}\" namespace:\\n')\n",
    "    if not query_result.matches:\n",
    "        print('no query result')\n",
    "    \n",
    "    matches = query_result.matches\n",
    "    ids = [res.id for res in matches]\n",
    "    scores = [res.score for res in matches]\n",
    "    df = pd.DataFrame({'id':ids, \n",
    "                       'score':scores,\n",
    "                       'title': [titles_mapped[_id] for _id in ids],\n",
    "                       'content': [content_mapped[_id] for _id in ids],\n",
    "                       })\n",
    "    \n",
    "    counter = 0\n",
    "    for k,v in df.iterrows():\n",
    "        counter += 1\n",
    "        print(f'{v.title} (score = {v.score})')\n",
    "    \n",
    "    print('\\n')\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3402b1b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_output = query_article('modern art in Europe','title')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64a3f90a",
   "metadata": {},
   "outputs": [],
   "source": [
    "content_query_output = query_article(\"Famous battles in Scottish history\",'content')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d939342f",
   "metadata": {},
   "source": [
    "## Weaviate\n",
    "\n",
    "Another vector database option we'll explore is **Weaviate**, which offers both a managed, [SaaS](https://console.weaviate.io/) option, as well as a self-hosted [open source](https://github.com/weaviate/weaviate) option. As we've already looked at a cloud vector database, we'll try the self-hosted option here.\n",
    "\n",
    "For this we will:\n",
    "- Set up a local deployment of Weaviate\n",
    "- Create indices in Weaviate\n",
    "- Store our data there\n",
    "- Fire some similarity search queries\n",
    "- Try a real use case\n",
    "\n",
    "\n",
    "### Bring your own vectors approach\n",
    "In this cookbook, we provide the data with already generated vectors. This is a good approach for scenarios, where your data is already vectorized.\n",
    "\n",
    "### Automated vectorization with OpenAI module\n",
    "For scenarios, where your data is not vectorized yet, you can delegate the vectorization task with OpenAI to Weaviate.\n",
    "Weaviate offers a built-in module [text2vec-openai](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/text2vec-openai), which takes care of the vectorization for you at:\n",
    "* import\n",
    "* for any CRUD operations\n",
    "* for semantic search\n",
    "\n",
    "Check out the [Getting Started with Weaviate and OpenAI module cookbook](./weaviate/getting-started-with-weaviate-and-openai.ipynb) to learn step by step how to import and vectorize data in one step."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bfdfe260",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "To run Weaviate locally, you'll need [Docker](https://www.docker.com/). Following the instructions contained in the Weaviate documentation [here](https://weaviate.io/developers/weaviate/installation/docker-compose), we created an example docker-compose.yml file in this repo saved at [./weaviate/docker-compose.yml](./weaviate/docker-compose.yml).\n",
    "\n",
    "After starting Docker, you can start Weaviate locally by navigating to the `examples/vector_databases/weaviate/` directory and running `docker-compose up -d`.\n",
    "\n",
    "#### SaaS\n",
    "Alternatively you can use [Weaviate Cloud Service](https://console.weaviate.io/) (WCS) to create a free Weaviate cluster.\n",
    "1. create a free account and/or login to [WCS](https://console.weaviate.io/)\n",
    "2. create a `Weaviate Cluster` with the following settings:\n",
    "    * Sandbox: `Sandbox Free`\n",
    "    * Weaviate Version: Use default (latest)\n",
    "    * OIDC Authentication: `Disabled`\n",
    "3. your instance should be ready in a minute or two\n",
    "4. make a note of the `Cluster Id`. The link will take you to the full path of your cluster (you will need it later to connect to it). It should be something like: `https://your-project-name-suffix.weaviate.network` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a78f95d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option #1 - Self-hosted - Weaviate Open Source \n",
    "client = weaviate.Client(\n",
    "    url=\"http://localhost:8080\",\n",
    "    additional_headers={\n",
    "        \"X-OpenAI-Api-Key\": os.getenv(\"OPENAI_API_KEY\")\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e00b7d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Option #2 - SaaS - (Weaviate Cloud Service)\n",
    "client = weaviate.Client(\n",
    "    url=\"https://your-wcs-instance-name.weaviate.network\",\n",
    "    additional_headers={\n",
    "        \"X-OpenAI-Api-Key\": os.getenv(\"OPENAI_API_KEY\")\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d370afa",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.is_ready()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "03a926b9",
   "metadata": {},
   "source": [
    "### Index data\n",
    "\n",
    "In Weaviate you create __schemas__ to capture each of the entities you will be searching. \n",
    "\n",
    "In this case we'll create a schema called **Article** with the **title** vector from above included for us to search by.\n",
    "\n",
    "The next few steps closely follow the documentation Weaviate provides [here](https://weaviate.io/developers/weaviate/quickstart).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e6175a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear up the schema, so that we can recreate it\n",
    "client.schema.delete_all()\n",
    "client.schema.get()\n",
    "\n",
    "# Define the Schema object to use `text-embedding-ada-002` on `title` and `content`, but skip it for `url`\n",
    "article_schema = {\n",
    "    \"class\": \"Article\",\n",
    "    \"description\": \"A collection of articles\",\n",
    "    \"vectorizer\": \"text2vec-openai\",\n",
    "    \"moduleConfig\": {\n",
    "        \"text2vec-openai\": {\n",
    "          \"model\": \"ada\",\n",
    "          \"modelVersion\": \"002\",\n",
    "          \"type\": \"text\"\n",
    "        }\n",
    "    },\n",
    "    \"properties\": [{\n",
    "        \"name\": \"title\",\n",
    "        \"description\": \"Title of the article\",\n",
    "        \"dataType\": [\"string\"]\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"content\",\n",
    "        \"description\": \"Contents of the article\",\n",
    "        \"dataType\": [\"text\"],\n",
    "        \"moduleConfig\": { \"text2vec-openai\": { \"skip\": True } }\n",
    "    }]\n",
    "}\n",
    "\n",
    "# add the Article schema\n",
    "client.schema.create_class(article_schema)\n",
    "\n",
    "# get the schema to make sure it worked\n",
    "client.schema.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea838e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Step 1 - configure Weaviate Batch, which optimizes CRUD operations in bulk\n",
    "# - starting batch size of 100\n",
    "# - dynamically increase/decrease based on performance\n",
    "# - add timeout retries if something goes wrong\n",
    "\n",
    "client.batch.configure(\n",
    "    batch_size=100,\n",
    "    dynamic=True,\n",
    "    timeout_retries=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c967ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Step 2 - import data\n",
    "\n",
    "print(\"Uploading data with vectors to Article schema..\")\n",
    "\n",
    "counter=0\n",
    "\n",
    "with client.batch as batch:\n",
    "    for k,v in article_df.iterrows():\n",
    "        \n",
    "        # print update message every 100 objects        \n",
    "        if (counter %100 == 0):\n",
    "            print(f\"Import {counter} / {len(article_df)} \")\n",
    "        \n",
    "        properties = {\n",
    "            \"title\": v[\"title\"],\n",
    "            \"content\": v[\"text\"]\n",
    "        }\n",
    "        \n",
    "        vector = v[\"title_vector\"]\n",
    "        \n",
    "        batch.add_data_object(properties, \"Article\", None, vector)\n",
    "        counter = counter+1\n",
    "\n",
    "print(f\"Importing ({len(article_df)}) Articles complete\")  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f826e1ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test that all data has loaded – get object count\n",
    "result = (\n",
    "    client.query.aggregate(\"Article\")\n",
    "    .with_fields(\"meta { count }\")\n",
    "    .do()\n",
    ")\n",
    "print(\"Object count: \", result[\"data\"][\"Aggregate\"][\"Article\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c09d483",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test one article has worked by checking one object\n",
    "test_article = (\n",
    "    client.query\n",
    "    .get(\"Article\", [\"title\", \"content\", \"_additional {id}\"])\n",
    "    .with_limit(1)\n",
    "    .do()\n",
    ")[\"data\"][\"Get\"][\"Article\"][0]\n",
    "\n",
    "print(test_article[\"_additional\"][\"id\"])\n",
    "print(test_article[\"title\"])\n",
    "print(test_article[\"content\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46050ca9",
   "metadata": {},
   "source": [
    "### Search data\n",
    "\n",
    "As above, we'll fire some queries at our new Index and get back results based on the closeness to our existing vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "add222d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_weaviate(query, collection_name, top_k=20):\n",
    "\n",
    "    # Creates embedding vector from user query\n",
    "    embedded_query = openai.Embedding.create(\n",
    "        input=query,\n",
    "        model=EMBEDDING_MODEL,\n",
    "    )[\"data\"][0]['embedding']\n",
    "    \n",
    "    near_vector = {\"vector\": embedded_query}\n",
    "\n",
    "    # Queries input schema with vectorised user query\n",
    "    query_result = (\n",
    "        client.query\n",
    "        .get(collection_name, [\"title\", \"content\", \"_additional {certainty distance}\"])\n",
    "        .with_near_vector(near_vector)\n",
    "        .with_limit(top_k)\n",
    "        .do()\n",
    "    )\n",
    "    \n",
    "    return query_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c888aa4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_result = query_weaviate(\"modern art in Europe\", \"Article\")\n",
    "counter = 0\n",
    "for article in query_result[\"data\"][\"Get\"][\"Article\"]:\n",
    "    counter += 1\n",
    "    print(f\"{counter}. { article['title']} (Certainty: {round(article['_additional']['certainty'],3) }) (Distance: {round(article['_additional']['distance'],3) })\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54cd8e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_result = query_weaviate(\"Famous battles in Scottish history\", \"Article\")\n",
    "counter = 0\n",
    "for article in query_result[\"data\"][\"Get\"][\"Article\"]:\n",
    "    counter += 1\n",
    "    print(f\"{counter}. {article['title']} (Score: {round(article['_additional']['certainty'],3) })\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "220b3e11",
   "metadata": {},
   "source": [
    "### Let Weaviate handle vector embeddings\n",
    "\n",
    "Weaviate has a [built-in module for OpenAI](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/text2vec-openai), which takes care of the steps required to generate a vector embedding for your queries and any CRUD operations.\n",
    "\n",
    "This allows you to run a vector query with the `with_near_text` filter, which uses your `OPEN_API_KEY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9425c882",
   "metadata": {},
   "outputs": [],
   "source": [
    "def near_text_weaviate(query, collection_name):\n",
    "    \n",
    "    nearText = {\n",
    "        \"concepts\": [query],\n",
    "        \"distance\": 0.7,\n",
    "    }\n",
    "\n",
    "    properties = [\n",
    "        \"title\", \"content\",\n",
    "        \"_additional {certainty distance}\"\n",
    "    ]\n",
    "\n",
    "    query_result = (\n",
    "        client.query\n",
    "        .get(collection_name, properties)\n",
    "        .with_near_text(nearText)\n",
    "        .with_limit(20)\n",
    "        .do()\n",
    "    )[\"data\"][\"Get\"][collection_name]\n",
    "    \n",
    "    print (f\"Objects returned: {len(query_result)}\")\n",
    "    \n",
    "    return query_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "501a16f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_result = near_text_weaviate(\"modern art in Europe\",\"Article\")\n",
    "counter = 0\n",
    "for article in query_result:\n",
    "    counter += 1\n",
    "    print(f\"{counter}. { article['title']} (Certainty: {round(article['_additional']['certainty'],3) }) (Distance: {round(article['_additional']['distance'],3) })\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839b26df",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_result = near_text_weaviate(\"Famous battles in Scottish history\",\"Article\")\n",
    "counter = 0\n",
    "for article in query_result:\n",
    "    counter += 1\n",
    "    print(f\"{counter}. { article['title']} (Certainty: {round(article['_additional']['certainty'],3) }) (Distance: {round(article['_additional']['distance'],3) })\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4dc3a0c0",
   "metadata": {},
   "source": [
    "## Milvus\n",
    "\n",
    "The next vector database we will take a look at is **Milvus**, which also offers a SaaS option like the previous two, as well as self-hosted options using either helm or docker-compose. Sticking to the idea of open source, we will show our self-hosted example here.\n",
    "\n",
    "In this example we will:\n",
    "- Set up a local docker-compose based deployment\n",
    "- Create the title and content collections\n",
    "- Store our data\n",
    "- Test out our system with real world searches"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "fe4914e9",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "There are many ways to run Milvus (take a look [here](https://milvus.io/docs/install_cluster-milvusoperator.md)), but for now we will stick to a simple standalone Milvus instance with docker-compose.\n",
    "\n",
    "A simple docker-file can be found at `./milvus/docker-compose.yaml` and can be run using `docker-compose up` if within that mentioned directory or using `docker-compose -f path/to/file up`\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e10f2ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import connections\n",
    "\n",
    "connections.connect(host='localhost', port=19530)  # Local instance defaults to port 19530"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "64ffed22",
   "metadata": {},
   "source": [
    "### Index data\n",
    "\n",
    "In Milvus data is stored in the form of collections, with each collection being able to store the vectors and any attributes that come with them.\n",
    "\n",
    "In this case we'll create a collection called **articles** which contains the url, title, text and the content_embedding.\n",
    "\n",
    "In addition to this we will also create an index on the content embedding. Milvus allows for the use of many SOTA indexing methods, but in this case, we are going to use HNSW.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfabc3db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import utility, Collection, FieldSchema, CollectionSchema, DataType\n",
    "\n",
    "# Remove the collection if it already exists.\n",
    "if utility.has_collection('articles'):\n",
    "    utility.drop_collection('articles')\n",
    "\n",
    "fields = [\n",
    "    FieldSchema(name='id', dtype=DataType.INT64),\n",
    "    FieldSchema(name='url', dtype=DataType.VARCHAR, max_length=1000),  # Strings have to specify a max length [1, 65535]\n",
    "    FieldSchema(name='title', dtype=DataType.VARCHAR, max_length=1000),\n",
    "    FieldSchema(name='text', dtype=DataType.VARCHAR, max_length=1000),\n",
    "    FieldSchema(name='content_vector', dtype=DataType.FLOAT_VECTOR, dim=len(article_df['content_vector'][0])),\n",
    "    FieldSchema(name='vector_id', dtype=DataType.INT64, is_primary=True, auto_id=False),\n",
    "]\n",
    "\n",
    "col_schema = CollectionSchema(fields)\n",
    "\n",
    "col = Collection('articles', col_schema)\n",
    "\n",
    "# Using a basic HNSW index for this example\n",
    "index = {\n",
    "    'index_type': 'HNSW',\n",
    "    'metric_type': 'L2',\n",
    "    'params': {\n",
    "        'M': 8,\n",
    "        'efConstruction': 64\n",
    "    },\n",
    "}\n",
    "\n",
    "col.create_index('content_vector', index)\n",
    "col.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51aaed1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the above provided batching function from Pinecone\n",
    "def to_batches(df: pd.DataFrame, batch_size: int) -> Iterator[pd.DataFrame]:\n",
    "    splits = df.shape[0] / batch_size\n",
    "    if splits <= 1:\n",
    "        yield df\n",
    "    else:\n",
    "        for chunk in np.array_split(df, splits):\n",
    "            yield chunk\n",
    "\n",
    "# Since we are storing the text within Milvus we need to clip any that are over our set limit.\n",
    "# We can also set the limit to be higher, but that slows down the search requests as more info \n",
    "# needs to be sent back.\n",
    "def shorten_text(text):\n",
    "    if len(text) >= 996:\n",
    "        return text[:996] + '...'\n",
    "    else:\n",
    "        return text\n",
    "\n",
    "for batch in to_batches(article_df, 1000):\n",
    "    batch = batch.drop(columns = ['title_vector'])\n",
    "    batch['text'] = batch.text.apply(shorten_text)\n",
    "    # Due to the vector_id being converted to a string for compatiblity for other vector dbs,\n",
    "    # we want to swap it back to its original form.\n",
    "    batch['vector_id'] = batch.vector_id.apply(int)\n",
    "    col.insert(batch) \n",
    "\n",
    "col.flush()   "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1f68a790",
   "metadata": {},
   "source": [
    "# Search\n",
    "Once the data is inserted into Milvus we can perform searches. For this example the search function takes one argument, top_k, how many closest matches to return. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02f21251",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_article(query, top_k=5):\n",
    "    # Generate the embedding with openai\n",
    "    embedded_query = openai.Embedding.create(\n",
    "        input=query,\n",
    "        model=EMBEDDING_MODEL,\n",
    "    )[\"data\"][0]['embedding']\n",
    "\n",
    "    # Using some basic params for HNSW\n",
    "    search_param = {\n",
    "        'ef': max(64, top_k)\n",
    "    }\n",
    "\n",
    "    # Perform the search.\n",
    "    res = col.search([embedded_query], 'content_vector', search_param, output_fields = ['title', 'url'], limit = top_k)\n",
    "\n",
    "    ret = []\n",
    "    for hit in res[0]:\n",
    "        # Get the id, distance, and title for the results\n",
    "        ret.append({'vector_id': hit.id, 'distance': hit.score, 'title': hit.entity.get('title'), 'url': hit.entity.get('url')})\n",
    "    return ret\n",
    "    \n",
    "\n",
    "for x in query_article('fastest plane ever made', 3):\n",
    "    print(x.items())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cfaed9d",
   "metadata": {},
   "source": [
    "## Qdrant\n",
    "\n",
    "The last vector database we'll consider is **[Qdrant](https://qdrant.tech/)**. This is a high-performant vector search database written in Rust. It offers both on-premise and cloud version, but for the purposes of that example we're going to use the local deployment mode.\n",
    "\n",
    "Setting everything up will require:\n",
    "- Spinning up a local instance of Qdrant\n",
    "- Configuring the collection and storing the data in it\n",
    "- Trying out with some queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38774565",
   "metadata": {},
   "source": [
    "### Setup\n",
    "\n",
    "For the local deployment, we are going to use Docker, according to the Qdrant documentation: https://qdrant.tech/documentation/quick_start/. Qdrant requires just a single container, but an example of the docker-compose.yaml file is available at `./qdrant/docker-compose.yaml` in this repo.\n",
    "\n",
    "You can start Qdrant instance locally by navigating to this directory and running `docker-compose up -d `"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d697e9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:28:38.928205Z",
     "start_time": "2023-01-18T09:28:38.913987Z"
    }
   },
   "outputs": [],
   "source": [
    "qdrant = qdrant_client.QdrantClient(host='localhost', prefer_grpc=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1deeb539",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:29:19.806639Z",
     "start_time": "2023-01-18T09:29:19.727897Z"
    }
   },
   "outputs": [],
   "source": [
    "qdrant.get_collections()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc006b6f",
   "metadata": {},
   "source": [
    "### Index data\n",
    "\n",
    "Qdrant stores data in __collections__ where each object is described by at least one vector and may contain an additional metadata called __payload__. Our collection will be called **Articles** and each object will be described by both **title** and **content** vectors.\n",
    "\n",
    "We'll be using an official [qdrant-client](https://github.com/qdrant/qdrant_client) package that has all the utility methods already built-in."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a84ee1d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:29:22.530121Z",
     "start_time": "2023-01-18T09:29:22.524604Z"
    }
   },
   "outputs": [],
   "source": [
    "from qdrant_client.http import models as rest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00876f92",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:31:14.413334Z",
     "start_time": "2023-01-18T09:31:13.619079Z"
    }
   },
   "outputs": [],
   "source": [
    "vector_size = len(article_df['content_vector'][0])\n",
    "\n",
    "qdrant.recreate_collection(\n",
    "    collection_name='Articles',\n",
    "    vectors_config={\n",
    "        'title': rest.VectorParams(\n",
    "            distance=rest.Distance.COSINE,\n",
    "            size=vector_size,\n",
    "        ),\n",
    "        'content': rest.VectorParams(\n",
    "            distance=rest.Distance.COSINE,\n",
    "            size=vector_size,\n",
    "        ),\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f24e76ab",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:36:28.597535Z",
     "start_time": "2023-01-18T09:36:24.108867Z"
    }
   },
   "outputs": [],
   "source": [
    "qdrant.upsert(\n",
    "    collection_name='Articles',\n",
    "    points=[\n",
    "        rest.PointStruct(\n",
    "            id=k,\n",
    "            vector={\n",
    "                'title': v['title_vector'],\n",
    "                'content': v['content_vector'],\n",
    "            },\n",
    "            payload=v.to_dict(),\n",
    "        )\n",
    "        for k, v in article_df.iterrows()\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1188a12",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:58:13.825886Z",
     "start_time": "2023-01-18T09:58:13.816248Z"
    }
   },
   "outputs": [],
   "source": [
    "# Check the collection size to make sure all the points have been stored\n",
    "qdrant.count(collection_name='Articles')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06ed119b",
   "metadata": {},
   "source": [
    "### Search Data\n",
    "\n",
    "Once the data is put into Qdrant we will start querying the collection for the closest vectors. We may provide an additional parameter `vector_name` to switch from title to content based search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1bac4ef",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:50:35.265647Z",
     "start_time": "2023-01-18T09:50:35.256065Z"
    }
   },
   "outputs": [],
   "source": [
    "def query_qdrant(query, collection_name, vector_name='title', top_k=20):\n",
    "\n",
    "    # Creates embedding vector from user query\n",
    "    embedded_query = openai.Embedding.create(\n",
    "        input=query,\n",
    "        model=EMBEDDING_MODEL,\n",
    "    )['data'][0]['embedding']\n",
    "    \n",
    "    query_results = qdrant.search(\n",
    "        collection_name=collection_name,\n",
    "        query_vector=(\n",
    "            vector_name, embedded_query\n",
    "        ),\n",
    "        limit=top_k,\n",
    "    )\n",
    "    \n",
    "    return query_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa92f3d3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:50:46.545145Z",
     "start_time": "2023-01-18T09:50:35.711020Z"
    }
   },
   "outputs": [],
   "source": [
    "query_results = query_qdrant('modern art in Europe', 'Articles')\n",
    "for i, article in enumerate(query_results):\n",
    "    print(f'{i + 1}. {article.payload[\"title\"]} (Score: {round(article.score, 3)})')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed116b8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-01-18T09:53:11.038910Z",
     "start_time": "2023-01-18T09:52:55.248029Z"
    }
   },
   "outputs": [],
   "source": [
    "# This time we'll query using content vector\n",
    "query_results = query_qdrant('Famous battles in Scottish history', 'Articles', 'content')\n",
    "for i, article in enumerate(query_results):\n",
    "    print(f'{i + 1}. {article.payload[\"title\"]} (Score: {round(article.score, 3)})')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "43bffd04",
   "metadata": {},
   "source": [
    "# Redis\n",
    "\n",
    "The last vector database covered in this tutorial is **[Redis](https://redis.io)**. You most likely already know Redis. What you might not be aware of is the [RediSearch module](https://github.com/RediSearch/RediSearch). Enterprises have been using Redis with the RediSearch module for years now across all major cloud providers, Redis Cloud, and on premise. Recently, the Redis team added vector storage and search capability to this module in addition to the features RediSearch already had.\n",
    "\n",
    "Given the large ecosystem around Redis, there are most likely client libraries in the language you need. You can use any standard Redis client library to run RediSearch commands, but it's easiest to use a library that wraps the RediSearch API. Below are a few examples, but you can find more client libraries [here](https://redis.io/resources/clients/).\n",
    "\n",
    "| Project | Language | License | Author | Stars |\n",
    "|----------|---------|--------|---------|-------|\n",
    "| [jedis][jedis-url] | Java | MIT | [Redis][redis-url] | ![Stars][jedis-stars] |\n",
    "| [redis-py][redis-py-url] | Python | MIT | [Redis][redis-url] | ![Stars][redis-py-stars] |\n",
    "| [node-redis][node-redis-url] | Node.js | MIT | [Redis][redis-url] | ![Stars][node-redis-stars] |\n",
    "| [nredisstack][nredisstack-url] | .NET | MIT | [Redis][redis-url] | ![Stars][nredisstack-stars] |\n",
    "| [redisearch-go][redisearch-go-url] | Go | BSD | [Redis][redisearch-go-author] | [![redisearch-go-stars]][redisearch-go-url] |\n",
    "| [redisearch-api-rs][redisearch-api-rs-url] | Rust | BSD | [Redis][redisearch-api-rs-author] | [![redisearch-api-rs-stars]][redisearch-api-rs-url] |\n",
    "\n",
    "[redis-url]: https://redis.com\n",
    "\n",
    "[redis-py-url]: https://github.com/redis/redis-py\n",
    "[redis-py-stars]: https://img.shields.io/github/stars/redis/redis-py.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "[redis-py-package]: https://pypi.python.org/pypi/redis\n",
    "\n",
    "[jedis-url]: https://github.com/redis/jedis\n",
    "[jedis-stars]: https://img.shields.io/github/stars/redis/jedis.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "[Jedis-package]: https://search.maven.org/artifact/redis.clients/jedis\n",
    "\n",
    "[nredisstack-url]: https://github.com/redis/nredisstack\n",
    "[nredisstack-stars]: https://img.shields.io/github/stars/redis/nredisstack.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "[nredisstack-package]: https://www.nuget.org/packages/nredisstack/\n",
    "\n",
    "[node-redis-url]: https://github.com/redis/node-redis\n",
    "[node-redis-stars]: https://img.shields.io/github/stars/redis/node-redis.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "[node-redis-package]: https://www.npmjs.com/package/redis\n",
    "\n",
    "[redis-om-python-url]: https://github.com/redis/redis-om-python\n",
    "[redis-om-python-author]: https://redis.com\n",
    "[redis-om-python-stars]: https://img.shields.io/github/stars/redis/redis-om-python.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "\n",
    "[redisearch-go-url]: https://github.com/RediSearch/redisearch-go\n",
    "[redisearch-go-author]: https://redis.com\n",
    "[redisearch-go-stars]: https://img.shields.io/github/stars/RediSearch/redisearch-go.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "\n",
    "[redisearch-api-rs-url]: https://github.com/RediSearch/redisearch-api-rs\n",
    "[redisearch-api-rs-author]: https://redis.com\n",
    "[redisearch-api-rs-stars]: https://img.shields.io/github/stars/RediSearch/redisearch-api-rs.svg?style=social&amp;label=Star&amp;maxAge=2592000\n",
    "\n",
    "\n",
    "In the below cells, we will walk you through using Redis as a vector database. Since many of you are likely already used to the Redis API, this should be familiar to most."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "698e24f6",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "There are many ways to deploy Redis with RediSearch. The easiest way to get started is to use Docker, but there are are many potential options for deployment. For other deployment options, see the [redis directory](./redis) in this repo.\n",
    "\n",
    "For this tutorial, we will use Redis Stack on Docker.\n",
    "\n",
    "Start a version of Redis with RediSearch (Redis Stack) by running the following docker command\n",
    "\n",
    "```bash\n",
    "$ cd redis\n",
    "$ docker compose up -d\n",
    "```\n",
    "This also includes the [RedisInsight](https://redis.com/redis-enterprise/redis-insight/) GUI for managing your Redis database which you can view at [http://localhost:8001](http://localhost:8001) once you start the docker container.\n",
    "\n",
    "You're all set up and ready to go! Next, we import and create our client for communicating with the Redis database we just created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "d2ce669a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 134,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import redis\n",
    "from redis.commands.search.indexDefinition import (\n",
    "    IndexDefinition,\n",
    "    IndexType\n",
    ")\n",
    "from redis.commands.search.query import Query\n",
    "from redis.commands.search.field import (\n",
    "    TextField,\n",
    "    VectorField\n",
    ")\n",
    "\n",
    "REDIS_HOST =  \"localhost\"\n",
    "REDIS_PORT = 6379\n",
    "REDIS_PASSWORD = \"\" # default for passwordless Redis\n",
    "\n",
    "# Connect to Redis\n",
    "redis_client = redis.Redis(\n",
    "    host=REDIS_HOST,\n",
    "    port=REDIS_PORT,\n",
    "    password=REDIS_PASSWORD\n",
    ")\n",
    "redis_client.ping()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3f6f0af9",
   "metadata": {},
   "source": [
    "## Creating a Search Index\n",
    "\n",
    "The below cells will show how to specify and create a search index in Redis. We will\n",
    "\n",
    "1. Set some constants for defining our index like the distance metric and the index name\n",
    "2. Define the index schema with RediSearch fields\n",
    "3. Create the index\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "a7c64cb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "VECTOR_DIM = len(article_df['title_vector'][0]) # length of the vectors\n",
    "VECTOR_NUMBER = len(article_df)                 # initial number of vectors\n",
    "INDEX_NAME = \"embeddings-index\"                 # name of the search index\n",
    "PREFIX = \"doc\"                                  # prefix for the document keys\n",
    "DISTANCE_METRIC = \"COSINE\"                      # distance metric for the vectors (ex. COSINE, IP, L2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "d95fcd06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define RediSearch fields for each of the columns in the dataset\n",
    "title = TextField(name=\"title\")\n",
    "url = TextField(name=\"url\")\n",
    "text = TextField(name=\"text\")\n",
    "title_embedding = VectorField(\"title_vector\",\n",
    "    \"FLAT\", {\n",
    "        \"TYPE\": \"FLOAT32\",\n",
    "        \"DIM\": VECTOR_DIM,\n",
    "        \"DISTANCE_METRIC\": DISTANCE_METRIC,\n",
    "        \"INITIAL_CAP\": VECTOR_NUMBER,\n",
    "    }\n",
    ")\n",
    "text_embedding = VectorField(\"content_vector\",\n",
    "    \"FLAT\", {\n",
    "        \"TYPE\": \"FLOAT32\",\n",
    "        \"DIM\": VECTOR_DIM,\n",
    "        \"DISTANCE_METRIC\": DISTANCE_METRIC,\n",
    "        \"INITIAL_CAP\": VECTOR_NUMBER,\n",
    "    }\n",
    ")\n",
    "fields = [title, url, text, title_embedding, text_embedding]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "7418480d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if index exists\n",
    "try:\n",
    "    redis_client.ft(INDEX_NAME).info()\n",
    "    print(\"Index already exists\")\n",
    "except:\n",
    "    # Create RediSearch Index\n",
    "    redis_client.ft(INDEX_NAME).create_index(\n",
    "        fields = fields,\n",
    "        definition = IndexDefinition(prefix=[PREFIX], index_type=IndexType.HASH)\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f3563eec",
   "metadata": {},
   "source": [
    "## Load Documents into the Index\n",
    "\n",
    "Now that we have a search index, we can load documents into it. We will use the same documents we used in the previous examples. In Redis, either the Hash or JSON (if using RedisJSON in addition to RediSearch) data types can be used to store documents. We will use the HASH data type in this example. The below cells will show how to load documents into the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "e98d63ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def index_documents(client: redis.Redis, prefix: str, documents: pd.DataFrame):\n",
    "    records = documents.to_dict(\"records\")\n",
    "    for doc in records:\n",
    "        key = f\"{prefix}:{str(doc['id'])}\"\n",
    "\n",
    "        # create byte vectors for title and content\n",
    "        title_embedding = np.array(doc[\"title_vector\"], dtype=np.float32).tobytes()\n",
    "        content_embedding = np.array(doc[\"content_vector\"], dtype=np.float32).tobytes()\n",
    "\n",
    "        # replace list of floats with byte vectors\n",
    "        doc[\"title_vector\"] = title_embedding\n",
    "        doc[\"content_vector\"] = content_embedding\n",
    "\n",
    "        client.hset(key, mapping = doc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "098d3c5a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 25000 documents in Redis search index with name: embeddings-index\n"
     ]
    }
   ],
   "source": [
    "index_documents(redis_client, PREFIX, article_df)\n",
    "print(f\"Loaded {redis_client.info()['db0']['keys']} documents in Redis search index with name: {INDEX_NAME}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f646bff4",
   "metadata": {},
   "source": [
    "## Running Search Queries\n",
    "\n",
    "Now that we have a search index and documents loaded into it, we can run search queries. Below we will provide a function that will run a search query and return the results. Using this function we run a few queries that will show how you can utilize Redis as a vector database. Each example will demonstrate specific features to keep in mind when developing your search application with Redis.\n",
    "\n",
    "1. **Return Fields**: You can specify which fields you want to return in the search results. This is useful if you only want to return a subset of the fields in your documents and doesn't require a separate call to retrieve documents. In the below example, we will only return the `title` field in the search results.\n",
    "2. **Hybrid Search**: You can combine vector search with any of the other RediSearch fields for hybrid search such as full text search, tag, geo, and numeric. In the below example, we will combine vector search with full text search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "id": "508d1f89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_redis(\n",
    "    redis_client: redis.Redis,\n",
    "    user_query: str,\n",
    "    index_name: str = \"embeddings-index\",\n",
    "    vector_field: str = \"title_vector\",\n",
    "    return_fields: list = [\"title\", \"url\", \"text\", \"vector_score\"],\n",
    "    hybrid_fields = \"*\",\n",
    "    k: int = 20,\n",
    ") -> List[dict]:\n",
    "\n",
    "    # Creates embedding vector from user query\n",
    "    embedded_query = openai.Embedding.create(input=user_query,\n",
    "                                            model=EMBEDDING_MODEL,\n",
    "                                            )[\"data\"][0]['embedding']\n",
    "\n",
    "    # Prepare the Query\n",
    "    base_query = f'{hybrid_fields}=>[KNN {k} @{vector_field} $vector AS vector_score]'\n",
    "    query = (\n",
    "        Query(base_query)\n",
    "         .return_fields(*return_fields)\n",
    "         .sort_by(\"vector_score\")\n",
    "         .paging(0, k)\n",
    "         .dialect(2)\n",
    "    )\n",
    "    params_dict = {\"vector\": np.array(embedded_query).astype(dtype=np.float32).tobytes()}\n",
    "\n",
    "    # perform vector search\n",
    "    results = redis_client.ft(index_name).search(query, params_dict)\n",
    "    for i, article in enumerate(results.docs):\n",
    "        score = 1 - float(article.vector_score)\n",
    "        print(f\"{i}. {article.title} (Score: {round(score ,3) })\")\n",
    "    return results.docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "1f0eef07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. Museum of Modern Art (Score: 0.875)\n",
      "1. Western Europe (Score: 0.867)\n",
      "2. Renaissance art (Score: 0.864)\n",
      "3. Pop art (Score: 0.86)\n",
      "4. Northern Europe (Score: 0.855)\n",
      "5. Hellenistic art (Score: 0.853)\n",
      "6. Modernist literature (Score: 0.847)\n",
      "7. Art film (Score: 0.843)\n",
      "8. Central Europe (Score: 0.843)\n",
      "9. European (Score: 0.841)\n"
     ]
    }
   ],
   "source": [
    "# For using OpenAI to generate query embedding\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\", \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\")\n",
    "results = search_redis(redis_client, 'modern art in Europe', k=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "id": "7b805a81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. Battle of Bannockburn (Score: 0.869)\n",
      "1. Wars of Scottish Independence (Score: 0.861)\n",
      "2. 1651 (Score: 0.853)\n",
      "3. First War of Scottish Independence (Score: 0.85)\n",
      "4. Robert I of Scotland (Score: 0.846)\n",
      "5. 841 (Score: 0.844)\n",
      "6. 1716 (Score: 0.844)\n",
      "7. 1314 (Score: 0.837)\n",
      "8. 1263 (Score: 0.836)\n",
      "9. William Wallace (Score: 0.835)\n"
     ]
    }
   ],
   "source": [
    "results = search_redis(redis_client, 'Famous battles in Scottish history', vector_field='content_vector', k=10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0ed0b34e",
   "metadata": {},
   "source": [
    "## Hybrid Queries with Redis\n",
    "\n",
    "The previous examples showed how run vector search queries with RediSearch. In this section, we will show how to combine vector search with other RediSearch fields for hybrid search. In the below example, we will combine vector search with full text search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c94d5cce",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_hybrid_field(field_name: str, value: str) -> str:\n",
    "    return f'@{field_name}:\"{value}\"'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bfcd31c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. First War of Scottish Independence (Score: 0.892)\n",
      "1. Wars of Scottish Independence (Score: 0.889)\n",
      "2. Second War of Scottish Independence (Score: 0.879)\n",
      "3. List of Scottish monarchs (Score: 0.873)\n",
      "4. Scottish Borders (Score: 0.863)\n"
     ]
    }
   ],
   "source": [
    "# search the content vector for articles about famous battles in Scottish history and only include results with Scottish in the title\n",
    "results = search_redis(redis_client,\n",
    "                       \"Famous battles in Scottish history\",\n",
    "                       vector_field=\"title_vector\",\n",
    "                       k=5,\n",
    "                       hybrid_fields=create_hybrid_field(\"title\", \"Scottish\")\n",
    "                       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "28ab1e30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0. Art (Score: 1.0)\n",
      "1. Paint (Score: 0.896)\n",
      "2. Renaissance art (Score: 0.88)\n",
      "3. Painting (Score: 0.874)\n",
      "4. Renaissance (Score: 0.846)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'In Europe, after the Middle Ages, there was a \"Renaissance\" which means \"rebirth\". People rediscovered science and artists were allowed to paint subjects other than religious subjects. People like Michelangelo and Leonardo da Vinci still painted religious pictures, but they also now could paint mythological pictures too. These artists also invented perspective where things in the distance look smaller in the picture. This was new because in the Middle Ages people would paint all the figures close up and just overlapping each other. These artists used nudity regularly in their art.'"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# run a hybrid query for articles about Art in the title vector and only include results with the phrase \"Leonardo da Vinci\" in the text\n",
    "results = search_redis(redis_client,\n",
    "                       \"Art\",\n",
    "                       vector_field=\"title_vector\",\n",
    "                       k=5,\n",
    "                       hybrid_fields=create_hybrid_field(\"text\", \"Leonardo da Vinci\")\n",
    "                       )\n",
    "\n",
    "# find specific mention of Leonardo da Vinci in the text that our full-text-search query returned\n",
    "mention = [sentence for sentence in results[0].text.split(\"\\n\") if \"Leonardo da Vinci\" in sentence][0]\n",
    "mention"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f94b5be2",
   "metadata": {},
   "source": [
    "For more example with Redis as a vector database, see the README and examples within the ``vector_databases/redis`` directory of this repository"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55afccbf",
   "metadata": {},
   "source": [
    "Thanks for following along, you're now equipped to set up your own vector databases and use embeddings to do all kinds of cool things - enjoy! For more complex use cases please continue to work through other cookbook examples in this repo."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "redisvl2",
   "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.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "9b1e6e9c2967143209c2f955cb869d1d3234f92dc4787f49f155f3abbdfb1316"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}