vector_databases/vector_databases.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb1537e6",
   "metadata": {},
   "source": [
    "# Vector databases in Azure"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ebc07204",
   "metadata": {},
   "source": [
    "### Databases used in this notebook\n",
    "- [Azure Cosmos DB for MongoDB vCore](#azure-cosmos-db-for-mongodb-vcore)\n",
    "- [Azure Database for PostgreSQL](#azure-database-for-postgresql)\n",
    "- [Azure Data Explorer (Kusto)](#azure-data-explorer-kusto)\n",
    "- [Azure Cognitive Search](#azure-cognitive-search)\n",
    "- [Redis](#redis)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b3e2d562",
   "metadata": {},
   "source": [
    "## Deploy infrastructure with Terraform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8616a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!terraform init -upgrade\n",
    "!terraform apply -auto-approve"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3018238f",
   "metadata": {},
   "source": [
    "## Install and import Python packages"
   ]
  },
  {
   "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 wget --user\n",
    "%pip install pymongo --user\n",
    "%pip install psycopg2 --user\n",
    "%pip install azure.kusto.data azure.kusto.ingest --user\n",
    "%pip install redis --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9ee4dc5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import wget\n",
    "import zipfile\n",
    "import pandas as pd\n",
    "import os\n",
    "import subprocess\n",
    "import numpy as np"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e5d9d2e1",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5dff8b55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'vector_database_wikipedia_articles_embedded.zip'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Download embeddings\n",
    "embeddings_url = 'https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip'\n",
    "wget.download(embeddings_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "21097972",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unzip embeddings\n",
    "with zipfile.ZipFile(\"vector_database_wikipedia_articles_embedded.zip\",\"r\") as zip_ref:\n",
    "    zip_ref.extractall(\"./data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "70bbd8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load file to Pandas dataframe\n",
    "article_df = pd.read_csv('./data/vector_database_wikipedia_articles_embedded.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dd07a928",
   "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": [
    "from ast import literal_eval\n",
    "\n",
    "# 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)\n",
    "\n",
    "article_df.info(show_counts=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9a95358",
   "metadata": {},
   "source": [
    "## Configurations and testing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "985f7157",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Login information\n",
    "user_name = \"tomas\"\n",
    "password = \"Azure12345678\"\n",
    "random_name = subprocess.check_output([\"terraform\", \"output\", \"-raw\", \"random_name\"]).decode(\"utf-8\")\n",
    "cs_api_key = subprocess.check_output([\"terraform\", \"output\", \"-raw\", \"cs_api_key\"]).decode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1b9b6016",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Testing parameters\n",
    "test_size = 100        # Number of vectors to test\n",
    "top_n = 3              # Number of results to return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "600312c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_vectors = np.random.rand(test_size, 1536)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e9451706",
   "metadata": {},
   "source": [
    "## Azure Cosmos DB for MongoDB vCore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "68c9358e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymongo import MongoClient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "951ed87a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get client\n",
    "CONNECTION_STRING = f\"mongodb+srv://{user_name}:{password}@{random_name}.mongocluster.cosmos.azure.com/mydb?tls=true&authMechanism=SCRAM-SHA-256&retrywrites=false&maxIdleTimeMS=120000\"\n",
    "client = MongoClient(CONNECTION_STRING)\n",
    "\n",
    "# Create database\n",
    "mydb = client[\"mydb\"]\n",
    "\n",
    "# Create collection\n",
    "mycol = mydb[\"mycol\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "ea7538a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Indexes are: ['_id_', 'vectorSearchIndex']\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create indexes\n",
    "indexes = [\n",
    "      {\n",
    "      \"name\": \"vectorSearchIndex\",\n",
    "      \"key\": {\n",
    "        \"content_vector\": \"cosmosSearch\"\n",
    "      },\n",
    "      \"cosmosSearchOptions\": {\n",
    "        \"kind\": \"vector-ivf\",\n",
    "        \"numLists\": 25,\n",
    "        \"similarity\": \"COS\",\n",
    "        \"dimensions\": 1536\n",
    "      }\n",
    "    }\n",
    "]\n",
    "mydb.command(\n",
    "    {\n",
    "        \"createIndexes\": \"mycol\",\n",
    "        \"indexes\": indexes,\n",
    "    }\n",
    ")\n",
    "print(\"Indexes are: {}\\n\".format(sorted(mycol.index_information())))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "4d755d61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pymongo.results.InsertManyResult at 0x7f1dd13836d0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Insert records\n",
    "dict_records = article_df.to_dict(\"records\")\n",
    "mycol.insert_many(dict_records)\n",
    "\n",
    "# time taken: 01:43"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "b1cf5006",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "# Query\n",
    "results = []\n",
    "for i in range(test_size):\n",
    "  query = {\n",
    "    \"$search\": {\n",
    "      \"cosmosSearch\": {\n",
    "        \"vector\": test_vectors[i].tolist(),\n",
    "        \"path\": \"content_vector\",\n",
    "        \"k\": top_n\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "  results.append(mycol.aggregate([query]))\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 03:51"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "f598520e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional commands to drop collection or database\n",
    "# mydb.drop_collection(\"mycol\")\n",
    "# client.drop_database(\"mydb\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f173cab5",
   "metadata": {},
   "source": [
    "## Azure Database for PostgreSQL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "57925f12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f1f54569",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get client\n",
    "CONNECTION_STRING_PSQL = f\"host={random_name}.postgres.database.azure.com port=5432 dbname=mydb user={user_name} password={password} sslmode=require\"\n",
    "conn = psycopg2.connect(CONNECTION_STRING_PSQL)\n",
    "cur = conn.cursor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "c40d26f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable pgvector extension\n",
    "cur.execute(\"CREATE EXTENSION IF NOT EXISTS vector;\")\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "ee1f119c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create table for article_df\n",
    "cur.execute(\n",
    "    \"\"\"\n",
    "CREATE TABLE IF NOT EXISTS mytable (\n",
    "    id INT PRIMARY KEY,\n",
    "    url TEXT,\n",
    "    title TEXT,\n",
    "    text TEXT,\n",
    "    title_vector vector(1536),\n",
    "    content_vector vector(1536),\n",
    "    vector_id TEXT\n",
    "    );\n",
    "    \"\"\"\n",
    ")\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "4f50b3c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Insert records\n",
    "insert_sql = \"INSERT INTO mytable (id, url, title, text, title_vector, content_vector, vector_id) VALUES (%s, %s, %s, %s, %s, %s, %s)\"\n",
    "data = [\n",
    "    tuple(x)\n",
    "    for x in article_df[\n",
    "        [\"id\", \"url\", \"title\", \"text\", \"title_vector\", \"content_vector\", \"vector_id\"]\n",
    "    ].to_numpy()\n",
    "]\n",
    "cur.executemany(insert_sql, data)\n",
    "conn.commit()\n",
    "\n",
    "# time taken: 04:17"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "f875cd48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "# Query (accurate search)\n",
    "results = []\n",
    "for i in range(test_size):\n",
    "    query = f\"\"\"\n",
    "    SELECT * FROM mytable\n",
    "    ORDER BY content_vector <=> '{test_vectors[i].tolist()}'\n",
    "    LIMIT {top_n};\n",
    "    \"\"\"\n",
    "    results.append(cur.execute(query))\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 00:16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fe1d0a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Enable vector cosine index for approximate search\n",
    "cur.execute(\n",
    "    \"CREATE INDEX IF NOT EXISTS content_vector_cosine_idx ON mytable USING ivfflat(content_vector) WITH (lists = 25);\"\n",
    ")\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "8543ec26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "# Query (approximate search -> using 2 probes, more probes, better recall)\n",
    "results = []\n",
    "for i in range(test_size):\n",
    "    query = f\"\"\"\n",
    "    SET ivfflat.probes = 2;\n",
    "    SELECT * FROM mytable\n",
    "    ORDER BY content_vector <=> '{test_vectors[i].tolist()}'\n",
    "    LIMIT {top_n};\n",
    "    \"\"\"\n",
    "    results.append(cur.execute(query))\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 00:16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "63ceb7a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop index\n",
    "cur.execute(\"DROP INDEX IF EXISTS content_vector_cosine_idx;\")\n",
    "conn.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "229d4c07",
   "metadata": {},
   "outputs": [],
   "source": [
    "cur.close()\n",
    "conn.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ae867d5f",
   "metadata": {},
   "source": [
    "## Azure Data Explorer (Kusto)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "b08d7637",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.kusto.data import KustoClient, KustoConnectionStringBuilder\n",
    "from azure.kusto.ingest import QueuedIngestClient, IngestionProperties, FileDescriptor\n",
    "from azure.kusto.data.data_format import DataFormat\n",
    "from azure.kusto.data.helpers import dataframe_from_result_table\n",
    "import io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "059979a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get client\n",
    "AAD_TENANT_ID = \"d6af5f85-2a50-4370-b4b5-9b9a55bcb0dc\"\n",
    "KUSTO_URI = f\"https://{random_name}.westeurope.kusto.windows.net/\"\n",
    "KUSTO_INGEST_URI = f\"https://ingest-{random_name}.westeurope.kusto.windows.net/\"\n",
    "KUSTO_DATABASE = \"mydb\"\n",
    "INGEST = KustoConnectionStringBuilder.with_interactive_login(KUSTO_INGEST_URI)\n",
    "DATA = KustoConnectionStringBuilder.with_interactive_login(KUSTO_URI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c51d9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create table\n",
    "client = KustoClient(DATA)\n",
    "query = \".create table mytable (id:int, url:string, title:string, text:string, title_vector:dynamic, content_vector:dynamic, vector_id:string)\"\n",
    "client.execute_mgmt(KUSTO_DATABASE, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "4eeb12e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ingestion client\n",
    "ingestion_client = QueuedIngestClient(INGEST)\n",
    "ingestion_props = IngestionProperties(database=KUSTO_DATABASE, table=\"mytable\", data_format=DataFormat.CSV)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "db047c15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create file without header\n",
    "with open(\"./data/vector_database_wikipedia_articles_embedded.csv\", \"r\", encoding=\"utf8\") as input_file, open(\"./data/kusto.csv\", \"w\", encoding=\"utf8\") as output_file:\n",
    "    lines = input_file.readlines()\n",
    "    output_file.writelines(lines[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6a2bdfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Ingest file\n",
    "ingestion_client.ingest_from_file(\"./data/kusto.csv\", ingestion_properties=ingestion_props)\n",
    "\n",
    "# time taken: submitting 05:00, but than background ingestion process tool another 10 minutes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "0f4d04f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25000"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Wait for ingestion to finish\n",
    "# We should see 25000 rows\n",
    "query = \"\"\"\n",
    "mytable\n",
    " | count\n",
    " \"\"\"\n",
    "\n",
    "result = client.execute_query(KUSTO_DATABASE, query)\n",
    "dataframe_from_result_table(result.primary_results[0]).Count.values[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "8abec375",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "# Query\n",
    "results = []\n",
    "for i in range(test_size):\n",
    "    query = f\"\"\"\n",
    "    let series_cosine_similarity_fl=(vec1:dynamic, vec2:dynamic, vec1_size:real=double(null), vec2_size:real=double(null))\n",
    "    {{\n",
    "        let dp = series_dot_product(vec1, vec2);\n",
    "        let v1l = iff(isnull(vec1_size), sqrt(series_dot_product(vec1, vec1)), vec1_size);\n",
    "        let v2l = iff(isnull(vec2_size), sqrt(series_dot_product(vec2, vec2)), vec2_size);\n",
    "        dp/(v1l*v2l)\n",
    "    }};\n",
    "    mytable\n",
    "    | extend similarity = series_cosine_similarity_fl(dynamic({list(test_vectors[i])}),content_vector,1536,1536) \n",
    "    | top {top_n} by similarity desc\n",
    "    \"\"\"\n",
    "\n",
    "    results.append(client.execute(KUSTO_DATABASE, query))\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 07:04"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d8eeabcf",
   "metadata": {},
   "source": [
    "## Azure Cognitive Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46eb941b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --index-url=https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/ azure-search-documents==11.4.0a20230509004"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "55037603",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.core.credentials import AzureKeyCredential \n",
    "from azure.search.documents import SearchClient  \n",
    "from azure.search.documents.indexes import SearchIndexClient  \n",
    "from azure.search.documents.models import Vector  \n",
    "from azure.search.documents.indexes.models import (  \n",
    "    SearchIndex,  \n",
    "    SearchField,  \n",
    "    SearchFieldDataType,  \n",
    "    SimpleField,  \n",
    "    SearchableField,  \n",
    "    SearchIndex,  \n",
    "    SemanticConfiguration,  \n",
    "    PrioritizedFields,  \n",
    "    SemanticField,  \n",
    "    SearchField,  \n",
    "    SemanticSettings,  \n",
    "    VectorSearch,  \n",
    "    VectorSearchAlgorithmConfiguration,  \n",
    ")  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "decab46f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configurations\n",
    "CS_ENDPOINT = f\"https://{random_name}.search.windows.net\"\n",
    "CS_API_KEY = cs_api_key\n",
    "INDEX_NAME = \"myindex\"\n",
    "\n",
    "credential = AzureKeyCredential(CS_API_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "b85a778d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " myindex created\n"
     ]
    }
   ],
   "source": [
    "# Create a search index\n",
    "index_client = SearchIndexClient(endpoint=CS_ENDPOINT, credential=credential)\n",
    "fields = [\n",
    "    SimpleField(name=\"vector_id\", type=SearchFieldDataType.String, key=True),\n",
    "    SimpleField(name=\"id\", type=SearchFieldDataType.Int64, key=False),\n",
    "    SearchableField(\n",
    "        name=\"title\", type=SearchFieldDataType.String, searchable=True, retrievable=True\n",
    "    ),\n",
    "    SearchableField(\n",
    "        name=\"url\", type=SearchFieldDataType.String, searchable=True, retrievable=True\n",
    "    ),\n",
    "    SearchableField(\n",
    "        name=\"text\", type=SearchFieldDataType.String, searchable=True, retrievable=True\n",
    "    ),\n",
    "    SearchField(\n",
    "        name=\"title_vector\",\n",
    "        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),\n",
    "        searchable=True,\n",
    "        dimensions=1536,\n",
    "        vector_search_configuration=\"my-vector-config\",\n",
    "    ),\n",
    "    SearchField(\n",
    "        name=\"content_vector\",\n",
    "        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),\n",
    "        searchable=True,\n",
    "        dimensions=1536,\n",
    "        vector_search_configuration=\"my-vector-config\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "vector_search = VectorSearch(\n",
    "    algorithm_configurations=[\n",
    "        VectorSearchAlgorithmConfiguration(\n",
    "            name=\"my-vector-config\",\n",
    "            kind=\"hnsw\",\n",
    "            hnsw_parameters={\n",
    "                \"m\": 4,\n",
    "                \"efConstruction\": 400,\n",
    "                \"efSearch\": 1000,\n",
    "                \"metric\": \"cosine\",\n",
    "            },\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "semantic_config = SemanticConfiguration(\n",
    "    name=\"my-semantic-config\",\n",
    "    prioritized_fields=PrioritizedFields(\n",
    "        title_field=SemanticField(field_name=\"title\"),\n",
    "        prioritized_content_fields=[SemanticField(field_name=\"text\")],\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Create the semantic settings with the configuration\n",
    "semantic_settings = SemanticSettings(configurations=[semantic_config])\n",
    "\n",
    "# Create the search index with the semantic settings\n",
    "index = SearchIndex(\n",
    "    name=INDEX_NAME,\n",
    "    fields=fields,\n",
    "    vector_search=vector_search,\n",
    "    semantic_settings=semantic_settings,\n",
    ")\n",
    "result = index_client.create_or_update_index(index)\n",
    "print(f\" {result.name} created\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bff23bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload some documents to the index\n",
    "search_client = SearchClient(endpoint=CS_ENDPOINT, index_name=INDEX_NAME, credential=credential)\n",
    "\n",
    "for index, row in article_df.iterrows():\n",
    "    search_client.upload_documents(row)\n",
    "\n",
    "print(f\"Uploaded {len(article_df)} documents\")\n",
    "\n",
    "# time taken: 36:50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "c740ba23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "# Query\n",
    "results = []\n",
    "for i in range(test_size):\n",
    "    results.append(\n",
    "        search_client.search(\n",
    "            search_text=\"\",\n",
    "            vector=Vector(\n",
    "                value=test_vectors[i], k=3, fields=\"content_vector\"\n",
    "            ),\n",
    "            select=[\"title\", \"text\", \"url\"],\n",
    "        )\n",
    "    )\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 00:00 (instant)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c216b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print results\n",
    "sindex = 0\n",
    "for result in results:\n",
    "    sindex += 1\n",
    "    print(f\"----- Search {sindex}\")\n",
    "    for row in result:\n",
    "        print(row)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bfbb3f51",
   "metadata": {},
   "source": [
    "## Redis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "a9bc35ff",
   "metadata": {},
   "outputs": [],
   "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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "c001c4d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "REDIS_HOST = f\"redis{random_name}.westeurope.azurecontainer.io\"\n",
    "REDIS_PORT = 6379\n",
    "REDIS_PASSWORD = password\n",
    "\n",
    "# Connect to Redis\n",
    "redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD)\n",
    "redis_client.ping()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "723d84f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constants\n",
    "VECTOR_DIM = 1536\n",
    "VECTOR_NUMBER = 25000\n",
    "INDEX_NAME = \"embeddings-index\"\n",
    "PREFIX = \"doc\"\n",
    "DISTANCE_METRIC = \"COSINE\"\n",
    "\n",
    "# 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(\n",
    "    \"title_vector\",\n",
    "    \"FLAT\",\n",
    "    {\n",
    "        \"TYPE\": \"FLOAT32\",\n",
    "        \"DIM\": VECTOR_DIM,\n",
    "        \"DISTANCE_METRIC\": DISTANCE_METRIC,\n",
    "        \"INITIAL_CAP\": VECTOR_NUMBER,\n",
    "    },\n",
    ")\n",
    "text_embedding = VectorField(\n",
    "    \"content_vector\",\n",
    "    \"FLAT\",\n",
    "    {\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": 112,
   "id": "4cbb977e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Index already exists\n"
     ]
    }
   ],
   "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",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "id": "88aceec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ingest data\n",
    "for i in range(len(article_df)):\n",
    "    key = f\"{PREFIX}:{str(article_df['id'][i])}\"\n",
    "    doc = {}\n",
    "    doc[\"url\"] = article_df['url'][i]\n",
    "    doc[\"title\"] = article_df['title'][i]\n",
    "    doc[\"text\"] = article_df['text'][i]\n",
    "    doc[\"title_vector\"] = np.array(article_df[\"title_vector\"][i], dtype=np.float32).tobytes()\n",
    "    doc[\"content_vector\"] = np.array(article_df[\"content_vector\"][i], dtype=np.float32).tobytes()\n",
    "    redis_client.hset(key, mapping = doc)\n",
    "\n",
    "# time taken: 02:27"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "cda094e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded 25000 documents in Redis search index with name: embeddings-index\n"
     ]
    }
   ],
   "source": [
    "# Check number of documents in index\n",
    "print(f\"Loaded {redis_client.info()['db0']['keys']} documents in Redis search index with name: {INDEX_NAME}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "62cf24f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Got 100 results\n"
     ]
    }
   ],
   "source": [
    "results = []\n",
    "for i in range(test_size):\n",
    "    base_query = f\"*=>[KNN {top_n} @content_vector $vector AS vector_score]\"\n",
    "    query = Query(base_query).sort_by(\"vector_score\").paging(0, top_n).dialect(2)\n",
    "    params_dict = {\n",
    "        \"vector\": np.array(test_vectors[i]).astype(dtype=np.float32).tobytes()\n",
    "    }\n",
    "    results.append(redis_client.ft(INDEX_NAME).search(query, params_dict))\n",
    "\n",
    "print(f\"Got {len(results)} results\")\n",
    "\n",
    "# time taken: 00:05"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "08ff3e08",
   "metadata": {},
   "source": [
    "# Destroy infrastructure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea0a396",
   "metadata": {},
   "outputs": [],
   "source": [
    "!terraform destroy -auto-approve"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "fd16a328ca3d68029457069b79cb0b38eb39a0f5ccc4fe4473d3047707df8207"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}