{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 1: Text Similarity Search\n",
        "\n",
        "This module uses OpenAI Embedding model and TiDB Serverless Vector Search."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F1fsS576izUl"
      },
      "source": [
        "## Install Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "pTpKX_lDizUp"
      },
      "outputs": [],
      "source": [
        "%pip install -q \\\n",
        "    openai==1.72.0 \\\n",
        "    peewee==3.17.9 \\\n",
        "    pymysql==1.1.1 \\\n",
        "    tidb_vector==0.0.14"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psEHGWiHizUq"
      },
      "source": [
        "## Preapre the environment\n",
        "\n",
        "> **Note:**\n",
        ">\n",
        "> - You can get the `OPENAI_API_KEY` from [OpenAI](https://platform.openai.com/docs/quickstart).\n",
        "> - We already set the `TIDB_HOST`, `TIDB_PORT`, `TIDB_USERNAME`, `TIDB_PASSWORD`, and `TIDB_DATABASE` in the environment parameters. If you want to use this code snippet out of TiDB Labs platform, please set them beforehand.\n",
        "\n",
        "Set the embedding model as `text-embedding-3-small`, and\n",
        "the amount of embedding dimensions is `1536`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "MgKOjwmYizUq"
      },
      "outputs": [],
      "source": [
        "import getpass\n",
        "\n",
        "OPENAI_API_KEY = getpass.getpass(\"Enter your OpenAI API key: \")\n",
        "\n",
        "embedding_model = \"text-embedding-3-small\"\n",
        "embedding_dimensions = 1536"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3WbH_BITizUr"
      },
      "source": [
        "## Initial the Clients of OpenAI and Database"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "UWtcs58-izUr"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "from openai import OpenAI\n",
        "from peewee import Model, MySQLDatabase, TextField, SQL\n",
        "from tidb_vector.peewee import VectorField\n",
        "\n",
        "client = OpenAI(api_key=OPENAI_API_KEY)\n",
        "db = MySQLDatabase(\n",
        "    os.getenv(\"TIDB_DATABASE\"),\n",
        "    user=os.getenv(\"TIDB_USERNAME\"),\n",
        "    password=os.getenv(\"TIDB_PASSWORD\"),\n",
        "    host=os.getenv(\"TIDB_HOST\"),\n",
        "    port=int(os.getenv(\"TIDB_PORT\")),\n",
        "    ssl_verify_cert=True,\n",
        "    ssl_verify_identity=True\n",
        ")\n",
        "db.connect()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uOyjrmWJizUr"
      },
      "source": [
        "## Prepare the Context\n",
        "\n",
        "In this case, contexts are the documents, use the openai embeddings model to get the embeddings of the documents, and store them in the TiDB."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "_e5P_m0MizUs"
      },
      "outputs": [],
      "source": [
        "documents = [\n",
        "   \"TiDB is an open-source distributed SQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.\",\n",
        "   \"TiFlash is the key component that makes TiDB essentially an Hybrid Transactional/Analytical Processing (HTAP) database. As a columnar storage extension of TiKV, TiFlash provides both good isolation level and strong consistency guarantee.\",\n",
        "   \"TiKV is a distributed and transactional key-value database, which provides transactional APIs with ACID compliance. With the implementation of the Raft consensus algorithm and consensus state stored in RocksDB, TiKV guarantees data consistency between multiple replicas and high availability. \",\n",
        "]\n",
        "\n",
        "class DocModel(Model):\n",
        "    text = TextField()\n",
        "    embedding = VectorField(dimensions=embedding_dimensions)\n",
        "\n",
        "    class Meta:\n",
        "        database = db\n",
        "        table_name = \"openai_embedding_test\"\n",
        "\n",
        "    def __str__(self):\n",
        "        return self.text\n",
        "\n",
        "db.drop_tables([DocModel])\n",
        "db.create_tables([DocModel])\n",
        "\n",
        "embeddings = [\n",
        "    r.embedding\n",
        "    for r in client.embeddings.create(\n",
        "      input=documents, model=embedding_model\n",
        "    ).data\n",
        "]\n",
        "data_source = [\n",
        "    {\"text\": doc, \"embedding\": emb}\n",
        "    for doc, emb in zip(documents, embeddings)\n",
        "]\n",
        "DocModel.insert_many(data_source).execute()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zMP-P1g8izUs"
      },
      "source": [
        "## Initial the Vector of Question\n",
        "\n",
        "Ask a question, use the openai embeddings model to get the embeddings of the question"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "-zrTOxs4izUt"
      },
      "outputs": [],
      "source": [
        "question = \"what is TiKV?\"\n",
        "question_embedding = client.embeddings.create(input=question, model=embedding_model).data[0].embedding"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "atc0gXVZizUt"
      },
      "source": [
        "## Retrieve by Cosine Distance of Vectors\n",
        "Get the relevant documents from the TiDB by comparing the embeddings of the question and the documents"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "DTtJRX64izUt"
      },
      "outputs": [],
      "source": [
        "related_docs = DocModel.select(\n",
        "    DocModel.text, DocModel.embedding.cosine_distance(question_embedding).alias(\"distance\")\n",
        ").order_by(SQL(\"distance\")).limit(3)\n",
        "\n",
        "print(\"Question:\", question)\n",
        "print(\"Related documents:\")\n",
        "for doc in related_docs:\n",
        "    print(doc.distance, doc.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bYBetPchmNUp"
      },
      "source": [
        "## Cleanup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "Lh27gC7gizUt"
      },
      "outputs": [],
      "source": [
        "db.close()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}