在 Retrieval Augmented Generation (RAG) 技术中,检索是直接影响生成输出质量的关键步骤。然而,基础 RAG 中的向量检索技术通常不足以满足所有情况。例如,传统的检索方法在处理大型私有文档存储库时往往表现不佳。许多研究团队已经开始将知识图谱整合到 RAG 中,来提高检索准确性,并取得了可喜的结果。今天,我们将探讨知识图谱的原理及其在 RAG 中的应用。
一、什么是知识图谱?
知识图谱是使用图结构来表示实体及其在现实世界中的关系并对其进行建模的一种技术方法。它将信息组织为节点(实体)和边(关系),形成一个有机网络,可以有效地存储、查询和分析复杂的知识。知识图谱的核心在于它使用三元组 (entity-relationship-entity) 来描述实体之间的关联。这种结构化数据表示形式捕获语义含义,便于理解和分析。
1.1 整体流程
在 RAG 中使用知识图谱一般流程如下图所示:
在数据摄取阶段,文档将被分块,并且知识图谱 RAG 系统将从这些块中提取实体和关系。这些实体和关系通常存储在图数据库中。
在检索过程中,知识图谱 RAG 从查询中提取实体,并从图数据库中检索相关实体和关系。检索结果通常会形成一个庞大的实体关系网络,然后将其与查询相结合,并提交给大型语言模型 (LLM) 以生成答案。
知识图谱 RAG 的一些实现结合了图检索和向量检索,利用两者的优势来提高检索的准确性和效率。
1.2 知识图谱解决RAG难点
在 RAG 中使用知识图谱主要解决了大型文档存储库的问答和理解挑战,尤其是基于 RAG 方法难以解决的全球性问题。Base RAG 在回答跨整个文档存储库的问题时通常表现不佳,例如“告诉我有关 XXX 的所有信息”。此类问题的上下文可能分散在大型存储库中,这使得通常使用 top-k 算法的矢量检索方法难以捕获所有相关的文档块,从而导致信息检索不完整。
此外,还有 LLMs上下文窗口问题。全局问题通常涉及大量上下文文档,将它们全部提交给 LLM 很容易超过其窗口限制。通过将文档提取到实体和关系中,知识图谱可以显著压缩文档块,从而可以将所有相关文档提交到LLM。
1.3 知识图谱RAG与Base RAG区别
- 知识图谱 RAG 使用图形结构来表示和存储信息,从而捕获实体之间的复杂关系,而Base RAG 通常使用矢量化文本数据。
- 知识图谱 RAG 通过图遍历和子图搜索来检索信息,而Base RAG 依赖于向量相似性搜索。
- 知识图谱 RAG 可以更好地理解实体之间的关系和层次结构,从而提供更丰富的上下文,而Base RAG 在处理复杂关系方面受到限制。
二、数据摄取
让我们看看知识图谱 RAG 的具体数据摄取过程。在Base RAG 中,文档块使用嵌入模型进行矢量化,并保存到矢量数据库中。相比之下,知识图谱 RAG 在摄取期间从文档块中提取实体和关系,并将其存储在图形数据库中。
传统的实体提取方法基于预定义的规则和字典、统计机器学习或深度学习。在 LLM 时代,实体提取越来越多地利用 LLMs,从而更好地理解文本语义并简化实现。
例如,在 LlamaIndex[1] 的 KnowledgeGraphIndex 类中,实体提取提示如下所示:
DEFAULT_KG_TRIPLET_EXTRACT_TMPL = (
"Some text is provided below. Given the text, extract up to "
"{max_knowledge_triplets} "
"knowledge triplets in the form of (subject, predicate, object). Avoid stopwords.\n"
"---------------------\n"
"Example:"
"Text: Alice is Bob's mother."
"Triplets:\n(Alice, is mother of, Bob)\n"
"Text: Philz is a coffee shop founded in Berkeley in 1982.\n"
"Triplets:\n"
"(Philz, is, coffee shop)\n"
"(Philz, founded in, Berkeley)\n"
"(Philz, founded in, 1982)\n"
"---------------------\n"
"Text: {text}\n"
"Triplets:\n"
)
该提示要求 LLM 从文本中提取 entities-relationships-entities 三元组。实体通常是表示文档实体的名词,而关系是表示实体之间连接的动词或介词。该提示包含一些示例,以帮助 LLM 更好地理解提取任务。
提取实体和关系后,通常存储在图形数据库中。某些知识图谱 RAG 实现可能会将此数据保存到文件中,并使用特定算法进行检索,例如 Microsoft 的 GraphRAG[2]。
图数据库是专门用于存储图形结构数据的数据库,支持图形数据的高效存储和查询。常见的图数据库包括 Neo4j 和 ArangoDB。不同的图形数据库使用不同的查询语言,例如 Cypher for Neo4j。如果您想使用 Neo4j 在 RAG 中存储知识图谱数据,那么学习一些基本的 Cypher 语法是必要的。
三、检索和生成
在了解了知识图谱 RAG 的数据摄取过程之后,我们再看一下它的检索和生成过程。Base RAG 通常对查询进行矢量化处理,并通过向量相似性搜索检索最相似的文档块,并将这些块提交到 LLM 以生成答案。相比之下,知识图谱 RAG 从查询中提取实体,并从图数据库中检索相关实体和关系。然后,这些实体和关系将提交给 LLM 以生成答案。
从查询中提取实体类似于数据摄取期间的实体提取,也是使用 LLMs。但是,只需提取查询中的实体,而无需形成三元组。下面是一个从查询中提取关键字的示例提示, LlamaIndex 的 KGTableRetriever 类中如下所示:
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 lookup answers to the question. Avoid stopwords.\n"
"---------------------\n"
"{question}\n"
"---------------------\n"
"Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'\n"
)
该提示要求 LLM 从查询中提取多个关键字,这些关键字通常是问题中的实体。提取实体后,它们将用于查询图形数据库。检索原则涉及使用图形数据库的查询语言搜索每个实体并检索相应的三元组。例如,Neo4j 图形数据库的简单 Cypher 查询可能如下所示:
MATCH (n {name: 'Alice'})-[r]-(m)
RETURN n, r, m
此查询在图数据库中搜索与实体 Alice 相关的所有实体和关系,检索所有相关的三元组。然后将检索到的数据转换为文本,形成问题的上下文,并将其提交给 LLM 以生成答案。
四、使用 LlamaIndex 实现知识图 RAG
了解了知识图谱 RAG 的原理后,我们看一下在实际项目中如何使用LlamaIndex 与 Neo4j 结合使用来实现知识图谱 RAG。
4.1 安装 Neo4j
Neo4j 是一种高性能图数据库,它将结构化数据存储在网络(图形)中,而不是传统的表中。这种设计使 Neo4j 在处理复杂的关系和连接方面具有显著的优势。Neo4j 使用 Cypher 作为其查询语言,这是一种类似于 SQL 的声明性图数据库查询语言,但专为图数据库量身定制。Cypher 语法简单直观,易于学习和用于编写复杂的图查询。Neo4j 支持多种检索方式,包括向量检索、全文检索等。
要安装 Neo4j 数据库,请使用 Docker 下载并启动 Neo4j 映像。安装命令如下:
docker run --name neo4j -d \
--publish=7474:7474 --publish=7687:7687 \
--volume=/your/host/path/neo4j-data/data:/data \
--env NEO4J_PLUGINS='["apoc"]' \
neo4j:5.21.0
- 使用 Neo4j Docker 镜像进行安装,版本 5.21.0。
- Neo4j 映像打开两个端口:7474 用于 Web 管理服务,7687 用于数据库服务。
- 将 Neo4j 数据目录 /your/host/path/neo4j-data/data 映射到主机上。
- 通过环境变量安装 Apoc 插件,确保 Python 程序可以使用用户名和密码连接到数据库。
服务启动成功后,打开浏览器并访问 http://localhost:7474 访问 Neo4j Web 管理界面,如下所示:
输入初始用户名和密码 neo4j/neo4j,然后设置新密码以访问 Neo4j 管理界面。
4.2 在 LlamaIndex 中使用 Neo4j
安装 Neo4j 数据库后,就可以在 LlamaIndex 中使用 Neo4jGraphStore 类连接到 Neo4j 数据库:
from llama_index.graph_stores.neo4j import Neo4jGraphStore
username = "neo4j"
password = "neo4j"
url = "bolt://localhost:7687"
database = "neo4j"
graph_store = Neo4jGraphStore(
username=username,
password=password,
url=url,
database=database,
)
- 创建一个使用 Neo4jGraphStore 连接到 Neo4j 数据库的存储对象,并传入用户名、密码、连接 URL 和数据库名称。
- Bolt 是 Neo4j 数据库用于客户端和服务器之间数据传输的高效二进制协议。
- Neo4j 社区版仅支持一个数据库,名称固定为 neo4j。
接下来,将文档保存到 Neo4j 数据库。在这里,我们使用 Wikipedia 中复仇者联盟电影[3]的情节作为测试文档。示例代码如下:
from llama_index.core import StorageContext, SimpleDirectoryReader KnowledgeGraphIndex
documents = SimpleDirectoryReader("./data").load_data()
storage_context = StorageContext.from_defaults(graph_store=graph_store)
index = KnowledgeGraphIndex.from_documents(
documents,
storage_context=storage_context,
max_triplets_per_chunk=2,
include_embeddings=True,
)
在文档分块、实体提取和嵌入之后,实体和关系将保存到 Neo4j 数据库中。数据摄取完成后,我们可以查看 Neo4j 数据库中的所有实体和关系,如下所示:
- 使用 SimpleDirectoryReader 加载文档数据。
- 使用 StorageContext 创建存储上下文对象,并传入图形存储对象。
- 使用 KnowledgeGraphIndex 从文档创建知识图谱索引对象。
- max_triplets_per_chunk=2 参数指定每个文档块将被提取为最多两个三元组。
- include_embeddings=True 参数表示提取的三元组将被转换为嵌入向量并保存。
- KnowledgeGraphIndex 默认使用 OpenAI 的 LLM 模型和嵌入模型进行实体提取和嵌入,需要在环境变量中设置 OpenAI API key。
最后,我们构建一个查询引擎并对查询执行检索和生成:
query_engine = index.as_query_engine(
include_text=True,
response_mode="tree_summarize",
embedding_mode="hybrid",
similarity_top_k=5,
verbose=True,
)
response = query_engine.query("Which two members of the Avengers created Ultron?")
print(f"Response: {response}")
- 使用 index.as_query_engine 创建查询引擎对象。
- response_mode="tree_summarize" 参数表示最终结果将以 tree-summary 模式生成。
- embedding_mode=“hybrid” 参数指定使用结合图形检索和向量检索的混合模式。
- similarity_top_k=5 参数表示最多返回 5 个类似的文档块,而 verbose=True 参数则启用检索期间的详细日志记录。
- 使用查询引擎对查询执行检索和生成,并打印生成的答案。
运行程序后,由于我们启用了调试模式,因此会打印出详细的检索信息,如下所示:
Extracted keywords: ['Avengers', 'Ultron', 'created', 'members']
KG context:
The following are knowledge sequence in max depth 2 in the form of directed graph like:
`subject -[predicate]->, object, <-[predicate_next_hop]-, object_next_hop ...`
['CAPTURES', 'Romanoff', 'USES', "Loki's scepter to close"]
['BATTLE', 'Chitauri', 'KNOWN_AS', 'Extraterrestrial race']
['CAPTURES', 'Romanoff', 'MAKES_HER_WAY_TO', 'Generator']
......
Response: Tony Stark and Bruce Banner.
- 从查询中提取关键字,例如 Avengers、Ultron、created、members。
- 然后,根据关键字打印检索到的实体关系三元组。
- 最后,从查询和上下文中生成答案。
五、微软GraphRAG
在使用 LlamaIndex 探索了知识图谱 RAG 实现之后,我们来看一下另一个实现。最近,Microsoft 基于他们早期的研究论文[4]开源了一个名为 GraphRAG[2] 的知识图谱 RAG 实现。与典型的知识图谱 RAG 不同,GraphRAG 不使用图数据库;相反,它会将知识图谱保存到文件中,并使用特定的图搜索算法检索信息。此外,GraphRAG 利用知识图谱的模块化,将它们划分为语义相关的单元,并为每个单元生成摘要报告。在响应用户查询时,GraphRAG 会搜索相关的单元摘要,并使用它们来生成最终答案。
以下是安装和使用 GraphRAG 的方法:
5.1 GraphRAG安装
要安装 GraphRAG,建议使用源码安装方法,通过修改源码,可以更轻松地调试和更好地理解 GraphRAG 原理。
首先,下载 GraphRAG 源代码:
git clone https://github.com/microsoft/graphrag.git
cd graphrag
然后,使用 Poetry[5] 安装 GraphRAG 依赖项。有关安装说明,请参阅 Poetry 网站上的安装指南[6]。安装命令如下:
Poetry 是 Python 项目中用于依赖项管理和打包的工具。Poetry 使用 pyproject.toml 文件来管理项目的所有依赖项和元数据,使项目配置更简单、更清晰。它会自动处理依赖项之间的版本冲突,并可以生成锁定文件 (poetry.lock) 以确保在不同环境中安装相同的依赖项版本。
# Create a Python environment using conda
conda create -n graphrag python=3.10
# Switch to this Python environment
conda activate graphrag
# Build a Poetry virtual environment
poetry env use python
# Enter the environment
poetry shell
# Install dependencies according to the GraphRAG poetry.lock file
poetry install
5.2 初始化配置
安装 GraphRAG 后,通过创建测试文件夹来准备测试文档:
mkdir -p ./ragtest/input
# Download the test document as suggested in the GraphRAG official documentation, or use other txt files
curl https://www.gutenberg.org/cache/epub/24022/pg24022.txt > ./ragtest/input/book.txt
接下来,使用 GraphRAG 初始化命令创建配置文件:
poetry run poe index --init --root ./ragtest
初始化后,看到在 ./ragtest 目录中生成的两个文件:settings.yaml 和 .env。使用 GRAPHRAG_API_KEY 密钥在 .env 文件中设置 OpenAI API 密钥。settings.yaml 文件用于保存 GraphRAG 的 pipeline 配置信息。
5.3 数据摄取过程
完成配置初始化,然后使用以下命令执行 GraphRAG 数据摄取管道:
poetry run poe index --root ./ragtest
此过程非常耗时,因为 GraphRAG 会对文档执行一系列操作,包括分块、实体提取、文本嵌入和生成单元报告。以下是 GraphRAG 索引流程图:
执行后,可以在 ./ragtest/output/{timestamp}/artifacts 目录中看到生成的索引文件,通常采用 parquet 格式。检索过程将从此处读取数据。
有关 GraphRAG 数据接入的更多信息,请参阅官方文档[7]。
5.4 检索过程
摄取文档后,使用以下命令使用 GraphRAG 进行检索和生成:
poetry run poe query --root ./ragtest --method local "Which two members of the Avengers created Ultron?"
GraphRAG 支持两种检索模式:local和global。使用 --method 参数指定模式。
local模式类似于传统的知识图谱 RAG,将知识图谱和原始文档块中的相关数据相结合来生成答案。global模式搜索所有单元报告,以类似 map-reduce 的方式生成答案。以下是 GraphRAG 的local检索流程图:
测试表明,GraphRAG 提供比Base RAG 检索的质量更好,尤其是对于较大的文档数据集。但是,GraphRAG 的检索速度比Base RAG 慢,因为它需要多种检索方法。
有关 GraphRAG 检索的更多信息,请参阅官方文档[8]。
参考文献:
[1] https://www.llamaindex.ai/
[2] https://microsoft.github.io/graphrag/
[3] https://en.wikipedia.org/wiki/Avenger
[4] https://arxiv.org/pdf/2404.16130
[5] https://python-poetry.org/
[6] https://python-poetry.org/docs/#installation
[7] https://microsoft.github.io/graphrag/posts/index/1-default_dataflow/
[8] https://microsoft.github.io/graphrag/posts/query/overview/