GraphRAG_v1_using_lamaindx.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ckw8AE2LRKvk"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/cookbooks/GraphRAG_v1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3adM3-bURKvq"
      },
      "source": [
        "# GraphRAG Implementation with LlamaIndex\n",
        "\n",
        "[GraphRAG (Graphs + Retrieval Augmented Generation)](https://www.microsoft.com/en-us/research/project/graphrag/) combines the strengths of Retrieval Augmented Generation (RAG) and Query-Focused Summarization (QFS) to effectively handle complex queries over large text datasets. While RAG excels in fetching precise information, it struggles with broader queries that require thematic understanding, a challenge that QFS addresses but cannot scale well. GraphRAG integrates these approaches to offer responsive and thorough querying capabilities across extensive, diverse text corpora.\n",
        "\n",
        "\n",
        "This notebook provides guidance on constructing the GraphRAG pipeline using the LlamaIndex PropertyGraph abstractions.\n",
        "\n",
        "\n",
        "**NOTE:** This is an approximate implementation of GraphRAG. We are currently developing a series of cookbooks that will detail the exact implementation of GraphRAG."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_x5JHhEERKvr"
      },
      "source": [
        "## GraphRAG Aproach\n",
        "\n",
        "The GraphRAG involves two steps:\n",
        "\n",
        "1. Graph Generation - Creates Graph, builds communities and its summaries over the given document.\n",
        "2. Answer to the Query - Use summaries of the communities created from step-1 to answer the query.\n",
        "\n",
        "**Graph Generation:**\n",
        "\n",
        "1. **Source Documents to Text Chunks:** Source documents are divided into smaller text chunks for easier processing.\n",
        "\n",
        "2. **Text Chunks to Element Instances:** Each text chunk is analyzed to identify and extract entities and relationships, resulting in a list of tuples that represent these elements.\n",
        "\n",
        "3. **Element Instances to Element Summaries:** The extracted entities and relationships are summarized into descriptive text blocks for each element using the LLM.\n",
        "\n",
        "4. **Element Summaries to Graph Communities:** These entities, relationships and summaries form a graph, which is subsequently partitioned into communities using algorithms using Heirarchical Leiden to establish a hierarchical structure.\n",
        "\n",
        "5. **Graph Communities to Community Summaries:** The LLM generates summaries for each community, providing insights into the dataset’s overall topical structure and semantics.\n",
        "\n",
        "**Answering the Query:**\n",
        "\n",
        "**Community Summaries to Global Answers:** The summaries of the communities are utilized to respond to user queries. This involves generating intermediate answers, which are then consolidated into a comprehensive global answer.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yufcJ4ZwRKvt"
      },
      "source": [
        "## GraphRAG Pipeline Components\n",
        "\n",
        "Here are the different components we implemented to build all of the processes mentioned above.\n",
        "\n",
        "1. **Source Documents to Text Chunks:** Implemented using `SentenceSplitter` with a chunk size of 1024 and chunk overlap of 20 tokens.\n",
        "\n",
        "2. **Text Chunks to Element Instances AND Element Instances to Element Summaries:** Implemented using `GraphRAGExtractor`.\n",
        "\n",
        "3. **Element Summaries to Graph Communities AND Graph Communities to Community Summaries:** Implemented using `GraphRAGStore`.\n",
        "\n",
        "4. **Community Summaries to Global Answers:** Implemented using `GraphQueryEngine`.\n",
        "\n",
        "\n",
        "Let's check into each of these components and build GraphRAG pipeline.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83kxSdwCRKvu"
      },
      "source": [
        "## Installation\n",
        "\n",
        "`graspologic` is used to use hierarchical_leiden for building communities."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g0OnltPIRKvv"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index graspologic numpy==1.24.4 scipy==1.12.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WzfFEQOORKvx"
      },
      "source": [
        "## Load Data\n",
        "\n",
        "We will use a sample news article dataset retrieved from Diffbot, which Tomaz has conveniently made available on GitHub for easy access.\n",
        "\n",
        "The dataset contains 2,500 samples; for ease of experimentation, we will use 50 of these samples, which include the `title` and `text` of news articles."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "upD3lvcERKvy",
        "outputId": "1d4840e4-4562-4f37-c08e-431898af9e32"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "summary": "{\n  \"name\": \"news\",\n  \"rows\": 50,\n  \"fields\": [\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 49,\n        \"samples\": [\n          \"NFL Rumors: What Patriots \\u2018Made Clear\\u2019 To DeAndre Hopkins\\u2019 Reps\",\n          \"Stellantis to close Illinois assembly plant, lay off workers\",\n          \"Uber to shut food delivery business in Italy, exit Israel\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"date\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"num_unique_values\": 38,\n        \"samples\": [\n          \"2023-06-15T12:14:00.000000000+00:00\",\n          \"2023-06-15T12:11:00.000000000+00:00\",\n          \"2023-06-15T13:50:00.000000000+00:00\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 50,\n        \"samples\": [\n          \"The Patriots apparently didn\\u2019t waste any time courting DeAndre Hopkins once he hit the open market.\\nHopkins officially became an NFL free agent when he was released by the Arizona Cardinals on May 26. Around that juncture, New England reportedly reached out to the star wide receiver\\u2019s team to express interest. Those conversations birthed a free-agent visit, which reportedly started Wednesday in Foxboro, Mass. and carried into Thursday morning.\\n\\u201cI just got official word: DeAndre Hopkins is in New England,\\u201d NFL insider Jeremy Fowler said Thursday on ESPN\\u2019s \\u201cGet Up.\\u201d \\u201cHe\\u2019s meeting with the Patriots. He will meet with the coaches, Bill Belichick. The Patriots\\u2019 interest remains high. They made that clear with Hopkins\\u2019 representatives from the very beginning when he became a free agent. So, we\\u2019ll see if they can close on this thing.\\u201d\\nHopkins has received co-signs from New England team leaders like Mac Jones and Matthew Judon, and the Patriots reportedly are optimistic they can add the five-time Pro Bowl selection. But even if Hopkins\\u2019 meeting with the Patriots goes well, don\\u2019t be surprised if the 31-year-old leaves New England without a deal.\\n\\u201cI\\u2019m also told Hopkins (is) not in a major rush,\\u201d Fowler said. \\u201cHe\\u2019s looking at more of training camp, even early August as more of a harder deadline for him to sign somewhere. So, it would have to be a very sweet offer for New England to make that happen today.\\u201d\\nThe Patriots appear to be the second team Hopkins met with as a free agent. Prior to his trip to New England, the veteran pass-catcher talked shop with the Titans in Nashville.\",\n          \"Generation Investment Management, an investment management firm, released its \\u201cGlobal Equity Strategy\\u201d first quarter 2023 investor letter. A copy of the same can be downloaded here. The strategy performed roughly in line with the benchmark on a rolling five-year net basis and it is about 1.3% below on a rolling three-year net basis. The main reasons for the underperformance of the strategy were overestimation of the quality of the holdings, change in the external environment, and overpayment for some companies. In addition, you can check the top 5 holdings of the fund to know its best picks in 2023.\\nGeneration Investment Management Global Equity Strategy highlighted stocks like Henry Schein, Inc. (NASDAQ:HSIC) in the first quarter 2023 investor letter. Headquartered in Melville, New York, Henry Schein, Inc. (NASDAQ:HSIC) is a healthcare products and services provider to medical, dental, and veterinary office-based practitioners. On June 14, 2023, Henry Schein, Inc. (NASDAQ:HSIC) stock closed at $74.81 per share. One-month return of Henry Schein, Inc. (NASDAQ:HSIC) was -2.27%, and its shares gained 2.24% of their value over the last 52 weeks. Henry Schein, Inc. (NASDAQ:HSIC) has a market capitalization of $9.8 billion.\\nGeneration Investment Management Global Equity Strategy made the following comment about Henry Schein, Inc. (NASDAQ:HSIC) in its first quarter 2023 investor letter:\\n\\\"We continue to evaluate new investment opportunities. Despite the volatility of recent months, we see a number of strong tailwinds in certain sectors \\u2014 particularly those that solve big societal problems. One such problem is growing healthcare costs, which will put more pressure on social-security systems and household budgets. Later in this letter we profile Henry Schein, Inc. (NASDAQ:HSIC), the largest distributor of dental and medical products and services globally. Companies like Henry Schein have the potential to control the growth in healthcare costs, while promoting access to underserved communities.\\nAt the same time, healthcare cost pressures are escalating. To make matters worse, healthcare has worse labour shortages than most industries, while productivity improvements are hard to come by.4 In almost every country, health spending is forecast to form an increasing share of GDP in the coming years.\\nAt Generation, we look for companies that can provide a solution to these long-term challenges. More specifically, we look for healthcare companies that reduce costs and drive efficiency; that improve clinical outcomes; and that improve access to care. Henry Schein, a long-term holding of your portfolio, is an example of a company that we believe can deliver on all three of these criteria...\\\" (Click here to read the full text)\\nHenry Schein, Inc. (NASDAQ:HSIC) is not on our list of . As per our database, 28 hedge fund portfolios held Henry Schein, Inc. (NASDAQ:HSIC) at the end of first quarter 2023 which was 30 in the previous quarter.\\nWe discussed Henry Schein, Inc. (NASDAQ:HSIC) in and shared the list of top gainers on November 1, 2022. In addition, please check out our page for more investor letters from hedge funds and other leading investors.\\nDisclosure: None. This article is originally published at .\",\n          \"Manchester City will begin their bid for a record fourth consecutive Premier League title away to Vincent Kompany\\u2019s Burnley.\\nThe Treble winners will visit Turf Moor to face their former captain\\u2019s newly-promoted side on the evening of Friday, August 11 to raise the curtain on the 2023-24 campaign.\\nIt will be the second time Kompany \\u2013 who won the title on four occasions as City skipper \\u2013 will have faced his old side as a manager, with City running out 6-0 winners at the Etihad in March\\u2019s FA Cup quarter-final.\\nPremier League debutants Luton will play their first top-flight fixture since 1992 away to Roberto De Zerbi\\u2019s Brighton on Saturday, August 12, having to wait until the following weekend for their first home game when Kenilworth Road will become the smallest ground to host a fixture in the competition for the visit of Burnley.\\nThe other promoted side Sheffield United kick off their season with a home game against Crystal Palace.\\nThe outstanding fixture of the opening weekend will be at Stamford Bridge where Mauricio Pochettino begins life as Chelsea manager against Liverpool on Sunday, August 13, with both sides looking to bounce back after disappointing campaigns.\\nFollow live updates and reaction below\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
              "type": "dataframe",
              "variable_name": "news"
            },
            "text/html": [
              "\n",
              "  <div id=\"df-e22eeaf4-31ad-42a9-a106-9fc8facafd6d\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>title</th>\n",
              "      <th>date</th>\n",
              "      <th>text</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>Chevron: Best Of Breed</td>\n",
              "      <td>2031-04-06T01:36:32.000000000+00:00</td>\n",
              "      <td>JHVEPhoto Like many companies in the O&amp;G secto...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>FirstEnergy (NYSE:FE) Posts Earnings Results</td>\n",
              "      <td>2030-04-29T06:55:28.000000000+00:00</td>\n",
              "      <td>FirstEnergy (NYSE:FE – Get Rating) posted its ...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>Dáil almost suspended after Sinn Féin TD put p...</td>\n",
              "      <td>2023-06-15T14:32:11.000000000+00:00</td>\n",
              "      <td>The Dáil was almost suspended on Thursday afte...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>Epic’s latest tool can animate hyperrealistic ...</td>\n",
              "      <td>2023-06-15T14:00:00.000000000+00:00</td>\n",
              "      <td>Today, Epic is releasing a new tool designed t...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>EU to Ban Huawei, ZTE from Internal Commission...</td>\n",
              "      <td>2023-06-15T13:50:00.000000000+00:00</td>\n",
              "      <td>The European Commission is planning to ban equ...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-e22eeaf4-31ad-42a9-a106-9fc8facafd6d')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-e22eeaf4-31ad-42a9-a106-9fc8facafd6d button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-e22eeaf4-31ad-42a9-a106-9fc8facafd6d');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-96ffab0a-0008-4260-bba5-3389a94b32df\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-96ffab0a-0008-4260-bba5-3389a94b32df')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-96ffab0a-0008-4260-bba5-3389a94b32df button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "text/plain": [
              "                                               title  \\\n",
              "0                             Chevron: Best Of Breed   \n",
              "1       FirstEnergy (NYSE:FE) Posts Earnings Results   \n",
              "2  Dáil almost suspended after Sinn Féin TD put p...   \n",
              "3  Epic’s latest tool can animate hyperrealistic ...   \n",
              "4  EU to Ban Huawei, ZTE from Internal Commission...   \n",
              "\n",
              "                                  date  \\\n",
              "0  2031-04-06T01:36:32.000000000+00:00   \n",
              "1  2030-04-29T06:55:28.000000000+00:00   \n",
              "2  2023-06-15T14:32:11.000000000+00:00   \n",
              "3  2023-06-15T14:00:00.000000000+00:00   \n",
              "4  2023-06-15T13:50:00.000000000+00:00   \n",
              "\n",
              "                                                text  \n",
              "0  JHVEPhoto Like many companies in the O&G secto...  \n",
              "1  FirstEnergy (NYSE:FE – Get Rating) posted its ...  \n",
              "2  The Dáil was almost suspended on Thursday afte...  \n",
              "3  Today, Epic is releasing a new tool designed t...  \n",
              "4  The European Commission is planning to ban equ...  "
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import pandas as pd\n",
        "from llama_index.core import Document\n",
        "\n",
        "news = pd.read_csv(\n",
        "    \"https://raw.githubusercontent.com/tomasonjo/blog-datasets/main/news_articles.csv\"\n",
        ")[:50]\n",
        "\n",
        "news.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PqE3r8KbRKv0"
      },
      "source": [
        "Prepare documents as required by LlamaIndex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vg2WAkrORKv1"
      },
      "outputs": [],
      "source": [
        "documents = [\n",
        "    Document(text=f\"{row['title']}: {row['text']}\")\n",
        "    for i, row in news.iterrows()\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdGgjoOQRKv2"
      },
      "source": [
        "## Setup API Key and LLM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XPZtcbZHRKv2"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
        "\n",
        "from llama_index.llms.openai import OpenAI\n",
        "\n",
        "llm = OpenAI(model=\"gpt-4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ERG0CIO3RKv2"
      },
      "source": [
        "## GraphRAGExtractor\n",
        "\n",
        "The GraphRAGExtractor class is designed to extract triples (subject-relation-object) from text and enrich them by adding descriptions for entities and relationships to their properties using an LLM.\n",
        "\n",
        "This functionality is similar to that of the `SimpleLLMPathExtractor`, but includes additional enhancements to handle entity, relationship descriptions. For guidance on implementation, you may look at similar existing [extractors](https://docs.llamaindex.ai/en/latest/examples/property_graph/Dynamic_KG_Extraction/?h=comparing).\n",
        "\n",
        "Here's a breakdown of its functionality:\n",
        "\n",
        "**Key Components:**\n",
        "\n",
        "1. `llm:` The language model used for extraction.\n",
        "2. `extract_prompt:` A prompt template used to guide the LLM in extracting information.\n",
        "3. `parse_fn:` A function to parse the LLM's output into structured data.\n",
        "4. `max_paths_per_chunk:` Limits the number of triples extracted per text chunk.\n",
        "5. `num_workers:` For parallel processing of multiple text nodes.\n",
        "\n",
        "\n",
        "**Main Methods:**\n",
        "\n",
        "1. `__call__:` The entry point for processing a list of text nodes.\n",
        "2. `acall:` An asynchronous version of __call__ for improved performance.\n",
        "3. `_aextract:` The core method that processes each individual node.\n",
        "\n",
        "\n",
        "**Extraction Process:**\n",
        "\n",
        "For each input node (chunk of text):\n",
        "1. It sends the text to the LLM along with the extraction prompt.\n",
        "2. The LLM's response is parsed to extract entities, relationships, descriptions for entities and relations.\n",
        "3. Entities are converted into EntityNode objects. Entity description is stored in metadata\n",
        "4. Relationships are converted into Relation objects. Relationship description is stored in metadata.\n",
        "5. These are added to the node's metadata under KG_NODES_KEY and KG_RELATIONS_KEY.\n",
        "\n",
        "**NOTE:** In the current implementation, we are using only relationship descriptions. In the next implementation, we will utilize entity descriptions during the retrieval stage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4QWPKt_iRKv3"
      },
      "outputs": [],
      "source": [
        "import asyncio\n",
        "import nest_asyncio\n",
        "\n",
        "nest_asyncio.apply()\n",
        "\n",
        "from typing import Any, List, Callable, Optional, Union, Dict\n",
        "from IPython.display import Markdown, display\n",
        "\n",
        "from llama_index.core.async_utils import run_jobs\n",
        "from llama_index.core.indices.property_graph.utils import (\n",
        "    default_parse_triplets_fn,\n",
        ")\n",
        "from llama_index.core.graph_stores.types import (\n",
        "    EntityNode,\n",
        "    KG_NODES_KEY,\n",
        "    KG_RELATIONS_KEY,\n",
        "    Relation,\n",
        ")\n",
        "from llama_index.core.llms.llm import LLM\n",
        "from llama_index.core.prompts import PromptTemplate\n",
        "from llama_index.core.prompts.default_prompts import (\n",
        "    DEFAULT_KG_TRIPLET_EXTRACT_PROMPT,\n",
        ")\n",
        "from llama_index.core.schema import TransformComponent, BaseNode\n",
        "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
        "\n",
        "\n",
        "class GraphRAGExtractor(TransformComponent):\n",
        "    \"\"\"Extract triples from a graph.\n",
        "\n",
        "    Uses an LLM and a simple prompt + output parsing to extract paths (i.e. triples) and entity, relation descriptions from text.\n",
        "\n",
        "    Args:\n",
        "        llm (LLM):\n",
        "            The language model to use.\n",
        "        extract_prompt (Union[str, PromptTemplate]):\n",
        "            The prompt to use for extracting triples.\n",
        "        parse_fn (callable):\n",
        "            A function to parse the output of the language model.\n",
        "        num_workers (int):\n",
        "            The number of workers to use for parallel processing.\n",
        "        max_paths_per_chunk (int):\n",
        "            The maximum number of paths to extract per chunk.\n",
        "    \"\"\"\n",
        "\n",
        "    llm: LLM\n",
        "    extract_prompt: PromptTemplate\n",
        "    parse_fn: Callable\n",
        "    num_workers: int\n",
        "    max_paths_per_chunk: int\n",
        "\n",
        "    def __init__(\n",
        "        self,\n",
        "        llm: Optional[LLM] = None,\n",
        "        extract_prompt: Optional[Union[str, PromptTemplate]] = None,\n",
        "        parse_fn: Callable = default_parse_triplets_fn,\n",
        "        max_paths_per_chunk: int = 10,\n",
        "        num_workers: int = 4,\n",
        "    ) -> None:\n",
        "        \"\"\"Init params.\"\"\"\n",
        "        from llama_index.core import Settings\n",
        "\n",
        "        if isinstance(extract_prompt, str):\n",
        "            extract_prompt = PromptTemplate(extract_prompt)\n",
        "\n",
        "        super().__init__(\n",
        "            llm=llm or Settings.llm,\n",
        "            extract_prompt=extract_prompt or DEFAULT_KG_TRIPLET_EXTRACT_PROMPT,\n",
        "            parse_fn=parse_fn,\n",
        "            num_workers=num_workers,\n",
        "            max_paths_per_chunk=max_paths_per_chunk,\n",
        "        )\n",
        "\n",
        "    @classmethod\n",
        "    def class_name(cls) -> str:\n",
        "        return \"GraphExtractor\"\n",
        "\n",
        "    def __call__(\n",
        "        self, nodes: List[BaseNode], show_progress: bool = False, **kwargs: Any\n",
        "    ) -> List[BaseNode]:\n",
        "        \"\"\"Extract triples from nodes.\"\"\"\n",
        "        return asyncio.run(\n",
        "            self.acall(nodes, show_progress=show_progress, **kwargs)\n",
        "        )\n",
        "\n",
        "    async def _aextract(self, node: BaseNode) -> BaseNode:\n",
        "        \"\"\"Extract triples from a node.\"\"\"\n",
        "        assert hasattr(node, \"text\")\n",
        "\n",
        "        text = node.get_content(metadata_mode=\"llm\")\n",
        "        try:\n",
        "            llm_response = await self.llm.apredict(\n",
        "                self.extract_prompt,\n",
        "                text=text,\n",
        "                max_knowledge_triplets=self.max_paths_per_chunk,\n",
        "            )\n",
        "            entities, entities_relationship = self.parse_fn(llm_response)\n",
        "        except ValueError:\n",
        "            entities = []\n",
        "            entities_relationship = []\n",
        "\n",
        "        existing_nodes = node.metadata.pop(KG_NODES_KEY, [])\n",
        "        existing_relations = node.metadata.pop(KG_RELATIONS_KEY, [])\n",
        "        metadata = node.metadata.copy()\n",
        "        for entity, entity_type, description in entities:\n",
        "            metadata[\n",
        "                \"entity_description\"\n",
        "            ] = description  # Not used in the current implementation. But will be useful in future work.\n",
        "            entity_node = EntityNode(\n",
        "                name=entity, label=entity_type, properties=metadata\n",
        "            )\n",
        "            existing_nodes.append(entity_node)\n",
        "\n",
        "        metadata = node.metadata.copy()\n",
        "        for triple in entities_relationship:\n",
        "            subj, rel, obj, description = triple\n",
        "            subj_node = EntityNode(name=subj, properties=metadata)\n",
        "            obj_node = EntityNode(name=obj, properties=metadata)\n",
        "            metadata[\"relationship_description\"] = description\n",
        "            rel_node = Relation(\n",
        "                label=rel,\n",
        "                source_id=subj_node.id,\n",
        "                target_id=obj_node.id,\n",
        "                properties=metadata,\n",
        "            )\n",
        "\n",
        "            existing_nodes.extend([subj_node, obj_node])\n",
        "            existing_relations.append(rel_node)\n",
        "\n",
        "        node.metadata[KG_NODES_KEY] = existing_nodes\n",
        "        node.metadata[KG_RELATIONS_KEY] = existing_relations\n",
        "        return node\n",
        "\n",
        "    async def acall(\n",
        "        self, nodes: List[BaseNode], show_progress: bool = False, **kwargs: Any\n",
        "    ) -> List[BaseNode]:\n",
        "        \"\"\"Extract triples from nodes async.\"\"\"\n",
        "        jobs = []\n",
        "        for node in nodes:\n",
        "            jobs.append(self._aextract(node))\n",
        "\n",
        "        return await run_jobs(\n",
        "            jobs,\n",
        "            workers=self.num_workers,\n",
        "            show_progress=show_progress,\n",
        "            desc=\"Extracting paths from text\",\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8IObJK4mRKv4"
      },
      "source": [
        "## GraphRAGStore\n",
        "\n",
        "The `GraphRAGStore` class is an extension of the `SimplePropertyGraphStore `class, designed to implement GraphRAG pipeline. Here's a breakdown of its key components and functions:\n",
        "\n",
        "\n",
        "The class uses community detection algorithms to group related nodes in the graph and then it generates summaries for each community using an LLM.\n",
        "\n",
        "\n",
        "**Key Methods:**\n",
        "\n",
        "`build_communities():`\n",
        "\n",
        "1. Converts the internal graph representation to a NetworkX graph.\n",
        "\n",
        "2. Applies the hierarchical Leiden algorithm for community detection.\n",
        "\n",
        "3. Collects detailed information about each community.\n",
        "\n",
        "4. Generates summaries for each community.\n",
        "\n",
        "`generate_community_summary(text):`\n",
        "\n",
        "1. Uses LLM to generate a summary of the relationships in a community.\n",
        "2. The summary includes entity names and a synthesis of relationship descriptions.\n",
        "\n",
        "`_create_nx_graph():`\n",
        "\n",
        "1. Converts the internal graph representation to a NetworkX graph for community detection.\n",
        "\n",
        "`_collect_community_info(nx_graph, clusters):`\n",
        "\n",
        "1. Collects detailed information about each node based on its community.\n",
        "2. Creates a string representation of each relationship within a community.\n",
        "\n",
        "`_summarize_communities(community_info):`\n",
        "\n",
        "1. Generates and stores summaries for each community using LLM.\n",
        "\n",
        "`get_community_summaries():`\n",
        "\n",
        "1. Returns the community summaries by building them if not already done."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dj6VbwFyRKv4",
        "outputId": "98d07d66-4fc9-4b2a-8992-769ff1022995"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/graspologic/models/edge_swaps.py:215: NumbaDeprecationWarning: The keyword argument 'nopython=False' was supplied. From Numba 0.59.0 the default is being changed to True and use of 'nopython=False' will raise a warning as the argument will have no effect. See https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit for details.\n",
            "  _edge_swap_numba = nb.jit(_edge_swap, nopython=False)\n"
          ]
        }
      ],
      "source": [
        "import re\n",
        "from llama_index.core.graph_stores import SimplePropertyGraphStore\n",
        "import networkx as nx\n",
        "from graspologic.partition import hierarchical_leiden\n",
        "\n",
        "from llama_index.core.llms import ChatMessage\n",
        "\n",
        "\n",
        "class GraphRAGStore(SimplePropertyGraphStore):\n",
        "    community_summary = {}\n",
        "    max_cluster_size = 5\n",
        "\n",
        "    def generate_community_summary(self, text):\n",
        "        \"\"\"Generate summary for a given text using an LLM.\"\"\"\n",
        "        messages = [\n",
        "            ChatMessage(\n",
        "                role=\"system\",\n",
        "                content=(\n",
        "                    \"You are provided with a set of relationships from a knowledge graph, each represented as \"\n",
        "                    \"entity1->entity2->relation->relationship_description. Your task is to create a summary of these \"\n",
        "                    \"relationships. The summary should include the names of the entities involved and a concise synthesis \"\n",
        "                    \"of the relationship descriptions. The goal is to capture the most critical and relevant details that \"\n",
        "                    \"highlight the nature and significance of each relationship. Ensure that the summary is coherent and \"\n",
        "                    \"integrates the information in a way that emphasizes the key aspects of the relationships.\"\n",
        "                ),\n",
        "            ),\n",
        "            ChatMessage(role=\"user\", content=text),\n",
        "        ]\n",
        "        response = OpenAI().chat(messages)\n",
        "        clean_response = re.sub(r\"^assistant:\\s*\", \"\", str(response)).strip()\n",
        "        return clean_response\n",
        "\n",
        "    def build_communities(self):\n",
        "        \"\"\"Builds communities from the graph and summarizes them.\"\"\"\n",
        "        nx_graph = self._create_nx_graph()\n",
        "        community_hierarchical_clusters = hierarchical_leiden(\n",
        "            nx_graph, max_cluster_size=self.max_cluster_size\n",
        "        )\n",
        "        community_info = self._collect_community_info(\n",
        "            nx_graph, community_hierarchical_clusters\n",
        "        )\n",
        "        self._summarize_communities(community_info)\n",
        "\n",
        "    def _create_nx_graph(self):\n",
        "        \"\"\"Converts internal graph representation to NetworkX graph.\"\"\"\n",
        "        nx_graph = nx.Graph()\n",
        "        for node in self.graph.nodes.values():\n",
        "            nx_graph.add_node(str(node))\n",
        "        for relation in self.graph.relations.values():\n",
        "            nx_graph.add_edge(\n",
        "                relation.source_id,\n",
        "                relation.target_id,\n",
        "                relationship=relation.label,\n",
        "                description=relation.properties[\"relationship_description\"],\n",
        "            )\n",
        "        return nx_graph\n",
        "\n",
        "    def _collect_community_info(self, nx_graph, clusters):\n",
        "        \"\"\"Collect detailed information for each node based on their community.\"\"\"\n",
        "        community_mapping = {item.node: item.cluster for item in clusters}\n",
        "        community_info = {}\n",
        "        for item in clusters:\n",
        "            cluster_id = item.cluster\n",
        "            node = item.node\n",
        "            if cluster_id not in community_info:\n",
        "                community_info[cluster_id] = []\n",
        "\n",
        "            for neighbor in nx_graph.neighbors(node):\n",
        "                if community_mapping[neighbor] == cluster_id:\n",
        "                    edge_data = nx_graph.get_edge_data(node, neighbor)\n",
        "                    if edge_data:\n",
        "                        detail = f\"{node} -> {neighbor} -> {edge_data['relationship']} -> {edge_data['description']}\"\n",
        "                        community_info[cluster_id].append(detail)\n",
        "        return community_info\n",
        "\n",
        "    def _summarize_communities(self, community_info):\n",
        "        \"\"\"Generate and store summaries for each community.\"\"\"\n",
        "        for community_id, details in community_info.items():\n",
        "            details_text = (\n",
        "                \"\\n\".join(details) + \".\"\n",
        "            )  # Ensure it ends with a period\n",
        "            self.community_summary[\n",
        "                community_id\n",
        "            ] = self.generate_community_summary(details_text)\n",
        "\n",
        "    def get_community_summaries(self):\n",
        "        \"\"\"Returns the community summaries, building them if not already done.\"\"\"\n",
        "        if not self.community_summary:\n",
        "            self.build_communities()\n",
        "        return self.community_summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aBtey8hfRKv5"
      },
      "source": [
        "## GraphRAGQueryEngine\n",
        "\n",
        "The GraphRAGQueryEngine class is a custom query engine designed to process queries using the GraphRAG approach. It leverages the community summaries generated by the GraphRAGStore to answer user queries. Here's a breakdown of its functionality:\n",
        "\n",
        "**Main Components:**\n",
        "\n",
        "`graph_store:` An instance of GraphRAGStore, which contains the community summaries.\n",
        "`llm:` A Language Model (LLM) used for generating and aggregating answers.\n",
        "\n",
        "\n",
        "**Key Methods:**\n",
        "\n",
        "`custom_query(query_str: str)`\n",
        "\n",
        "1. This is the main entry point for processing a query. It retrieves community summaries, generates answers from each summary, and then aggregates these answers into a final response.\n",
        "\n",
        "`generate_answer_from_summary(community_summary, query):`\n",
        "\n",
        "1. Generates an answer for the query based on a single community summary.\n",
        "Uses the LLM to interpret the community summary in the context of the query.\n",
        "\n",
        "`aggregate_answers(community_answers):`\n",
        "\n",
        "1. Combines individual answers from different communities into a coherent final response.\n",
        "2. Uses the LLM to synthesize multiple perspectives into a single, concise answer.\n",
        "\n",
        "\n",
        "**Query Processing Flow:**\n",
        "\n",
        "1. Retrieve community summaries from the graph store.\n",
        "2. For each community summary, generate a specific answer to the query.\n",
        "3. Aggregate all community-specific answers into a final, coherent response.\n",
        "\n",
        "\n",
        "**Example usage:**\n",
        "\n",
        "```\n",
        "query_engine = GraphRAGQueryEngine(graph_store=graph_store, llm=llm)\n",
        "\n",
        "response = query_engine.query(\"query\")\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7FUQQd9WRKv5"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.query_engine import CustomQueryEngine\n",
        "from llama_index.core.llms import LLM\n",
        "\n",
        "\n",
        "class GraphRAGQueryEngine(CustomQueryEngine):\n",
        "    graph_store: GraphRAGStore\n",
        "    llm: LLM\n",
        "\n",
        "    def custom_query(self, query_str: str) -> str:\n",
        "        \"\"\"Process all community summaries to generate answers to a specific query.\"\"\"\n",
        "        community_summaries = self.graph_store.get_community_summaries()\n",
        "        community_answers = [\n",
        "            self.generate_answer_from_summary(community_summary, query_str)\n",
        "            for _, community_summary in community_summaries.items()\n",
        "        ]\n",
        "\n",
        "        final_answer = self.aggregate_answers(community_answers)\n",
        "        return final_answer\n",
        "\n",
        "    def generate_answer_from_summary(self, community_summary, query):\n",
        "        \"\"\"Generate an answer from a community summary based on a given query using LLM.\"\"\"\n",
        "        prompt = (\n",
        "            f\"Given the community summary: {community_summary}, \"\n",
        "            f\"how would you answer the following query? Query: {query}\"\n",
        "        )\n",
        "        messages = [\n",
        "            ChatMessage(role=\"system\", content=prompt),\n",
        "            ChatMessage(\n",
        "                role=\"user\",\n",
        "                content=\"I need an answer based on the above information.\",\n",
        "            ),\n",
        "        ]\n",
        "        response = self.llm.chat(messages)\n",
        "        cleaned_response = re.sub(r\"^assistant:\\s*\", \"\", str(response)).strip()\n",
        "        return cleaned_response\n",
        "\n",
        "    def aggregate_answers(self, community_answers):\n",
        "        \"\"\"Aggregate individual community answers into a final, coherent response.\"\"\"\n",
        "        # intermediate_text = \" \".join(community_answers)\n",
        "        prompt = \"Combine the following intermediate answers into a final, concise response.\"\n",
        "        messages = [\n",
        "            ChatMessage(role=\"system\", content=prompt),\n",
        "            ChatMessage(\n",
        "                role=\"user\",\n",
        "                content=f\"Intermediate answers: {community_answers}\",\n",
        "            ),\n",
        "        ]\n",
        "        final_response = self.llm.chat(messages)\n",
        "        cleaned_final_response = re.sub(\n",
        "            r\"^assistant:\\s*\", \"\", str(final_response)\n",
        "        ).strip()\n",
        "        return cleaned_final_response"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EexUuC0JRKv6"
      },
      "source": [
        "##  Build End to End GraphRAG Pipeline\n",
        "\n",
        "Now that we have defined all the necessary components, let’s construct the GraphRAG pipeline:\n",
        "\n",
        "1. Create nodes/chunks from the text.\n",
        "2. Build a PropertyGraphIndex using `GraphRAGExtractor` and `GraphRAGStore`.\n",
        "3. Construct communities and generate a summary for each community using the graph built above.\n",
        "4. Create a `GraphRAGQueryEngine` and begin querying."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6UhQNz1RKv6"
      },
      "source": [
        "### Create nodes/ chunks from the text."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QffRgYOwRKv7"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.node_parser import SentenceSplitter\n",
        "\n",
        "splitter = SentenceSplitter(\n",
        "    chunk_size=1024,\n",
        "    chunk_overlap=20,\n",
        ")\n",
        "nodes = splitter.get_nodes_from_documents(documents)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s5QbLyBMRKv7",
        "outputId": "9d344910-4dce-4e4f-9971-8e970ed3a2ce"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "50"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(nodes)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBjbLFysRKv8"
      },
      "source": [
        "### Build ProperGraphIndex using `GraphRAGExtractor` and `GraphRAGStore`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ciminh3RKv8"
      },
      "outputs": [],
      "source": [
        "KG_TRIPLET_EXTRACT_TMPL = \"\"\"\n",
        "-Goal-\n",
        "Given a text document, identify all entities and their entity types from the text and all relationships among the identified entities.\n",
        "Given the text, extract up to {max_knowledge_triplets} entity-relation triplets.\n",
        "\n",
        "-Steps-\n",
        "1. Identify all entities. For each identified entity, extract the following information:\n",
        "- entity_name: Name of the entity, capitalized\n",
        "- entity_type: Type of the entity\n",
        "- entity_description: Comprehensive description of the entity's attributes and activities\n",
        "Format each entity as (\"entity\"$$$$<entity_name>$$$$<entity_type>$$$$<entity_description>)\n",
        "\n",
        "2. From the entities identified in step 1, identify all pairs of (source_entity, target_entity) that are *clearly related* to each other.\n",
        "For each pair of related entities, extract the following information:\n",
        "- source_entity: name of the source entity, as identified in step 1\n",
        "- target_entity: name of the target entity, as identified in step 1\n",
        "- relation: relationship between source_entity and target_entity\n",
        "- relationship_description: explanation as to why you think the source entity and the target entity are related to each other\n",
        "\n",
        "Format each relationship as (\"relationship\"$$$$<source_entity>$$$$<target_entity>$$$$<relation>$$$$<relationship_description>)\n",
        "\n",
        "3. When finished, output.\n",
        "\n",
        "-Real Data-\n",
        "######################\n",
        "text: {text}\n",
        "######################\n",
        "output:\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c-qEDOt6RKv9"
      },
      "outputs": [],
      "source": [
        "entity_pattern = r'\\(\"entity\"\\$\\$\\$\\$\"(.+?)\"\\$\\$\\$\\$\"(.+?)\"\\$\\$\\$\\$\"(.+?)\"\\)'\n",
        "relationship_pattern = r'\\(\"relationship\"\\$\\$\\$\\$\"(.+?)\"\\$\\$\\$\\$\"(.+?)\"\\$\\$\\$\\$\"(.+?)\"\\$\\$\\$\\$\"(.+?)\"\\)'\n",
        "\n",
        "\n",
        "def parse_fn(response_str: str) -> Any:\n",
        "    entities = re.findall(entity_pattern, response_str)\n",
        "    relationships = re.findall(relationship_pattern, response_str)\n",
        "    return entities, relationships\n",
        "\n",
        "\n",
        "kg_extractor = GraphRAGExtractor(\n",
        "    llm=llm,\n",
        "    extract_prompt=KG_TRIPLET_EXTRACT_TMPL,\n",
        "    max_paths_per_chunk=2,\n",
        "    parse_fn=parse_fn,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uv86kYVBRKv9",
        "outputId": "3e933385-4399-4cce-eb26-93392bb15557"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Extracting paths from text: 100%|██████████| 50/50 [04:30<00:00,  5.41s/it]\n",
            "Generating embeddings: 100%|██████████| 1/1 [00:01<00:00,  1.24s/it]\n",
            "Generating embeddings: 100%|██████████| 4/4 [00:00<00:00,  4.22it/s]\n"
          ]
        }
      ],
      "source": [
        "from llama_index.core import PropertyGraphIndex\n",
        "\n",
        "index = PropertyGraphIndex(\n",
        "    nodes=nodes,\n",
        "    property_graph_store=GraphRAGStore(),\n",
        "    kg_extractors=[kg_extractor],\n",
        "    show_progress=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ69cyTPRKv-",
        "outputId": "a348e0c2-56c9-446d-c153-d9b9fc2688f9"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "EntityNode(label='entity', embedding=None, properties={'relationship_description': 'Gett Taxi is a competitor of Uber in the Israeli taxi market.', 'triplet_source_id': 'e4f765e3-fdfd-48d0-92a9-36f75b5865aa'}, name='Competition')"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "list(index.property_graph_store.graph.nodes.values())[-1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vHoa2cO1RKv-",
        "outputId": "fd4edfc7-c8db-4579-e8ab-0c917f68522a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Relation(label='O&G sector', source_id='Chevron', target_id='Operates in', properties={'relationship_description': 'Chevron operates in the O&G sector, as evidenced by the text mentioning that it is a company in this industry.', 'triplet_source_id': '6a28dc67-0dc0-486f-8dd6-70a3502f1c8e'})"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "list(index.property_graph_store.graph.relations.values())[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tUQkeWcZRKv_",
        "outputId": "8cbf6e83-dd8c-4bcd-e15a-1e9c9fef1421"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'Chevron operates in the O&G sector, as evidenced by the text mentioning that it is a company in this industry.'"
            ]
          },
          "execution_count": null,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "list(index.property_graph_store.graph.relations.values())[0].properties[\n",
        "    \"relationship_description\"\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "muSEoUP6RKv_"
      },
      "source": [
        "### Build communities\n",
        "\n",
        "This will create communities and summary for each community."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uLcPvNLmRKwA"
      },
      "outputs": [],
      "source": [
        "index.property_graph_store.build_communities()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xoxeehUkRKwA"
      },
      "source": [
        "### Create QueryEngine"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qR0oKUzgRKwA"
      },
      "outputs": [],
      "source": [
        "query_engine = GraphRAGQueryEngine(\n",
        "    graph_store=index.property_graph_store, llm=llm\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aJJBkgdsRKwA"
      },
      "source": [
        "### Querying"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HbQ9VVB7RKwB",
        "outputId": "1f9d7272-4d74-45b1-c77f-e860da53b6b2"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "The document discusses various news topics across different sectors. In the business sector, it mentions FirstEnergy being a publicly traded company on the New York Stock Exchange and State Street Corporation being listed on the NYSE. It also discusses Coinbase Global Inc.'s repurchase of $64.5 million worth of 0.50% convertible senior notes and the closure of the startup Protonn. In the political sphere, it highlights a theatrical act performed by Sinn Féin TD John Brady during a debate on retained firefighters. In the tech industry, it discusses the European Commission's actions against ZTE Corp. and TikTok Inc. due to security concerns. In the sports sector, it mentions Manchester United's interest in Harry Kane, the transfer of Jude Bellingham from Borussia Dortmund to Real Madrid, and the negotiation process for Maliek Collins' contract extension with the Houston Texans. In the music industry, it discusses the acquisition of The Hollies' recording catalog by BMG and the distribution pact between ADA Worldwide and Rostrum Records. In the hospitality sector, it mentions the partnership between Supplier.io and Hyatt Hotels. In the energy sector, it discusses the partnership between GE Vernova and Amplus Solar. In the gaming industry, it discusses the creation of the unannounced game \"Star Ocean: The Second Story R\" by Square Enix. In the automotive industry, it mentions the upcoming launch of the Hyundai Exter in India and Stellantis' plans to shut down the Belvidere Assembly Plant. In the airline industry, it discusses Deutsche Bank's decision to upgrade Allegiant Travel's status from Hold to Buy. In the football sector, it discusses the rejected bids made by Arsenal for Rice and the rejected bid received by Chelsea for Mason Mount. In the space industry, it mentions MDA Ltd.'s participation in the Jefferies Virtual Space Summit. In the transportation industry, it discusses Uber's strategic decision to exit the Israeli market and the emergence of Yango as a key player in the Israeli taxi market."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = query_engine.query(\n",
        "    \"What are the main news discussed in the document?\"\n",
        ")\n",
        "display(Markdown(f\"{response.response}\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oTpPfZPiRKwB",
        "outputId": "19af80df-0fd7-4d29-b84d-7e3ea48332cc"
      },
      "outputs": [
        {
          "data": {
            "text/markdown": [
              "The recent news related to the financial sector includes Morgan Stanley hiring Thomas Christl to co-head its coverage of consumer and retail clients in Europe. KeyBank has expanded its presence in the Western U.S. by opening a new branch in American Fork and donated $10,000 to the Five.12 Foundation. BMG has acquired the recording catalog of The Hollies, and Matt Pincus led a $15 million pre-growth round of investment for Soundtrack Your Brand. Hyatt Hotels and Supplier.io have been honored with the Supply & Demand Chain Executive 2023 Top Supply Chain Projects award. Bank of America Corp. reported a decline in uninsured deposits, while JPMorgan Chase & Co. reported a 1.9% increase in uninsured deposits. Coinbase Global Inc. repurchased $64.5 million worth of 0.50% convertible senior notes and also decided to repurchase its 0.50% Convertible Senior Notes due 2026 for approximately $45.5 million. Deutsche Bank upgraded Allegiant Travel's status from Hold to Buy and increased the price target to $145. Lastly, Tesla Inc.'s stock performance was analyzed by Ihor Dusaniwsky, a managing director at S3 Partners, and the company formed a significant partnership with General Motors Co. in the electric vehicle industry."
            ],
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "response = query_engine.query(\"What are news related to financial sector?\")\n",
        "display(Markdown(f\"{response.response}\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MAno-gAmRKwB"
      },
      "source": [
        "## Future Work:\n",
        "\n",
        "This cookbook is an approximate implementation of GraphRAG. In future cookbooks, we plan to extend it as follows:\n",
        "\n",
        "1. Implement retrieval using entity description embeddings.\n",
        "2. Inegrate with Neo4JPropertyGraphStore.\n",
        "3. Calculate a helpfulness score for each answer generated from the community summaries and filter out answers where the helpfulness score is zero.\n",
        "4. Perform entity disambiguation to remove duplicate entities.\n",
        "5. Implement claims or covariate information extraction, Local Search and Global Search techniques."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "llamaindex",
      "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"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}