{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "655b8f55-2089-4733-8b09-35dea9580695",
   "metadata": {},
   "source": [
    "# Google Vertex AI MatchingEngine\n",
    "\n",
    "This notebook shows how to use functionality related to the `GCP Vertex AI MatchingEngine` vector database.\n",
    "\n",
    "> Vertex AI [Matching Engine](https://cloud.google.com/vertex-ai/docs/matching-engine/overview) provides the industry's leading high-scale low latency vector database. These vector databases are commonly referred to as vector similarity-matching or an approximate nearest neighbor (ANN) service.\n",
    "\n",
    "**Note**: This module expects an endpoint and deployed index already created as the creation time takes close to one hour. To see how to create an index refer to the section [Create Index and deploy it to an Endpoint](#create-index-and-deploy-it-to-an-endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9971578-0ae9-4809-9e80-e5f9d3dcc98a",
   "metadata": {},
   "source": [
    "## Create VectorStore from texts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c96da4-8d97-4f69-8c13-d2fcafc03b05",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import MatchingEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58b70880-edd9-46f3-b769-f26c2bcc8395",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [\n",
    "    \"The cat sat on\",\n",
    "    \"the mat.\",\n",
    "    \"I like to\",\n",
    "    \"eat pizza for\",\n",
    "    \"dinner.\",\n",
    "    \"The sun sets\",\n",
    "    \"in the west.\",\n",
    "]\n",
    "\n",
    "\n",
    "vector_store = MatchingEngine.from_components(\n",
    "    texts=texts,\n",
    "    project_id=\"<my_project_id>\",\n",
    "    region=\"<my_region>\",\n",
    "    gcs_bucket_uri=\"<my_gcs_bucket>\",\n",
    "    index_id=\"<my_matching_engine_index_id>\",\n",
    "    endpoint_id=\"<my_matching_engine_endpoint_id>\",\n",
    ")\n",
    "\n",
    "vector_store.add_texts(texts=texts)\n",
    "\n",
    "vector_store.similarity_search(\"lunch\", k=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e76e05c-d4ef-49a1-b1b9-2ea989a0eda3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Create Index and deploy it to an Endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61935a91-5efb-48af-bb40-ea1e83e24974",
   "metadata": {},
   "source": [
    "### Imports, Constants and Configs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "421b66c9-5b8f-4ef7-821e-12886a62b672",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing dependencies.\n",
    "!pip install tensorflow \\\n",
    "            google-cloud-aiplatform \\\n",
    "            tensorflow-hub \\\n",
    "            tensorflow-text "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4e9cc02-371e-40a1-bce9-37ac8efdf2cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "from google.cloud import aiplatform\n",
    "import tensorflow_hub as hub\n",
    "import tensorflow_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "352a05df-6532-4aba-a36f-603327a5bc5b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "PROJECT_ID = \"<my_project_id>\"\n",
    "REGION = \"<my_region>\"\n",
    "VPC_NETWORK = \"<my_vpc_network_name>\"\n",
    "PEERING_RANGE_NAME = \"ann-langchain-me-range\"  # Name for creating the VPC peering.\n",
    "BUCKET_URI = \"gs://<bucket_uri>\"\n",
    "# The number of dimensions for the tensorflow universal sentence encoder.\n",
    "# If other embedder is used, the dimensions would probably need to change.\n",
    "DIMENSIONS = 512\n",
    "DISPLAY_NAME = \"index-test-name\"\n",
    "EMBEDDING_DIR = f\"{BUCKET_URI}/banana\"\n",
    "DEPLOYED_INDEX_ID = \"endpoint-test-name\"\n",
    "\n",
    "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
    "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
    "VPC_NETWORK_FULL = f\"projects/{PROJECT_NUMBER}/global/networks/{VPC_NETWORK}\"\n",
    "\n",
    "# Change this if you need the VPC to be created.\n",
    "CREATE_VPC = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "076e7931-f83e-4597-8748-c8004fd8de96",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the project id\n",
    "! gcloud config set project {PROJECT_ID}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4265081b-a5b7-491e-8ac5-1e26975b9974",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove the if condition to run the encapsulated code\n",
    "if CREATE_VPC:\n",
    "    # Create a VPC network\n",
    "    ! gcloud compute networks create {VPC_NETWORK} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}\n",
    "\n",
    "    # Add necessary firewall rules\n",
    "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-icmp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow icmp\n",
    "\n",
    "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-internal --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9\n",
    "\n",
    "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-rdp --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:3389\n",
    "\n",
    "    ! gcloud compute firewall-rules create {VPC_NETWORK}-allow-ssh --network {VPC_NETWORK} --priority 65534 --project {PROJECT_ID} --allow tcp:22\n",
    "\n",
    "    # Reserve IP range\n",
    "    ! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={VPC_NETWORK} --purpose=VPC_PEERING --project={PROJECT_ID} --description=\"peering range\"\n",
    "\n",
    "    # Set up peering with service networking\n",
    "    # Your account must have the \"Compute Network Admin\" role to run the following.\n",
    "    ! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={VPC_NETWORK} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dfbb847-fc53-48c1-b0f2-00d1c4330b01",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating bucket.\n",
    "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9698068-3d2f-471b-90c3-dae3e4ca6f63",
   "metadata": {},
   "source": [
    "### Using Tensorflow Universal Sentence Encoder as an Embedder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "144007e2-ddf8-43cd-ac45-848be0458ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the Universal Sentence Encoder module\n",
    "module_url = \"https://tfhub.dev/google/universal-sentence-encoder-multilingual/3\"\n",
    "model = hub.load(module_url)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94a2bdcb-c7e3-4fb0-8c97-cc1f2263f06c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate embeddings for each word\n",
    "embeddings = model([\"banana\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a4e6e99-5e42-4e55-90f6-c03aae4fbf14",
   "metadata": {},
   "source": [
    "### Inserting a test embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "024c78f3-4663-4d8f-9f3c-b7d82073ada4",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_config = {\n",
    "    \"id\": \"banana_id\",\n",
    "    \"embedding\": [float(x) for x in list(embeddings.numpy()[0])],\n",
    "}\n",
    "\n",
    "with open(\"data.json\", \"w\") as f:\n",
    "    json.dump(initial_config, f)\n",
    "\n",
    "!gsutil cp data.json {EMBEDDING_DIR}/file.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a11489f4-5904-4fc2-9178-f32c2df0406d",
   "metadata": {},
   "outputs": [],
   "source": [
    "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3c6953b-11f6-4803-bf2d-36fa42abf3c7",
   "metadata": {},
   "source": [
    "### Creating Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31c3c56-bfe0-49ec-9901-cd146f592da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_index = aiplatform.MatchingEngineIndex.create_tree_ah_index(\n",
    "    display_name=DISPLAY_NAME,\n",
    "    contents_delta_uri=EMBEDDING_DIR,\n",
    "    dimensions=DIMENSIONS,\n",
    "    approximate_neighbors_count=150,\n",
    "    distance_measure_type=\"DOT_PRODUCT_DISTANCE\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50770669-edf6-4796-9563-d1ea59cfa8e8",
   "metadata": {},
   "source": [
    "### Creating Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20c93d1b-a7d5-47b0-9c95-1aec1c62e281",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint.create(\n",
    "    display_name=f\"{DISPLAY_NAME}-endpoint\",\n",
    "    network=VPC_NETWORK_FULL,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b52df797-28db-4b4a-b79c-e8a274293a6a",
   "metadata": {},
   "source": [
    "### Deploy Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "019a7043-ad11-4a48-bec7-18928547b2ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_index_endpoint = my_index_endpoint.deploy_index(\n",
    "    index=my_index, deployed_index_id=DEPLOYED_INDEX_ID\n",
    ")\n",
    "\n",
    "my_index_endpoint.deployed_indexes"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "common-cpu.m107",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/base-cpu:m107"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}