{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 3: GraphRAG Retrieval\n",
    "\n",
    "This module uses TiDB Serverless Vector Search to search an exist graph.\n",
    "\n",
    "The the dimension of vector is 512. We will use it to find the initial node, and use graph to expand 1 degree."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q \\\n",
    "    PyMySQL==1.1.0 \\\n",
    "    openai==1.27.0 \\\n",
    "    SQLAlchemy==2.0.30 \\\n",
    "    mycli"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create TiDB Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from sqlalchemy import create_engine, URL\n",
    "\n",
    "engine = create_engine(URL(\n",
    "    \"mysql+pymysql\",\n",
    "    username=os.getenv(\"TIDB_USERNAME\"),\n",
    "    password=os.getenv(\"TIDB_PASSWORD\"),\n",
    "    host=os.getenv(\"TIDB_HOST\"),\n",
    "    port=int(os.getenv(\"TIDB_PORT\")),\n",
    "    database=os.getenv(\"TIDB_DATABASE\"),\n",
    "    query={\"ssl_verify_cert\": True, \"ssl_verify_identity\": True},\n",
    "))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Prebuilt Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from sqlalchemy import text\n",
    "\n",
    "url = 'https://gist.githubusercontent.com/Icemap/a0eefdc6951a30c0c2d2df713021055e/raw/61d8af5e7ff0cf85fbef76570fb007ad4fd7624f/prebuilt-data.sql'\n",
    "response = requests.get(url)\n",
    "sql_lines = response.content.decode('utf-8').splitlines()\n",
    "\n",
    "with engine.connect() as connection:\n",
    "    connection.execute(text(\"DROP TABLE IF EXISTS `relationships`;\"))\n",
    "    connection.execute(text(\"DROP TABLE IF EXISTS `entities`;\"))\n",
    "    connection.commit()\n",
    "\n",
    "with engine.connect() as connection:\n",
    "    for line in sql_lines:\n",
    "        line = line.strip()\n",
    "        if line:\n",
    "            connection.execute(text(line))\n",
    "    connection.commit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query Data\n",
    "\n",
    "After importing data into TiDB Serverless, we can query it by asking a question, like \"What is TiKV?\" or another you like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import text\n",
    "import openai\n",
    "import getpass\n",
    "\n",
    "# TiDB Connection String Pattern:\n",
    "# mysql+pymysql://{TIDB_USER}:{TIDB_PASSWORD}@{TIDB_HOST}:{TIDB_PORT}/{TIDB_DB_NAME}?ssl_verify_cert=True&ssl_verify_identity=True\n",
    "\n",
    "oai_cli = openai.OpenAI(api_key=getpass.getpass(\"Input your OpenAI API Key:\"))\n",
    "question = input(\"Enter your question:\")\n",
    "embedding = str(oai_cli.embeddings.create(input=[question], model=\"text-embedding-3-small\").data[0].embedding)\n",
    "\n",
    "with engine.connect() as conn:\n",
    "    result = conn.execute(text(\"\"\"\n",
    "    WITH initial_entity AS (\n",
    "        SELECT id FROM `entities`\n",
    "        ORDER BY VEC_Cosine_Distance(description_vec, :embedding) LIMIT 1\n",
    "    ), entities_ids AS (\n",
    "        SELECT source_entity_id i FROM relationships r INNER JOIN initial_entity i ON r.target_entity_id = i.id\n",
    "        UNION SELECT target_entity_id i FROM relationships r INNER JOIN initial_entity i ON r.source_entity_id = i.id\n",
    "        UNION SELECT initial_entity.id i FROM initial_entity\n",
    "    ) SELECT description FROM `entities` WHERE id IN (SELECT i FROM entities_ids);\"\"\"), {\"embedding\": embedding}).fetchall()\n",
    "\n",
    "    print(oai_cli.chat.completions.create(model=\"gpt-4o\", messages=[\n",
    "        {\"role\": \"system\", \"content\": f\"Please carefully answer the question by {str(result)}\"},\n",
    "        {\"role\": \"user\", \"content\": question}]).choices[0].message.content)"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}