{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xv-ufMosI466"
   },
   "source": [
    "# Module 4: Build to Read: A Complete GraphRAG Tutorial\n",
    "\n",
    "In this module, we will use DSPy to extract entities and relationships from text into a graph. And store entities and relationships in TiDB Serverless. At the end, we will query the graph within TiDB Serverless to boost RAG performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uCddkwYubxYz"
   },
   "source": [
    "## Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "J1twp_Mabzuw"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "%pip install PyMySQL==1.1.0 \\\n",
    "    SQLAlchemy==2.0.30 \\\n",
    "    tidb-vector==0.0.14 \\\n",
    "    pydantic==2.7.1 \\\n",
    "    pydantic_core==2.18.2 \\\n",
    "    dspy-ai==2.4.9 \\\n",
    "    langchain-community==0.2.0 \\\n",
    "    wikipedia==1.4.0 \\\n",
    "    pyvis==0.3.1 \\\n",
    "    openai==1.72.0\n",
    "\n",
    "import pymysql\n",
    "import dspy\n",
    "import enum\n",
    "import openai\n",
    "import tidb_vector\n",
    "import os\n",
    "\n",
    "from pymysql import Connection\n",
    "from pymysql.cursors import DictCursor\n",
    "from dspy.functional import TypedPredictor\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import Mapping, Any, Optional, List\n",
    "from langchain_community.document_loaders import WikipediaLoader\n",
    "from pyvis.network import Network\n",
    "from IPython.display import HTML\n",
    "\n",
    "from sqlalchemy import (\n",
    "    Column,\n",
    "    Integer,\n",
    "    String,\n",
    "    Text,\n",
    "    JSON,\n",
    "    ForeignKey,\n",
    "    BLOB,\n",
    "    Enum as SQLEnum,\n",
    "    DateTime,\n",
    "    URL,\n",
    "    create_engine,\n",
    "    or_,\n",
    ")\n",
    "from sqlalchemy.orm import relationship, Session, declarative_base, joinedload\n",
    "from tidb_vector.sqlalchemy import VectorType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "\n",
    "OPENAI_API_KEY = getpass.getpass(\"Enter your OpenAI API key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RATtBpbvNpWb"
   },
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "UPYR-X_e78N-"
   },
   "outputs": [],
   "source": [
    "# DSPy Part\n",
    "\n",
    "class Entity(BaseModel):\n",
    "    \"\"\"List of entities extracted from the text to form the knowledge graph\"\"\"\n",
    "\n",
    "    name: str = Field(\n",
    "        description=\"Name of the entity, it should be a clear and concise term\"\n",
    "    )\n",
    "    description: str = Field(\n",
    "        description=(\n",
    "            \"Description of the entity, it should be a complete and comprehensive sentence, not few words. \"\n",
    "            \"Sample description of entity 'TiDB in-place upgrade': \"\n",
    "            \"'Upgrade TiDB component binary files to achieve upgrade, generally use rolling upgrade method'\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "class Relationship(BaseModel):\n",
    "    \"\"\"List of relationships extracted from the text to form the knowledge graph\"\"\"\n",
    "\n",
    "    source_entity: str = Field(\n",
    "        description=\"Source entity name of the relationship, it should an existing entity in the Entity list\"\n",
    "    )\n",
    "    target_entity: str = Field(\n",
    "        description=\"Target entity name of the relationship, it should an existing entity in the Entity list\"\n",
    "    )\n",
    "    relationship_desc: str = Field(\n",
    "        description=(\n",
    "            \"Description of the relationship, it should be a complete and comprehensive sentence, not few words. \"\n",
    "            \"Sample relationship description: 'TiDB will release a new LTS version every 6 months.'\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "class KnowledgeGraph(BaseModel):\n",
    "    \"\"\"Graph representation of the knowledge for text.\"\"\"\n",
    "\n",
    "    entities: List[Entity] = Field(\n",
    "        description=\"List of entities in the knowledge graph\"\n",
    "    )\n",
    "    relationships: List[Relationship] = Field(\n",
    "        description=\"List of relationships in the knowledge graph\"\n",
    "    )\n",
    "\n",
    "class ExtractGraphTriplet(dspy.Signature):\n",
    "    \"\"\"Carefully analyze the provided text from database documentation and community blogs to thoroughly identify all entities related to database technologies, including both general concepts and specific details.\n",
    "\n",
    "    Follow these Step-by-Step Analysis:\n",
    "\n",
    "    1. Extract Meaningful Entities:\n",
    "      - Identify all significant nouns, proper nouns, and technical terminologies that represent database-related concepts, objects, components, features, issues, key steps, execute order, user case, locations, versions, or any substantial entities.\n",
    "      - Ensure that you capture entities across different levels of detail, from high-level overviews to specific technical specifications, to create a comprehensive representation of the subject matter.\n",
    "      - Choose names for entities that are specific enough to indicate their meaning without additional context, avoiding overly generic terms.\n",
    "      - Consolidate similar entities to avoid redundancy, ensuring each represents a distinct concept at appropriate granularity levels.\n",
    "\n",
    "    2. Establish Relationships:\n",
    "      - Carefully examine the text to identify all relationships between clearly-related entities, ensuring each relationship is correctly captured with accurate details about the interactions.\n",
    "      - Analyze the context and interactions between the identified entities to determine how they are interconnected, focusing on actions, associations, dependencies, or similarities.\n",
    "      - Clearly define the relationships, ensuring accurate directionality that reflects the logical or functional dependencies among entities. \\\n",
    "         This means identifying which entity is the source, which is the target, and what the nature of their relationship is (e.g., $source_entity depends on $target_entity for $relationship).\n",
    "\n",
    "    Some key points to consider:\n",
    "      - Please endeavor to extract all meaningful entities and relationships from the text, avoid subsequent additional gleanings.\n",
    "\n",
    "    Objective: Produce a detailed and comprehensive knowledge graph that captures the full spectrum of entities mentioned in the text, along with their interrelations, reflecting both broad concepts and intricate details specific to the database domain.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    text = dspy.InputField(\n",
    "        desc=\"a paragraph of text to extract entities and relationships to form a knowledge graph\"\n",
    "    )\n",
    "    knowledge: KnowledgeGraph = dspy.OutputField(\n",
    "        desc=\"Graph representation of the knowledge extracted from the text.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class Extractor(dspy.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.prog_graph = TypedPredictor(ExtractGraphTriplet)\n",
    "\n",
    "    def forward(self, text):\n",
    "        return self.prog_graph(\n",
    "            text=text,\n",
    "            config={\n",
    "                \"response_format\": {\"type\": \"json_object\"},\n",
    "            },\n",
    "        )\n",
    "\n",
    "def jupyter_interactive_graph(kg: KnowledgeGraph) -> str:\n",
    "    net = Network(notebook=True, cdn_resources='remote')\n",
    "\n",
    "    node_map = {}\n",
    "    for index in range(len(kg.entities)):\n",
    "        node_map[kg.entities[index].name] = index\n",
    "        net.add_node(\n",
    "            index,\n",
    "            label=kg.entities[index].name,\n",
    "            title=kg.entities[index].description\n",
    "        )\n",
    "\n",
    "    for index in range(len(kg.relationships)):\n",
    "        relation = kg.relationships[index]\n",
    "        src_index = node_map[relation.source_entity]\n",
    "        target_index = node_map[relation.target_entity]\n",
    "        net.add_edge(src_index, target_index)\n",
    "\n",
    "    filename = \"kg.html\"\n",
    "    net.save_graph(filename)\n",
    "\n",
    "    return filename\n",
    "\n",
    "\n",
    "# OpenAI Part\n",
    "\n",
    "def get_query_embedding(query: str):\n",
    "    open_ai_client = openai.OpenAI(api_key=OPENAI_API_KEY)\n",
    "    response = open_ai_client.embeddings.create(input=[query], model=\"text-embedding-3-small\")\n",
    "    return response.data[0].embedding\n",
    "\n",
    "\n",
    "def generate_result(query: str, entities, relationships):\n",
    "    open_ai_client = openai.OpenAI(api_key=OPENAI_API_KEY)\n",
    "    entities_prompt = '\\n'.join(map(lambda e: f'(Name: \"{e.name}\", Description: \"{e.description}\")', entities))\n",
    "    relationships_prompt = '\\n'.join(map(lambda r: f'\"{r.relationship_desc}\"', relationships))\n",
    "\n",
    "    response = open_ai_client.chat.completions.create(\n",
    "        model=\"gpt-4o\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"Please carefully think the user's \" +\n",
    "             \"question and ONLY use the content below to generate answer:\\n\" +\n",
    "             f\"Entities: {entities_prompt}, Relationships: {relationships_prompt}\"},\n",
    "            {\"role\": \"user\", \"content\": query}\n",
    "        ])\n",
    "\n",
    "    return response.choices[0].message.content\n",
    "\n",
    "\n",
    "# TiDB Serverless Database Part\n",
    "\n",
    "def get_db_url():\n",
    "    return 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",
    "    )\n",
    "\n",
    "engine = create_engine(get_db_url(), pool_recycle=300)\n",
    "Base = declarative_base()\n",
    "\n",
    "class DatabaseEntity(Base):\n",
    "    id = Column(Integer, primary_key=True)\n",
    "    name = Column(String(512))\n",
    "    description = Column(Text)\n",
    "    description_vec = Column(VectorType(1536))\n",
    "\n",
    "    __tablename__ = \"entities\"\n",
    "\n",
    "\n",
    "class DatabaseRelationship(Base):\n",
    "    id = Column(Integer, primary_key=True)\n",
    "    source_entity_id = Column(Integer, ForeignKey(\"entities.id\"))\n",
    "    target_entity_id = Column(Integer, ForeignKey(\"entities.id\"))\n",
    "    relationship_desc = Column(Text)\n",
    "\n",
    "    source_entity = relationship(\"DatabaseEntity\", foreign_keys=[source_entity_id])\n",
    "    target_entity = relationship(\"DatabaseEntity\", foreign_keys=[target_entity_id])\n",
    "\n",
    "    __tablename__ = \"relationships\"\n",
    "\n",
    "def clean_knowledge_graph(kg: KnowledgeGraph) -> KnowledgeGraph:\n",
    "    entity_name_set = set(map(lambda e: e.name, kg.entities))\n",
    "    kg.relationships = list(filter(\n",
    "        lambda r: r.source_entity in entity_name_set and r.target_entity in entity_name_set,\n",
    "        kg.relationships))\n",
    "\n",
    "    return kg\n",
    "\n",
    "def save_knowledge_graph(kg: KnowledgeGraph):\n",
    "    data_entities = list(map(lambda e: DatabaseEntity(\n",
    "        name = e.name,\n",
    "        description = e.description,\n",
    "        description_vec = get_query_embedding(e.description)\n",
    "    ), kg.entities))\n",
    "\n",
    "    with Session(engine) as session:\n",
    "        session.add_all(data_entities)\n",
    "        # get increment ids\n",
    "        session.flush()\n",
    "\n",
    "        entity_id_map = dict(map(lambda e: (e.name, e.id), data_entities))\n",
    "        data_relationships = list(map(lambda r: DatabaseRelationship(\n",
    "            source_entity_id = entity_id_map[r.source_entity],\n",
    "            target_entity_id = entity_id_map[r.target_entity],\n",
    "            relationship_desc = r.relationship_desc\n",
    "        ), kg.relationships))\n",
    "\n",
    "        session.add_all(data_relationships)\n",
    "        session.commit()\n",
    "\n",
    "def retrieve_entities_relationships(question_embedding) -> (List[DatabaseEntity], List[DatabaseRelationship]) :\n",
    "    with Session(engine) as session:\n",
    "        entity = session.query(DatabaseEntity) \\\n",
    "            .order_by(DatabaseEntity.description_vec.cosine_distance(question_embedding)) \\\n",
    "            .limit(1).first()\n",
    "        entities = {entity.id: entity}\n",
    "\n",
    "        relationships = session.query(DatabaseRelationship).options(\n",
    "            joinedload(DatabaseRelationship.source_entity),\n",
    "            joinedload(DatabaseRelationship.target_entity),\n",
    "        ).filter(\n",
    "            or_(\n",
    "                DatabaseRelationship.source_entity == entity,\n",
    "                DatabaseRelationship.target_entity == entity\n",
    "            )\n",
    "        )\n",
    "\n",
    "        for r in relationships:\n",
    "            entities.update({\n",
    "                r.source_entity.id: r.source_entity,\n",
    "                r.target_entity.id: r.target_entity,\n",
    "            })\n",
    "\n",
    "        return entities.values(), relationships\n",
    "\n",
    "# Initial\n",
    "\n",
    "extractor = Extractor()\n",
    "Base.metadata.drop_all(engine)\n",
    "Base.metadata.create_all(engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mal3vFr5ZhKd"
   },
   "source": [
    "## Core Code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7hPnNXv0OMqT"
   },
   "source": [
    "### Part 1. Indexing\n",
    "\n",
    "Indexing in terms of RAG is the process of organizing a vast amount of text data in a way that allows the RAG system to quickly find the most relevant pieces of information for a given query. [\\[1\\]](https://medium.com/@j13mehul/rag-part-4-indexing-1985f4000f72#)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GhlzJOqHZmGV"
   },
   "source": [
    "#### Set OpenAI and DSPy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "GalPqH4bZn5k"
   },
   "outputs": [],
   "source": [
    "open_ai_client = dspy.OpenAI(model=\"gpt-4o\", api_key=OPENAI_API_KEY, max_tokens=4096)\n",
    "dspy.settings.configure(lm=open_ai_client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fxMzegLUZuNx"
   },
   "source": [
    "#### Load Raw Wikipedia Page"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true,
    "id": "XxASIQfcDxk_"
   },
   "outputs": [],
   "source": [
    "wiki = WikipediaLoader(query=\"TiDB\").load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mv8fv4VPZ6mJ"
   },
   "source": [
    "#### Extract Raw Wikipedia Page to Knowledge Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "V6KywhEJMe_E"
   },
   "outputs": [],
   "source": [
    "pred = extractor(text = wiki[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clean Knowledge Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "knowledge_graph = clean_knowledge_graph(pred.knowledge)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Q1aZdrdaIX1"
   },
   "source": [
    "#### Let's Show the Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8cM2hau4S3kO"
   },
   "outputs": [],
   "source": [
    "HTML(filename=jupyter_interactive_graph(pred.knowledge))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j139JoygEgc6"
   },
   "source": [
    "#### Save Graph to TiDB Serverless"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "YXvTxGbIEnPy"
   },
   "outputs": [],
   "source": [
    "save_knowledge_graph(pred.knowledge)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y0zLR_ZuqsXy"
   },
   "source": [
    "### Part 2. Retrieve\n",
    "\n",
    "After indexing, we can retrieve data from the graph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aCAeZfZd992I"
   },
   "source": [
    "#### Ask Question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "cellView": "form",
    "id": "DWw2lyLmEf-P"
   },
   "outputs": [],
   "source": [
    "question = \"What is the relationship between PD and Raft?\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kcvp8wi9En-7"
   },
   "source": [
    "#### Find Entites and Relationships\n",
    "\n",
    "In this case, we will get the nearest entities, by using embedding vector which the feature offered by TiDB Serverless. Then, get the nearest neighbors of this node, and the relationships between them.\n",
    "\n",
    "![!40](https://public-pingcap-innovation.s3.us-west-2.amazonaws.com/demo-408/408-tidb-graphrag-0410.png)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "nbtaoDbtxEUB"
   },
   "outputs": [],
   "source": [
    "question_embedding = get_query_embedding(question)\n",
    "entities, relationships = retrieve_entities_relationships(question_embedding)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RP6w4W2p-zz3"
   },
   "source": [
    "### Part 3. Generate Answer\n",
    "\n",
    "Once we got the entities and relationships, we can generate the answer by laveraging the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yHyc9qFb-zQJ"
   },
   "outputs": [],
   "source": [
    "result = generate_result(question, entities, relationships)\n",
    "result"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "uCddkwYubxYz",
    "RATtBpbvNpWb"
   ],
   "provenance": [],
   "toc_visible": true
  },
  "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": 0
}