LightRAG: Simple and Fast Retrieval-Augmented Generation

https://github.com/HKUDS/LightRAG

LightRAG 是一个用于处理知识图谱和向量数据库的框架,主要用于信息检索和知识管理。以下是对其核心组件、功能和流程的全面解析。

1. 核心组件

  • LightRAG 类:

    • 负责框架的初始化和管理,包括存储、日志、LLM(大语言模型)等配置。

    • 提供 queryaquery 方法,支持不同的查询模式(本地、全局、混合、简单)。

  • 存储类:

    • JsonKVStorage: 用于存储键值对数据。

    • NanoVectorDBStorage: 用于存储向量数据,支持高效的相似性查询。

    • NetworkXStorage: 用于图形数据存储,支持图形操作。

  • 查询参数:

    • QueryParam: 定义查询的参数,包括查询模式、返回类型、最大令牌数等。

2. 主要功能

  • 关键词提取:

    • 使用 LLM 提取查询中的关键词,以便后续构建上下文。

  • 上下文构建:

    • buildlocal_query_contextbuildglobal_query_context 函数分别用于构建本地和全局查询上下文。

    • combine_contexts 函数用于合并不同层次的上下文。

  • 信息提取:

    • extract_entities 函数从文本块中提取实体和关系,并将其存储到知识图谱和向量数据库中。

  • 查询执行:

    • local_queryglobal_query 函数分别处理本地和全局查询,返回相应的结果。

3. 查询流程

  • 接收查询:

    • 用户通过 query 方法提交查询请求,框架根据查询参数选择相应的查询模式。

  • 关键词提取:

    • 在查询过程中,首先提取查询中的关键词,以便后续构建上下文。

  • 上下文构建:

    • 根据提取的关键词,调用 buildlocal_query_contextbuildglobal_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_queryglobal_queryhybrid_querynaive_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_contexthigh_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_contextTrue,则直接返回合并后的上下文。

    • 如果上下文为 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 则构建与关键词相关的全局上下文。

返回结果: 两者都返回格式化的字符串,但内容不同。前者返回的是与查询相关的实体和文本单元,后者返回的是与关键词相关的关系和实体。

前者专注于具体的查询,而后者则关注更广泛的关键词和相关的关系。