LightRAG: Simple and Fast Retrieval-Augmented Generation
LightRAG 是一个用于处理知识图谱和向量数据库的框架,主要用于信息检索和知识管理。以下是对其核心组件、功能和流程的全面解析。
1. 核心组件
LightRAG 类:
负责框架的初始化和管理,包括存储、日志、LLM(大语言模型)等配置。
提供 query 和 aquery 方法,支持不同的查询模式(本地、全局、混合、简单)。
存储类:
JsonKVStorage: 用于存储键值对数据。
NanoVectorDBStorage: 用于存储向量数据,支持高效的相似性查询。
NetworkXStorage: 用于图形数据存储,支持图形操作。
查询参数:
QueryParam: 定义查询的参数,包括查询模式、返回类型、最大令牌数等。
2. 主要功能
关键词提取:
使用 LLM 提取查询中的关键词,以便后续构建上下文。
上下文构建:
buildlocal_query_context 和 buildglobal_query_context 函数分别用于构建本地和全局查询上下文。
combine_contexts 函数用于合并不同层次的上下文。
信息提取:
extract_entities 函数从文本块中提取实体和关系,并将其存储到知识图谱和向量数据库中。
查询执行:
local_query 和 global_query 函数分别处理本地和全局查询,返回相应的结果。
3. 查询流程
接收查询:
用户通过 query 方法提交查询请求,框架根据查询参数选择相应的查询模式。
关键词提取:
在查询过程中,首先提取查询中的关键词,以便后续构建上下文。
上下文构建:
根据提取的关键词,调用 buildlocal_query_context 或 buildglobal_query_context 函数构建上下文。
生成响应:
使用构建的上下文和 LLM 生成最终的响应。
存储和更新:
在信息提取过程中,提取的实体和关系会被存储到相应的数据库中,以便后续查询使用。
aquery 方法解析
aquery 方法是 LightRAG 类中的一个异步方法,负责处理用户的查询请求。它根据传入的查询参数选择不同的查询模式,并调用相应的查询函数。以下是对该方法及其关联代码的详细解析。
方法定义
async def aquery(self, query: str, param: QueryParam = QueryParam()):
- 参数:
query: 用户输入的查询字符串。
param: 查询参数,默认为 QueryParam() 的实例,包含查询的配置(如模式、返回类型等)。
方法流程
查询模式选择:
根据 param.mode 的值,选择不同的查询方式:
local: 调用本地查询函数 local_query。
global: 调用全局查询函数 global_query。
hybrid: 调用混合查询函数 hybrid_query。
naive: 调用简单查询函数 naive_query。
如果 param.mode 的值不在上述选项中,抛出 ValueError。
调用查询函数:
每种查询模式都会调用相应的查询函数,并传入必要的参数,如知识图谱实例、实体向量数据库、关系向量数据库、文本块数据库、查询参数等。
完成查询后处理:
调用 querydone 方法,执行查询完成后的操作(如更新缓存等)。
关联代码
以下是与 aquery 方法相关的代码片段:
1. 查询模式的实现:
local_query、global_query、hybrid_query 和 naive_query 函数分别实现了不同的查询逻辑。
查询参数类:
QueryParam 类定义了查询的参数,包括模式、返回类型、最大令牌数等。
3. 查询完成处理:
await self._query_done()
querydone 方法用于处理查询完成后的操作,可能涉及更新缓存或其他后续处理。
hybrid_query 方法解析
hybrid_query 方法是 LightRAG 框架中的一个异步函数,旨在结合本地和全局上下文来处理用户的查询。它通过提取关键词并构建相应的上下文,最终生成一个响应。以下是对该方法的详细解析:
方法定义
async def hybrid_query(
query,
knowledge_graph_inst: BaseGraphStorage,
entities_vdb: BaseVectorStorage,
relationships_vdb: BaseVectorStorage,
text_chunks_db: BaseKVStorage[TextChunkSchema],
query_param: QueryParam,
global_config: dict,
) -> str:
参数:
query: 用户输入的查询字符串。
knowledge_graph_inst: 知识图谱的实例,提供图形存储和查询功能。
entities_vdb: 存储实体的向量数据库实例。
relationships_vdb: 存储关系的向量数据库实例。
text_chunks_db: 存储文本块的键值存储实例。
query_param: 查询参数,包含查询的配置(如模式、返回类型等)。
global_config: 全局配置字典,包含模型函数和其他设置。
方法流程
初始化上下文:
low_level_context 和 high_level_context 初始化为 None,用于存储后续构建的上下文。
从 global_config 中获取用于调用 LLM 的函数 use_model_func。
关键词提取:
使用预定义的提示模板 PROMPTS["keywords_extraction"] 生成关键词提取的提示。
调用 LLM 函数 use_model_func,传入关键词提取的提示,获取结果。
解析关键词:
尝试将 LLM 返回的结果解析为 JSON 格式,提取高层关键词(high_level_keywords)和低层关键词(low_level_keywords)。
如果解析失败,进行错误处理,尝试从结果中提取关键词。
构建上下文:
如果存在低层关键词(ll_keywords),调用 buildlocal_query_context 函数构建本地查询上下文。
如果存在高层关键词(hl_keywords),调用 buildglobal_query_context 函数构建全局查询上下文。
合并上下文:
使用 combine_contexts 函数将高层和低层上下文合并为一个完整的上下文。
返回上下文或生成响应:
如果 query_param.only_need_context 为 True,则直接返回合并后的上下文。
如果上下文为 None,返回失败提示。
否则,使用合并后的上下文生成系统提示,并调用 LLM 函数生成最终响应。
处理响应:
如果生成的响应长度大于系统提示的长度,进行清理,去除多余的部分(如系统提示、用户输入等)。
hybrid 查询流程详解
关键词提取
query:自建组合的分红方式
通过提示词:
---Role---
You are a helpful assistant tasked with identifying both high-level and low-level keywords in the user's query.
---Goal---
Given the query, list both high-level and low-level keywords. High-level keywords focus on overarching concepts or themes, while low-level keywords focus on specific entities, details, or concrete terms.
---Instructions---
- Output the keywords in JSON format.
- The JSON should have two keys:
- "high_level_keywords" for overarching concepts or themes.
- "low_level_keywords" for specific entities or details.
######################
-Examples-
######################
Example 1:
Query: "How does international trade influence global economic stability?"
################
Output:
{
"high_level_keywords": ["International trade", "Global economic stability", "Economic impact"],
"low_level_keywords": ["Trade agreements", "Tariffs", "Currency exchange", "Imports", "Exports"]
}
#############################
Example 2:
Query: "What are the environmental consequences of deforestation on biodiversity?"
################
Output:
{
"high_level_keywords": ["Environmental consequences", "Deforestation", "Biodiversity loss"],
"low_level_keywords": ["Species extinction", "Habitat destruction", "Carbon emissions", "Rainforest", "Ecosystem"]
}
#############################
Example 3:
Query: "What is the role of education in reducing poverty?"
################
Output:
{
"high_level_keywords": ["Education", "Poverty reduction", "Socioeconomic development"],
"low_level_keywords": ["School access", "Literacy rates", "Job training", "Income inequality"]
}
#############################
-Real Data-
######################
Query: 自建组合的分红方式
######################
Output:
从模型获取到如下数据:
{
"high_level_keywords": ["自建组合", "分红方式"],
"low_level_keywords": ["投资策略", "收益分配", "股票组合", "财务管理"]
}
关键词检索
分别用关键词调用_build_local_query_context 和 _build_global_query_context 方法获取检索内容。
async def _build_local_query_context(
query,
knowledge_graph_inst: BaseGraphStorage, # 这是一个知识图谱的实例,提供对图形数据的存储和查询功能。
entities_vdb: BaseVectorStorage, # 这是一个存储实体的向量数据库实例,用于根据查询获取相关的实体。
text_chunks_db: BaseKVStorage[TextChunkSchema], # 这是一个存储文本块的键值存储实例,用于查找与实体相关的文本单元。
query_param: QueryParam,
):
results = await entities_vdb.query(query, top_k=query_param.top_k) # 这是从 entities_vdb 中查询得到的结果,包含与输入查询相关的实体信息。使用 top_k 参数限制返回的实体数量。
node_datas = await asyncio.gather( # 包含从知识图谱中获取的节点数据。通过 asyncio.gather 并行获取每个实体的详细信息。
*[knowledge_graph_inst.get_node(r["entity_name"]) for r in results]
)
node_degrees = await asyncio.gather( # 包含每个实体的度(即与该实体相连的边的数量)
*[knowledge_graph_inst.node_degree(r["entity_name"]) for r in results]
)
node_datas = [ # 包含每个节点的详细信息、实体名称和排名。通过 zip 函数将 results、node_datas 和 node_degrees 组合在一起。
{**n, "entity_name": k["entity_name"], "rank": d}
for k, n, d in zip(results, node_datas, node_degrees)
if n is not None
]
use_text_units = await _find_most_related_text_unit_from_entities( # 获取的与实体相关的文本单元
node_datas, query_param, text_chunks_db, knowledge_graph_inst
)
use_relations = await _find_most_related_edges_from_entities( # 获取的与实体相关的关系
node_datas, query_param, knowledge_graph_inst
)
......
_build_global_query_context 的区别
数据来源: buildlocal_query_context 主要从实体向量数据库中查询相关的实体,而 buildglobal_query_context 则从关系向量数据库中查询相关的关系。
处理的对象: buildlocal_query_context 处理的是节点(实体),而 buildglobal_query_context 处理的是边(关系)。
上下文构建: buildlocal_query_context 主要构建与特定查询相关的上下文,而 buildglobal_query_context 则构建与关键词相关的全局上下文。
返回结果: 两者都返回格式化的字符串,但内容不同。前者返回的是与查询相关的实体和文本单元,后者返回的是与关键词相关的关系和实体。
前者专注于具体的查询,而后者则关注更广泛的关键词和相关的关系。