GraphRAG Comprehensive Exploration: Principle + Integration + Cost + Applicability
This blog will introduce GraphRAG step by step from the principle, integration method, cost and applicability, etc. Considering that some readers are not professionals, this article will start from the most basic Retrieval Augmented Generation (RAG), Knowledge Graph (KG), and so on.
This article will introduce GraphRAG in three ways:
- Technical Principles of GraphRAG.
- GraphRAG Programming Examples.
- GraphRAG Enterprise Integration Thinking.
Retrieval-Augmented Generation (RAG)
Today GPT, LLama, and other large-scale language models (LLM) are very popular, they use all the world’s publicly available data in their models, and they usually perform very intelligently for known data, but there is still a part of the non-disclosable data is unknown for LLM. For newly generated data, unknown data, or questions, they don’t perform well in answering, or even talk nonsense, which is often called the illusion problem. Retrieval Augmented Generation (RAG) is the technology generated to solve this kind of problem, in simple terms, it is the unpublished data produced into a knowledge base, in the user query, the relevant content will be found in it as a context, and finally let the big model combined with the content to answer.
The RAG system performs well in the face of knowledge-intensive tasks (e.g., technical Q&A, academic question answering), reduces the phantom problem of large models in answering questions, ensures timeliness of answers, and enhances answering of long-tail questions.
RAG system building
Recall the process of answering questions in reading comprehension: in the first step, we have to read the whole text to understand the main idea of each passage and keep it in mind. The second step is to read the question, understand it, and then go to the corresponding passage in the text to find the answer according to the meaning. The third step is to combine the key information and questions in the passage to answer.
This is similar to the process of the RAG system.
The first step is to load the document, chunk the document (text chunks correspond to paragraphs), and then embed and store the chunks of text in the database (corresponding to the understanding of the content of each paragraph and remembering it in your head).
The second step is to embed the problem into the vector space and retrieve the relevant contents of the database through vector similarity. (This corresponds to the process of finding the content of the passages according to the questions).
The third step is to use the LLM to answer the question based on the retrieved content and the question, which corresponds to the process of filling in the answer at the end of the process above.
Notice that the vector similarity algorithm is used for retrieval, which is different from the complexity of the human brain retrieval based on content in reading comprehension, which is a much rougher retrieval process in RAG.
Due to flaws in the retrieval methodology, naive RAG systems usually suffer from three problems:
- Retrieval recall is limited: although vector retrieval can find semantically similar content, it may retrieve irrelevant information when faced with synonyms and strong context dependencies.
- Noise and irrelevant information interference: retrieved documents may contain a large number of content unrelated to the user query, the generation module may therefore be interfered with by the wrong or irrelevant information, resulting in a decline in the quality of the generated results.
- Lack of Contextual Understanding: Retrieval modules usually compute similarity sequentially and in blocks of text, and do not capture connections between contextual instances.
Many RAG enhancement techniques have emerged to address the above issues, such as HyDE, hybrid search, hierarchical search, etc. Among them Knowledge Graph (KG) is specifically used to solve the problem of lack of contextual understanding.
Knowledge Graph (KG)
The Knowledge Graph (KG) is a structured data model for representing and organizing knowledge. It builds a complex knowledge network by representing entities (e.g., people, places, things) and their relationships (e.g., ‘is’, ‘is located in’, ‘belongs to’) in the form of graphs.
The core of knowledge mapping is to describe the association between things through the ternary form (entity-relationship-entity), which is a structured data representation that not only captures the semantic meaning of the data but also facilitates understanding and analysis. Here is an example.
Text: Zhang San is Li Si's Mum
Triad:
(Zhang San, is mum, Li Si)
Text: Starbucks is a coffee shop established in Hangzhou in 2000.
Triad:
(Starbucks, is, a coffee shop)
(Starbucks, founded in, Hangzhou)
(Starbucks, founded, 2000)Knowledge graph extraction methods
Knowledge Graph (KG) has been performed in the past using Machine Learning and Deep Learning. Conditional Random Field, Deep Learning (LSTM, BERT, etc.) for Named Entity Recognition (NER) and Relationship Classification.
Currently, knowledge graph extraction has entered a new phase using LLM for extraction with higher extraction accuracy and more accurate relationships. Prompts on Knowledge Graph Extraction in LlamaIndex:
DEFAULT_QUERY_KEYWORD_EXTRACT_TEMPLATE_TMPL = (
"A question is provided below. Given the question, extract up to {max_keywords} "
"keywords from the text. Focus on extracting the keywords that we can use "
"to best look up answers to the question. Avoid stopwords.\n"
"---------------------\n"
"{question}\n"
"---------------------\n"
"Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'\n"
)RAG system building with knowledge graph
After understanding the knowledge graph, we can explain how the knowledge graph is applied in RAG. Compared with the vector retrieval-based RAG system, the knowledge graph replaces vector retrieval with graph retrieval, accordingly, in the stage of storing the knowledge base it is necessary to extract the entities and corresponding relationships of all documents to be stored, and to conduct the Q&A to be searched through the entities as the source of retrieval in the graph database, and to retrieve the relevant entities and textual content to be handed over to LLM as the answer.
Vector Search && Knowledge Graph Search
So what exactly is the difference between vector retrieval and knowledge graph retrieval? 🤔
- Data entry mode: vector retrieval is usually embedded into the text as a numerical matrix stored in vector databases (e.g., faiss, ChromeDB, Milvus, etc.). Knowledge graphs through entity identification, relationship extraction attribute extraction, and other technologies, build a knowledge graph, entities, and relationships are stored in a structured way, usually using graph databases (such as Neo4j, OrientDB, etc.) or relational databases.
- Retrieval method: Vector retrieval looks for similar vectors through vector similarity, cosine distance, and other methods, with poor interpretability. Knowledge graph retrieval through graph traversal, multi-pick inference, etc., and strong interpretability.
- The retrieval effect: vector search can retrieve content semantically similar to the query, provide more relevant search results, and provide high relevance retrieval results, suitable for application scenarios that require precise matching. Knowledge graph supports complex query and reasoning tasks, can handle complex problems involving multiple entities and relationships, have high factuality and accuracy, and are suitable for application scenarios that require reliable information.
Let’s use a vivid example to illustrate the difference between the two retrieval results. 😆
The content of the book ‘Journey to the West’ is taken as an example, and the book will be retrieved according to the two ways of vectorization and knowledge graph in the library.
Question: Who is the master of the Monkey King?
Vector retrieval result: searching the corresponding cut documents by vector similarity may return multiple passages related to ‘Monkey King’ in ‘Journey to the West’, including his birth, growing up, learning the art, fetching scriptures, etc. After sending it to LLM for analysis, we can’t know who his master is.
Knowledge graph result: first identify the entity ‘Sun Wukong’, search for the related entity ‘Master Bodhi’ through the graph, the relationship is the master, and send it to LLM for the answer, then it will return the answer directly without contextual information.
In practical applications, we can choose the appropriate method according to the specific needs and scenarios, or combine the advantages of both to build a more powerful search system. For example, precise entities and relationships can be obtained through knowledge graph search first, and then relevant contextual information can be obtained through vector search, to achieve a balance between global and accuracy.
GraphRAG
GraphRAG has an additional concept of communities compared to the normal knowledge graph RAG. It performs hierarchical clustering on the initial knowledge graph and assigns entities in each graph to different communities for deeper analysis.
There exist two types of retrieval in the retrieval phase of GraphRAG, global retrieval, and local retrieval.
What do we need to do in the indexing phase? 🤔
The first three steps are similar to the Knowledge Graph process, which involves splitting the text, followed by entity and relationship extraction to build the Knowledge Graph. The difference is that it will be summarised once for each text chunk. Then a clustering algorithm called Leiden is performed based on the initial knowledge graph to aggregate similar graphical relationships to form a community. A community report summary is then performed based on the entities, relationships, and corresponding summaries of that community.
While GraphRAG uses the idea of knowledge graphs, both graph and vector databases are used in the underlying implementation. Graph databases store data about entities and the relationships between them, and the vector database stores the embedded vector space corresponding to the text units.
How the stored vector and graph databases and community reports are used will be described below. Until then, let’s use a vivid example to understand what each step is doing. Take the original text of “Journey to the West” as an example:
- “The Journey to the West” is divided into text chunks.
- Take out one of the blocks for entity and relationship extraction. Let’s take the text and assume the contents:
- chunk content — regarding “The Monkey King let his horse loose in the Heavenly River when he was a Pilgrim and got into a fight with the Heavenly Puff Marshal”.
- entities & relationships — (Sun Wukong, was a colleague of Tian Peng Yuan Shuai during his time in Heaven, Tian Peng Yuan Shuai)
- summary — regarding “Summary of the relationship descriptions of the two entities Sun Wukong and Tian Peng Yuan Shuai”.
- claim(optional) — regarding “Entity-related claims with evaluation status and time-bounds”. These are emitted as a primary artifact called Covariates.
3. Data entry (vector database and graph database). Communities are then formed using a clustering algorithm on the initial mapping to form a community report. The relevant content is as follows:
- community — Communities related to the content of the Monkey King during his rebellion in Heaven
- community report — 1) how its entities are related to each other, and significant information associated with its entities. 2) Importance scores of entities such as Sun Wukong in the community and explanation of ratings of importance scores. 3) Key insight related to the Celestial Court community.
Global Search Pipeline
Global search focuses on summarising the overall information of the article with a higher density of information. The process of searching is as follows:
- The community reports were randomly divided into N batches, and each batch of reports was chunked according to a pre-defined text length, and an answer was generated for each batch using LLM, which consisted of the relevance scores of the summaries of the different text chunks to the questions asked.
- Filtered by score ranking as contextual information for LLM’s last answer.
- Final Answer Generation.
RIR’s detail process is as follows:
As hierarchical clustering allows for the division of communities into multiple tiers, community tiers are also specified when performing global searches.
Get Nau Cxj’s stories in your inbox
Join Medium for free to get updates from this writer.
Longer documents (divided into multiple tiers) can be selected according to the need to answer the quality of the lower the specified tier, the corresponding batch will be more, the answer will be more thorough and detailed, and the corresponding time will be longer. The actual use can be adjusted according to the needs.
Local Search
Local retrieval is at the lowest level and it goes through the knowledge graph to associate contents related to the entities. The process of local search is as follows:
- Entity retrieval: Embed the problem directly into a vector space and use vectors to retrieve information about the most relevant chunks. Mapping to entities via chunk information.
- Entity-Text Unit Mapping: Take out the text chunk related to the entity and delete the information of the searched query.
- Entity-Relationship Extraction: Retrieve graph data to extract neighboring entities and relationships from it.
- Entity-Community Report Mapping: Extract the report of the community to which the entity belongs
- Entity-Covariate Mapping: Extract meta-information about the text units associated with the entity, which may include custom information such as the text ID of the text block.
- Utilization of Conversation History: If provided, the system uses conversation history to better understand the user’s intent and context.
- Response Generation: Generate answers based on all of the above.
Global Search vs Local Search
Scenarios for global search usage:
- Macro-thematic or integrative analyses: usually thematically related content inquiry.
- Integrated understanding of cross-entity relationships: usually involves open-ended questions such as who has what effect on whom
- Broad questions that lack clarity on specific entities: usually asked using relationships between entities and will not address specific entity information.
Scenarios for local search usage:
- Detailed queries for specific entities: local search is appropriate when you want to delve into specific information about a particular entity, its related attributes, specific relationships with other entities, and other issues.
- Precise search within a limited scope: the scope of the query is more clearly defined, and it is desirable to obtain information within a relatively small scope related to a specific entity.
Run GraphRAG
First, we start by installing the GraphRAG environment.
# Use conda env to create python environment, python version >= 3.11
conda create -n graphrag python=3.11
conda activate graphrag
pip install graphragThen we can start initializing the project and setting up the relevant parameters and keys.
mkdir graphrag
cd graphrag
python -m graphrag.index --init --root .After success, the following will appear in the folder:
├── prompts
│ ├── claim_extraction.txt
│ ├── community_report.txt
│ ├── entity_extraction.txt
│ └── summarize_descriptions.txt
├── .env
└── settings.yamlentity_extraction.txt is used to generate entities and relationships in text units. summarize_descriptions.txt uses the result of entity_extraction.txt to generate the summary of the text unit. claim_extraction.txtacts as a standalone workflow to extract the claim from the source text, the content of the claim is related to the entity, and claim extraction is not used in the default setup. community_report.txt is used to generate the report of communities with summaries of text units.
The .env file holds model-related API_KEY. Modify the relevant configurations in the settings.yaml file, taking GPT-4o-mini and text-embedding-small as examples:
llm:
api_key: ${GRAPHRAG_API_KEY}
type: openai_chat # or azure_openai_chat
model: gpt-4o-mini
model_supports_json: true # recommended if this is available for your model.
# max_tokens: 8000
# request_timeout: 180.0
api_base: https://api.zhizengzeng.com/v1
# api_version: 2024-02-15-preview
# organization: <organization_id>
# deployment_name: <azure_model_deployment_name>
tokens_per_minute: 50_000 # set a leaky bucket throttle
requests_per_minute: 2000 # set a leaky bucket throttle
# max_retries: 10
# max_retry_wait: 10.0
# sleep_on_rate_limit_recommendation: true # whether to sleep when azure suggests wait-times
# concurrent_requests: 25 # the number of parallel inflight requests that may be made
# temperature: 0 # temperature for sampling
# top_p: 1 # top-p sampling
# n: 1 # Number of completions to generate
embeddings:
## parallelization: override the global parallelization settings for embeddings
async_mode: threaded # or asyncio
llm:
api_key: ${GRAPHRAG_API_KEY}
type: openai_embedding # or azure_openai_embedding
model: text-embedding-3-small
api_base: https://api.zhizengzeng.com/v1
# api_version: 2024-02-15-preview
# organization: <organization_id>
# deployment_name: <azure_model_deployment_name>
# tokens_per_minute: 150_000 # set a leaky bucket throttle
# requests_per_minute: 10_000 # set a leaky bucket throttle
# max_retries: 10
# max_retry_wait: 10.0
# sleep_on_rate_limit_recommendation: true # whether to sleep when azure suggests wait-times
# concurrent_requests: 25 # the number of parallel inflight requests that may be made
# batch_size: 16 # the number of documents to send in a single request
# batch_max_tokens: 8191 # the maximum number of tokens to send in a single request
# target: required # or optionalCreate an input folder and place the knowledge base document (default use of txt files) in this folder. We can use LLM to automatically adjust the content of the prompts according to the content of the document. For more information on the parameters, please refer to the official GraphRAG documentation.
python -m graphrag.prompt_tune --limit 10 --root . --output ./improved_promptsNote that to use the tuned prompts, you need to modify the four prompt-related paths in settings.yaml .
With that ready, we can start the indexing phase. This stage takes a little longer and requires patience.
python -m graphrag.index --root .After waiting for the indexing to complete, you can start using the two methods mentioned above to start asking questions.
- global search:
python -m graphrag.query --root . --method global "your question"- local search:
python -m graphrag.query --root . --method local "your question"After indexing against the complete text of the Westward Journey, questions were asked based on what the two search methods mentioned above are good at handling. The results were as expected and are not shown here due to too much content.
Visualization
All knowledge is stored in pandas-readable .parquet ending files. Implementing the visualization requires the data to be converted into a CSV file by code and imported into a graph database. Here we have used the Neo4j database.
- Neo4j Installation (need Java 17 environment)
curl -O https://dist.neo4j.org/neo4j-community-5.21.2-unix.tar.gz
tar -zxvf neo4j-community-5.21.2.tar.gz
cd neo4j-community-5.21.2
./bin/neo4j start
./bin/neo4j startGo to the URL http://localhost:7474, enter the initial account and password neo4j and you will be taken to the administration page.
- Data conversion and data import
This section references Mervin Praison’s approach, which can be found in his blog.
GraphRAG Enterprise Integration Thinking
GraphRAG integration method
- Deploy GraphRAG as a separate service with no external interaction.
Features: Out of the box, using GraphRAG comes with packages and related code for integration.
Take a demo of a milvus database integration:
import os
import pandas as pd
import tiktoken
from dotenv import load_dotenv
from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
from graphrag.query.indexer_adapters import (
read_indexer_covariates,
read_indexer_entities,
read_indexer_relationships,
read_indexer_reports,
read_indexer_text_units,
)
from graphrag.query.input.loaders.dfs import (
store_entity_semantic_embeddings,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores import MilvusVectorStore
import asyncio
# 设置API_KEY
load_dotenv()
API_KEY = os.getenv("OPENAI_API_KEY")
index_root = "graphrag"
# 取到最后一次存储的结果
output_dir = os.path.join(index_root, "output")
subdirs = [os.path.join(output_dir, d) for d in os.listdir(output_dir)]
latest_subdir = max(subdirs, key=os.path.getmtime) # Get latest output directory
INPUT_DIR = os.path.join(latest_subdir, "artifacts")
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"
COMMUNITY_LEVEL = 2
# Read entities
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet") # 读取实体表
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet") # 读取实体embedding表
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
description_embedding_store = MilvusVectorStore(
collection_name="entity_description_embeddings",
)
# description_embedding_store.connect(uri="http://localhost:19530") # For Milvus docker service
description_embedding_store.connect(uri="./milvus.db") # For Milvus Lite
# 存储Milvus向量
entity_description_embeddings = store_entity_semantic_embeddings(
entities=entities, vectorstore=description_embedding_store
)
print(f"Entity count: {len(entity_df)}")
print(entity_df.head()) # 打印前五行数据
# Read relationships
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)
print(f"Relationship count: {len(relationship_df)}")
print(relationship_df.head()) # 打印前五行数据
# Read community reports
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
print(f"Report records: {len(report_df)}")
print(report_df.head()) # 打印前五行数据
# Read text units
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)
print(f"Text unit records: {len(text_unit_df)}")
print(text_unit_df.head()) # 打印前五行数据
llm = ChatOpenAI(
api_key=API_KEY,
api_base="https://api.zhizengzeng.com/v1",
model="gpt-4o-mini",
api_type=OpenaiApiType.OpenAI,
max_retries=20,
)
token_encoder = tiktoken.get_encoding("cl100k_base")
text_embedder = OpenAIEmbedding(
api_key=API_KEY,
api_base="https://api.zhizengzeng.com/v1",
api_type=OpenaiApiType.OpenAI,
model="text-embedding-3-small",
deployment_name="text-embedding-3-small",
max_retries=20,
)
# Create local search
context_builder = LocalSearchMixedContext(
community_reports=reports,
text_units=text_units,
entities=entities,
relationships=relationships,
covariates=None, #covariates,#todo
entity_text_embeddings=description_embedding_store,
embedding_vectorstore_key=EntityVectorStoreKey.ID, # if the vectorstore uses entity title as ids, set this to EntityVectorStoreKey.TITLE
text_embedder=text_embedder,
token_encoder=token_encoder,
)
local_context_params = {
"text_unit_prop": 0.5,
"community_prop": 0.1,
"conversation_history_max_turns": 5,
"conversation_history_user_turns_only": True,
"top_k_mapped_entities": 10,
"top_k_relationships": 10,
"include_entity_rank": True,
"include_relationship_weight": True,
"include_community_rank": False,
"return_candidate_context": False,
"embedding_vectorstore_key": EntityVectorStoreKey.ID, # set this to EntityVectorStoreKey.TITLE if the vectorstore uses entity title as ids
"max_tokens": 10000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)
}
llm_params = {
"max_tokens": 3000, # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 1000=1500)
"temperature": 0.0,
}
search_engine = LocalSearch(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
response_type="multiple paragraphs", # free form text describing the response type and format, can be anything, e.g. prioritized list, single paragraph, multiple paragraphs, multiple-page report
)
question_generator = LocalQuestionGen(
llm=llm,
context_builder=context_builder,
token_encoder=token_encoder,
llm_params=llm_params,
context_builder_params=local_context_params,
)
question_history = [
"Tell me what is knowledge distillation method",
"What knowledge distillation methods are used in this article?",
]
async def local_search(query, engine):
result = await engine.asearch(query, history=None)
print(result.response)
async def question_gen(query_history, context_data):
candidate_questions = await question_generator.agenerate(
question_history=query_history, context_data=context_data, question_count=5
)
return candidate_questions.response
print("=========Local search=========")
print()
asyncio.run(local_search("Tell me what is knowledge distillation method. What knowledge distillation methods are used in this article?", search_engine))
print()
print("=========Local search with question generation=========")
print()
new_query_list = asyncio.run(question_gen(question_history, None))
for new_query in new_query_list:
print(f"New query: {new_query}")
asyncio.run(local_search(new_query, search_engine))Capability Requirements: concurrent LLM with good language capability (at least 32b, 8K length), the embedding model with good language support.
Advantages and disadvantages: The advantages of integration are relatively simple, and low coupling with other modules of the RAG system. The disadvantage is that GraphRAG itself is costly and cannot be reorganized.
- Integration using external frameworks (e.g. Langchain)
Features: Local retrieval is achieved through the Cypher language using GraphRAG’s process query graph database, and the global searcher is achieved through code customization.
Requirements: Neo4J database integration, concurrent LLM, embedding model.
Advantages and disadvantages: The internal process can be modified according to business needs and integrated into the existing RAG system, more flexible. The disadvantage is that integration requires more development workload.
The Chinese adaptation of GraphRAG can refer to https://github.com/Airmomo/graphrag-practice-chinese.
GraphRAG Cost and Inference Efficiency
The most costly and time-consuming aspect of RAG is interacting with LLM!
Indexing cost: GraphRAG is very expensive to update the knowledge base, suitable for systems with infrequent changes in the knowledge base, and the cost can be reduced to some extent by splitting the knowledge base into different domains.
Retrieval cost: If you use a local big model, it is a matter of how much the electricity cost is, and it should be cheaper than a token.
Global retrieval inference efficiency: the number of interactions in the first sorting step is linked to the specified retrieval community level, document length, and chunk cut length. The lower the specified community level, the longer the document length, and the shorter the chunk slice length, the more interactions. However, they are performed concurrently, the second time the final answer is generated for 1 time, and the final interaction time >= the interaction time for 2 times.
Local retrieval inference efficiency: final big model answer = 1 interaction
Global retrieval is generally more time-consuming and requires limited use.
Applicability
GraphRAG is more suitable for long text documents with complex text (complex relationships between entities), e.g., literary works, novels, financial domain knowledge base, and legal domain.
Ordinary RAG + Advanced RAG: For incremental knowledge bases in purely business domains, although the number of product documents may also be large, it can be improved by some techniques.
Reliability & Scalability
GraphRAG can set the number of retries in LLM and embedding models in both the Indexing and Retrieval phases to reduce the probability of errors.
GraphRAG is externally extensible and can be directly integrated into an RAG system using a call interface.




