wehub-resource-sync
a0c8464e58
Build Package / build (ubuntu-latest) (push) Failing after 1s
CodeQL / Analyze (python) (push) Failing after 1s
Core Typecheck / core-typecheck (push) Failing after 1s
Linting / lint (push) Failing after 1s
llama-dev tests / test-llama-dev (push) Failing after 1s
Publish Sub-Package to PyPI if Needed / publish_subpackage_if_needed (push) Has been skipped
Sync Docs to Developer Hub / sync-docs (push) Failing after 0s
Build Package / build (windows-latest) (push) Has been cancelled
1143 lines
44 KiB
Plaintext
1143 lines
44 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"# GraphRAG Implementation with LlamaIndex - V2\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",
|
||
"This notebook provides guidance on constructing the GraphRAG pipeline using the LlamaIndex PropertyGraph abstractions using Neo4J.\n",
|
||
"\n",
|
||
"This notebook updates the GraphRAG pipeline to v2. If you haven’t checked v1 yet, you can find it [here](https://github.com/run-llama/llama_index/blob/main/docs/examples/cookbooks/GraphRAG_v1.ipynb). Following are the updates to the existing implementation:\n",
|
||
"\n",
|
||
"1. Integrate with Neo4J Graph database.\n",
|
||
"2. Embedding based retrieval.\n",
|
||
"\n"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Installation\n",
|
||
"\n",
|
||
"`graspologic` is used to use hierarchical_leiden for building communities."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!pip install llama-index llama-index-graph-stores-neo4j graspologic numpy==1.24.4 scipy==1.12.0 future"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"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": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/html": [
|
||
"<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&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>"
|
||
],
|
||
"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": {},
|
||
"source": [
|
||
"Prepare documents as required by LlamaIndex"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"documents = [\n",
|
||
" Document(text=f\"{row['title']}: {row['text']}\")\n",
|
||
" for i, row in news.iterrows()\n",
|
||
"]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Setup API Key and LLM"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"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": {},
|
||
"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": {},
|
||
"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",
|
||
" entity_metadata = node.metadata.copy()\n",
|
||
" for entity, entity_type, description in entities:\n",
|
||
" entity_metadata[\"entity_description\"] = description\n",
|
||
" entity_node = EntityNode(\n",
|
||
" name=entity, label=entity_type, properties=entity_metadata\n",
|
||
" )\n",
|
||
" existing_nodes.append(entity_node)\n",
|
||
"\n",
|
||
" relation_metadata = node.metadata.copy()\n",
|
||
" for triple in entities_relationship:\n",
|
||
" subj, obj, rel, description = triple\n",
|
||
" relation_metadata[\"relationship_description\"] = description\n",
|
||
" rel_node = Relation(\n",
|
||
" label=rel,\n",
|
||
" source_id=subj,\n",
|
||
" target_id=obj,\n",
|
||
" properties=relation_metadata,\n",
|
||
" )\n",
|
||
"\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": {},
|
||
"source": [
|
||
"## GraphRAGStore\n",
|
||
"\n",
|
||
"The `GraphRAGStore` class is an extension of the `Neo4jPropertyGraphStore`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": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import re\n",
|
||
"import networkx as nx\n",
|
||
"from graspologic.partition import hierarchical_leiden\n",
|
||
"from collections import defaultdict\n",
|
||
"\n",
|
||
"from llama_index.core.llms import ChatMessage\n",
|
||
"from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore\n",
|
||
"\n",
|
||
"\n",
|
||
"class GraphRAGStore(Neo4jPropertyGraphStore):\n",
|
||
" community_summary = {}\n",
|
||
" entity_info = None\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",
|
||
" self.entity_info, 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",
|
||
" triplets = self.get_triplets()\n",
|
||
" for entity1, relation, entity2 in triplets:\n",
|
||
" nx_graph.add_node(entity1.name)\n",
|
||
" nx_graph.add_node(entity2.name)\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",
|
||
" \"\"\"\n",
|
||
" Collect information for each node based on their community,\n",
|
||
" allowing entities to belong to multiple clusters.\n",
|
||
" \"\"\"\n",
|
||
" entity_info = defaultdict(set)\n",
|
||
" community_info = defaultdict(list)\n",
|
||
"\n",
|
||
" for item in clusters:\n",
|
||
" node = item.node\n",
|
||
" cluster_id = item.cluster\n",
|
||
"\n",
|
||
" # Update entity_info\n",
|
||
" entity_info[node].add(cluster_id)\n",
|
||
"\n",
|
||
" for neighbor in nx_graph.neighbors(node):\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",
|
||
"\n",
|
||
" # Convert sets to lists for easier serialization if needed\n",
|
||
" entity_info = {k: list(v) for k, v in entity_info.items()}\n",
|
||
"\n",
|
||
" return dict(entity_info), dict(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": {},
|
||
"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": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from llama_index.core.query_engine import CustomQueryEngine\n",
|
||
"from llama_index.core.llms import LLM\n",
|
||
"from llama_index.core import PropertyGraphIndex\n",
|
||
"\n",
|
||
"import re\n",
|
||
"\n",
|
||
"\n",
|
||
"class GraphRAGQueryEngine(CustomQueryEngine):\n",
|
||
" graph_store: GraphRAGStore\n",
|
||
" index: PropertyGraphIndex\n",
|
||
" llm: LLM\n",
|
||
" similarity_top_k: int = 20\n",
|
||
"\n",
|
||
" def custom_query(self, query_str: str) -> str:\n",
|
||
" \"\"\"Process all community summaries to generate answers to a specific query.\"\"\"\n",
|
||
"\n",
|
||
" entities = self.get_entities(query_str, self.similarity_top_k)\n",
|
||
"\n",
|
||
" community_ids = self.retrieve_entity_communities(\n",
|
||
" self.graph_store.entity_info, entities\n",
|
||
" )\n",
|
||
" community_summaries = self.graph_store.get_community_summaries()\n",
|
||
" community_answers = [\n",
|
||
" self.generate_answer_from_summary(community_summary, query_str)\n",
|
||
" for id, community_summary in community_summaries.items()\n",
|
||
" if id in community_ids\n",
|
||
" ]\n",
|
||
"\n",
|
||
" final_answer = self.aggregate_answers(community_answers)\n",
|
||
" return final_answer\n",
|
||
"\n",
|
||
" def get_entities(self, query_str, similarity_top_k):\n",
|
||
" nodes_retrieved = self.index.as_retriever(\n",
|
||
" similarity_top_k=similarity_top_k\n",
|
||
" ).retrieve(query_str)\n",
|
||
"\n",
|
||
" enitites = set()\n",
|
||
" pattern = (\n",
|
||
" r\"^(\\w+(?:\\s+\\w+)*)\\s*->\\s*([a-zA-Z\\s]+?)\\s*->\\s*(\\w+(?:\\s+\\w+)*)$\"\n",
|
||
" )\n",
|
||
"\n",
|
||
" for node in nodes_retrieved:\n",
|
||
" matches = re.findall(\n",
|
||
" pattern, node.text, re.MULTILINE | re.IGNORECASE\n",
|
||
" )\n",
|
||
"\n",
|
||
" for match in matches:\n",
|
||
" subject = match[0]\n",
|
||
" obj = match[2]\n",
|
||
" enitites.add(subject)\n",
|
||
" enitites.add(obj)\n",
|
||
"\n",
|
||
" return list(enitites)\n",
|
||
"\n",
|
||
" def retrieve_entity_communities(self, entity_info, entities):\n",
|
||
" \"\"\"\n",
|
||
" Retrieve cluster information for given entities, allowing for multiple clusters per entity.\n",
|
||
"\n",
|
||
" Args:\n",
|
||
" entity_info (dict): Dictionary mapping entities to their cluster IDs (list).\n",
|
||
" entities (list): List of entity names to retrieve information for.\n",
|
||
"\n",
|
||
" Returns:\n",
|
||
" List of community or cluster IDs to which an entity belongs.\n",
|
||
" \"\"\"\n",
|
||
" community_ids = []\n",
|
||
"\n",
|
||
" for entity in entities:\n",
|
||
" if entity in entity_info:\n",
|
||
" community_ids.extend(entity_info[entity])\n",
|
||
"\n",
|
||
" return list(set(community_ids))\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": {},
|
||
"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": {},
|
||
"source": [
|
||
"### Create nodes/ chunks from the text."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"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": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"50"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"len(nodes)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Build ProperGraphIndex using `GraphRAGExtractor` and `GraphRAGStore`"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"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",
|
||
"\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",
|
||
"3. Output Formatting:\n",
|
||
"- Return the result in valid JSON format with two keys: 'entities' (list of entity objects) and 'relationships' (list of relationship objects).\n",
|
||
"- Exclude any text outside the JSON structure (e.g., no explanations or comments).\n",
|
||
"- If no entities or relationships are identified, return empty lists: { \"entities\": [], \"relationships\": [] }.\n",
|
||
"\n",
|
||
"-An Output Example-\n",
|
||
"{\n",
|
||
" \"entities\": [\n",
|
||
" {\n",
|
||
" \"entity_name\": \"Albert Einstein\",\n",
|
||
" \"entity_type\": \"Person\",\n",
|
||
" \"entity_description\": \"Albert Einstein was a theoretical physicist who developed the theory of relativity and made significant contributions to physics.\"\n",
|
||
" },\n",
|
||
" {\n",
|
||
" \"entity_name\": \"Theory of Relativity\",\n",
|
||
" \"entity_type\": \"Scientific Theory\",\n",
|
||
" \"entity_description\": \"A scientific theory developed by Albert Einstein, describing the laws of physics in relation to observers in different frames of reference.\"\n",
|
||
" },\n",
|
||
" {\n",
|
||
" \"entity_name\": \"Nobel Prize in Physics\",\n",
|
||
" \"entity_type\": \"Award\",\n",
|
||
" \"entity_description\": \"A prestigious international award in the field of physics, awarded annually by the Royal Swedish Academy of Sciences.\"\n",
|
||
" }\n",
|
||
" ],\n",
|
||
" \"relationships\": [\n",
|
||
" {\n",
|
||
" \"source_entity\": \"Albert Einstein\",\n",
|
||
" \"target_entity\": \"Theory of Relativity\",\n",
|
||
" \"relation\": \"developed\",\n",
|
||
" \"relationship_description\": \"Albert Einstein is the developer of the theory of relativity.\"\n",
|
||
" },\n",
|
||
" {\n",
|
||
" \"source_entity\": \"Albert Einstein\",\n",
|
||
" \"target_entity\": \"Nobel Prize in Physics\",\n",
|
||
" \"relation\": \"won\",\n",
|
||
" \"relationship_description\": \"Albert Einstein won the Nobel Prize in Physics in 1921.\"\n",
|
||
" }\n",
|
||
" ]\n",
|
||
"}\n",
|
||
"\n",
|
||
"-Real Data-\n",
|
||
"######################\n",
|
||
"text: {text}\n",
|
||
"######################\n",
|
||
"output:\"\"\""
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import json\n",
|
||
"\n",
|
||
"\n",
|
||
"def parse_fn(response_str: str) -> Any:\n",
|
||
" json_pattern = r\"\\{.*\\}\"\n",
|
||
" match = re.search(json_pattern, response_str, re.DOTALL)\n",
|
||
" entities = []\n",
|
||
" relationships = []\n",
|
||
" if not match:\n",
|
||
" return entities, relationships\n",
|
||
" json_str = match.group(0)\n",
|
||
" try:\n",
|
||
" data = json.loads(json_str)\n",
|
||
" entities = [\n",
|
||
" (\n",
|
||
" entity[\"entity_name\"],\n",
|
||
" entity[\"entity_type\"],\n",
|
||
" entity[\"entity_description\"],\n",
|
||
" )\n",
|
||
" for entity in data.get(\"entities\", [])\n",
|
||
" ]\n",
|
||
" relationships = [\n",
|
||
" (\n",
|
||
" relation[\"source_entity\"],\n",
|
||
" relation[\"target_entity\"],\n",
|
||
" relation[\"relation\"],\n",
|
||
" relation[\"relationship_description\"],\n",
|
||
" )\n",
|
||
" for relation in data.get(\"relationships\", [])\n",
|
||
" ]\n",
|
||
" return entities, relationships\n",
|
||
" except json.JSONDecodeError as e:\n",
|
||
" print(\"Error parsing JSON:\", e)\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": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Docker Setup And Neo4J setup\n",
|
||
"\n",
|
||
"To launch Neo4j locally, first ensure you have docker installed. Then, you can launch the database with the following docker command.\n",
|
||
"\n",
|
||
"```\n",
|
||
"docker run \\\n",
|
||
" -p 7474:7474 -p 7687:7687 \\\n",
|
||
" -v $PWD/data:/data -v $PWD/plugins:/plugins \\\n",
|
||
" --name neo4j-apoc \\\n",
|
||
" -e NEO4J_apoc_export_file_enabled=true \\\n",
|
||
" -e NEO4J_apoc_import_file_enabled=true \\\n",
|
||
" -e NEO4J_apoc_import_file_use__neo4j__config=true \\\n",
|
||
" -e NEO4JLABS_PLUGINS=\\[\\\"apoc\\\"\\] \\\n",
|
||
" neo4j:latest\n",
|
||
"```\n",
|
||
"From here, you can open the db at http://localhost:7474/. On this page, you will be asked to sign in. Use the default username/password of neo4j and neo4j.\n",
|
||
"\n",
|
||
"Once you login for the first time, you will be asked to change the password."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stderr",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Received notification from DBMS server: {severity: WARNING} {code: Neo.ClientNotification.Statement.FeatureDeprecationWarning} {category: DEPRECATION} {title: This feature is deprecated and will be removed in future versions.} {description: The procedure has a deprecated field. ('config' used by 'apoc.meta.graphSample' is deprecated.)} {position: line: 1, column: 1, offset: 0} for query: \"CALL apoc.meta.graphSample() YIELD nodes, relationships RETURN nodes, [rel in relationships | {name:apoc.any.property(rel, 'type'), count: apoc.any.property(rel, 'count')}] AS relationships\"\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore\n",
|
||
"\n",
|
||
"# Note: used to be `Neo4jPGStore`\n",
|
||
"graph_store = GraphRAGStore(\n",
|
||
" username=\"neo4j\", password=\"<PASSWORD>\", url=\"bolt://localhost:7687\"\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stderr",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"Extracting paths from text: 100%|██████████| 50/50 [05:45<00:00, 6.90s/it]\n",
|
||
"Generating embeddings: 100%|██████████| 1/1 [00:02<00:00, 2.59s/it]\n",
|
||
"Generating embeddings: 100%|██████████| 2/2 [00:03<00:00, 1.90s/it]\n",
|
||
"Received notification from DBMS server: {severity: WARNING} {code: Neo.ClientNotification.Statement.FeatureDeprecationWarning} {category: DEPRECATION} {title: This feature is deprecated and will be removed in future versions.} {description: The procedure has a deprecated field. ('config' used by 'apoc.meta.graphSample' is deprecated.)} {position: line: 1, column: 1, offset: 0} for query: \"CALL apoc.meta.graphSample() YIELD nodes, relationships RETURN nodes, [rel in relationships | {name:apoc.any.property(rel, 'type'), count: apoc.any.property(rel, 'count')}] AS relationships\"\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"from llama_index.core import PropertyGraphIndex\n",
|
||
"\n",
|
||
"index = PropertyGraphIndex(\n",
|
||
" nodes=nodes,\n",
|
||
" kg_extractors=[kg_extractor],\n",
|
||
" property_graph_store=graph_store,\n",
|
||
" show_progress=True,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"[EntityNode(label='Software', embedding=None, properties={'id': 'Unreal Engine', 'entity_description': \"Unreal Engine is a game engine developed by Epic. It is used in conjunction with Epic's MetaHuman Animator tool to animate hyperrealistic MetaHumans.\", 'triplet_source_id': 'b6fbbdc0-cc13-4342-a70e-b0d86f3fd2ad'}, name='MetaHuman Animator'),\n",
|
||
" Relation(label='Integrated', source_id='MetaHuman Animator', target_id='Unreal Engine', properties={'relationship_description': 'The MetaHuman Animator tool developed by Epic is integrated with the Unreal Engine. It applies the captured actor’s facial performance to a hyperrealistic “MetaHuman” in the Unreal Engine.', 'triplet_source_id': 'a6f5c123-65a8-4278-8e24-e103e767b82f'}),\n",
|
||
" EntityNode(label='Software', embedding=None, properties={'id': 'MetaHuman Animator', 'entity_description': 'MetaHuman Animator is a tool developed by Epic that captures an actor’s facial performance using a device as simple as an iPhone and applies it to a MetaHuman in the Unreal Engine. It is designed to produce results quickly and efficiently.', 'triplet_source_id': 'b6fbbdc0-cc13-4342-a70e-b0d86f3fd2ad'}, name='Unreal Engine')]"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"index.property_graph_store.get_triplets()[10]"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"{'id': 'Unreal Engine',\n",
|
||
" 'entity_description': \"Unreal Engine is a game engine developed by Epic. It is used in conjunction with Epic's MetaHuman Animator tool to animate hyperrealistic MetaHumans.\",\n",
|
||
" 'triplet_source_id': 'b6fbbdc0-cc13-4342-a70e-b0d86f3fd2ad'}"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"index.property_graph_store.get_triplets()[10][0].properties"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"{'relationship_description': 'The MetaHuman Animator tool developed by Epic is integrated with the Unreal Engine. It applies the captured actor’s facial performance to a hyperrealistic “MetaHuman” in the Unreal Engine.',\n",
|
||
" 'triplet_source_id': 'a6f5c123-65a8-4278-8e24-e103e767b82f'}"
|
||
]
|
||
},
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"index.property_graph_store.get_triplets()[10][1].properties"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Build communities\n",
|
||
"\n",
|
||
"This will create communities and summary for each community."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"index.property_graph_store.build_communities()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Create QueryEngine"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"query_engine = GraphRAGQueryEngine(\n",
|
||
" graph_store=index.property_graph_store,\n",
|
||
" llm=llm,\n",
|
||
" index=index,\n",
|
||
" similarity_top_k=10,\n",
|
||
")"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Querying"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/markdown": [
|
||
"The document discusses several key business news: FirstEnergy's earnings results, Tram Nguyen's appointment as the Global Head of Strategic and Sustainable Investments at Bank of America, Thomas Christl's hiring by Morgan Stanley to co-head its coverage of consumer and retail clients in Europe alongside Imran Ansari, and the significant impacts of the COVID-19 pandemic on Delta Air Lines and Southwest Airlines, including the suspension and reinstatement of their dividend payouts."
|
||
],
|
||
"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": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/markdown": [
|
||
"The main news in the energy sector is that GE Vernova and Amplus Solar have entered into a Supplier-Client relationship. GE Vernova has been selected by Amplus Solar to provide and install 40 units of its 2.7-132 onshore wind turbines for a 108 MW wind power project. This means that GE Vernova will be supplying the necessary equipment and services for the successful execution of the project."
|
||
],
|
||
"text/plain": [
|
||
"<IPython.core.display.Markdown object>"
|
||
]
|
||
},
|
||
"metadata": {},
|
||
"output_type": "display_data"
|
||
}
|
||
],
|
||
"source": [
|
||
"response = query_engine.query(\"What are the main news in energy sector?\")\n",
|
||
"display(Markdown(f\"{response.response}\"))"
|
||
]
|
||
}
|
||
],
|
||
"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
|
||
}
|