lightrag-comments/lightrag/lightrag.py

"""
LightRAG - 轻量级检索增强生成系统
该模块实现了一个基于图的文档检索和问答系统，支持文档的存储、检索和知识图谱构建
"""

# 导入异步IO模块，用于处理异步编程
import asyncio

# 导入操作系统接口模块，用于处理文件路径和环境变量
import os

# 从dataclasses模块导入数据类相关工具
from dataclasses import (
    asdict,     # 将数据类实例转换为字典的函数
    dataclass,  # 数据类装饰器
    field,      # 用于定义数据类字段的函数
)

# 导入日期时间处理模块
from datetime import datetime

# 从functools导入partial函数，用于创建偏函数
from functools import partial

# 从typing模块导入类型提示工具
from typing import (
    Type,  # 用于类型注解中表示类型的类型
    cast,  # 用于类型转换的函数
)

# 导入LLM相关功能
from .llm import (
    gpt_4o_mini_complete,  # GPT模型完成功能
    openai_embedding,      # OpenAI文本嵌入功能
)

# 导入核心操作功能
from .operate import (
    chunking_by_token_size,  # 文本分块
    extract_entities,        # 实体提取
    local_query,            # 本地查询
    global_query,           # 全局查询
    hybrid_query,           # 混合查询
    naive_query,            # 简单查询
)

# 导入存储实现
from .storage import (
    JsonKVStorage,         # JSON键值存储
    NanoVectorDBStorage,   # 向量数据库存储
    NetworkXStorage,       # 图数据库存储
)

from .kg.neo4j_impl import Neo4JStorage  # Neo4j图数据库实现
# 未来可能的图数据库集成
# from .kg.ArangoDB_impl import (
#     GraphStorage as ArangoDBStorage
# )

# 导入工具函数
from .utils import (
    EmbeddingFunc,           # 嵌入函数类型
    compute_mdhash_id,       # 计算MD5哈希ID
    limit_async_func_call,   # 限制异步函数调用
    convert_response_to_json, # 响应转JSON
    logger,                  # 日志记录器
    set_logger,              # 设置日志
)

# 导入基类
from .base import (
    BaseGraphStorage,     # 图存储基类
    BaseKVStorage,        # 键值存储基类
    BaseVectorStorage,    # 向量存储基类
    StorageNameSpace,     # 存储命名空间
    QueryParam,           # 查询参数
)


def always_get_an_event_loop() -> asyncio.AbstractEventLoop:
    """
    获取或创建事件循环
    如果当前线程没有事件循环，则创建一个新的
    返回值:
        asyncio.AbstractEventLoop: 事件循环实例
    """
    try:
        return asyncio.get_event_loop()
    except RuntimeError:
        logger.info("Creating a new event loop in main thread.")
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        return loop


@dataclass
class LightRAG:
    """
    轻量级检索增强生成(LightRAG)系统的主类
    实现了文档的存储、检索、知识图谱构建和问答功能
    """

    # 工作目录配置，用存储所有缓存文件
    working_dir: str = field(
        default_factory=lambda: f"./lightrag_cache_{datetime.now().strftime('%Y-%m-%d-%H:%M:%S')}"
    )

    # 知识图谱存储类型，默认使用NetworkX实现
    kg: str = field(default="NetworkXStorage")

    # 日志级别设置
    current_log_level = logger.level
    log_level: str = field(default=current_log_level)

    # 文本分块参数配置
    chunk_token_size: int = 1200  # 每个文本块的目标token数
    chunk_overlap_token_size: int = 100  # 相邻文本块的重叠token数
    tiktoken_model_name: str = "gpt-4o-mini"  # 用于计算token的模型名称

    # 实体提取参数
    entity_extract_max_gleaning: int = 1  # 最大实体提取次数
    entity_summary_to_max_tokens: int = 500  # 实体摘要的最大token数

    # 节点嵌入配置
    node_embedding_algorithm: str = "node2vec"  # 节点嵌入算法选择
    node2vec_params: dict = field(
        default_factory=lambda: {
            "dimensions": 1536,     # 嵌入向量维度
            "num_walks": 10,        # 每个节点的随机游走次数
            "walk_length": 40,      # 每次随机游走的长度
            "window_size": 2,       # 上下文窗口大小
            "iterations": 3,        # 训练迭代次数
            "random_seed": 3,       # 随机种子
        }
    )

    # 文本嵌入配置
    embedding_func: EmbeddingFunc = field(default_factory=lambda: openai_embedding)  # 默认使用OpenAI的嵌入模型
    embedding_batch_num: int = 32  # 批处理大小
    embedding_func_max_async: int = 16  # 最大并发请求数

    # 语言模型配置
    llm_model_func: callable = gpt_4o_mini_complete  # 默认使用的语言模型
    llm_model_name: str = "meta-llama/Llama-3.2-1B-Instruct"  # 模型名称
    llm_model_max_token_size: int = 32768  # 模型最大token限制
    llm_model_max_async: int = 16  # 最大并发请求数
    llm_model_kwargs: dict = field(default_factory=dict)  # 模型额外参数

    # 存储配置
    key_string_value_json_storage_cls: Type[BaseKVStorage] = JsonKVStorage  # 键值存储类
    vector_db_storage_cls: Type[BaseVectorStorage] = NanoVectorDBStorage  # 向量存储类
    vector_db_storage_cls_kwargs: dict = field(default_factory=dict)  # 向量存储额外参数
    enable_llm_cache: bool = True  # 是否启用语言模型缓存

    # 扩展配置
    addon_params: dict = field(default_factory=dict)  # 附加参数
    convert_response_to_json_func: callable = convert_response_to_json  # JSON转换函数

    def __post_init__(self):
        """
        初始化方法，在对象创建后自动调用
        负责设置日志、初始化存储系统和配置各种功能组件
        """
        # 配置日志系统
        log_file = os.path.join(self.working_dir, "lightrag.log")
        set_logger(log_file)
        logger.setLevel(self.log_level)
        logger.info(f"Logger initialized for working directory: {self.working_dir}")

        # 记录初始化参数
        _print_config = ",\n  ".join([f"{k} = {v}" for k, v in asdict(self).items()])
        logger.debug(f"LightRAG init with param:\n  {_print_config}\n")

        # 根据配置选择图存储实现类
        self.graph_storage_cls: Type[BaseGraphStorage] = self._get_storage_class()[
            self.kg
        ]

        # 确保工作目录存在
        if not os.path.exists(self.working_dir):
            logger.info(f"Creating working directory {self.working_dir}")
            os.makedirs(self.working_dir)

        # 初始化文档存储系统
        self.full_docs = self.key_string_value_json_storage_cls(
            namespace="full_docs", global_config=asdict(self)
        )

        # 初始化文本块存储系统
        self.text_chunks = self.key_string_value_json_storage_cls(
            namespace="text_chunks", global_config=asdict(self)
        )

        # 初始化语言模型响应缓存（如果启用）
        self.llm_response_cache = (
            self.key_string_value_json_storage_cls(
                namespace="llm_response_cache", global_config=asdict(self)
            )
            if self.enable_llm_cache
            else None
        )

        # 初始化实体关系图存储
        self.chunk_entity_relation_graph = self.graph_storage_cls(
            namespace="chunk_entity_relation", global_config=asdict(self)
        )

        # 配置嵌入函数的并发限制
        self.embedding_func = limit_async_func_call(self.embedding_func_max_async)(
            self.embedding_func
        )

        # 初始化向量数据库存储系统
        # 用于存储实体的向量表示
        self.entities_vdb = self.vector_db_storage_cls(
            namespace="entities",
            global_config=asdict(self),
            embedding_func=self.embedding_func,
            meta_fields={"entity_name"},
        )
        # 用于存储关系的向量表示
        self.relationships_vdb = self.vector_db_storage_cls(
            namespace="relationships",
            global_config=asdict(self),
            embedding_func=self.embedding_func,
            meta_fields={"src_id", "tgt_id"},
        )
        # 用于存储文本块的向量表示
        self.chunks_vdb = self.vector_db_storage_cls(
            namespace="chunks",
            global_config=asdict(self),
            embedding_func=self.embedding_func,
        )

        # 配置语言模型函数的并发限制和缓存
        self.llm_model_func = limit_async_func_call(self.llm_model_max_async)(
            partial(
                self.llm_model_func,
                hashing_kv=self.llm_response_cache,
                **self.llm_model_kwargs,
            )
        )

    def _get_storage_class(self) -> Type[BaseGraphStorage]:
        """
        获取图存储类的实现
        根据配置选择合适的图存储后端（Neo4J或NetworkX）
        
        返回值:
            Type[BaseGraphStorage]: 图存储类
        """
        return {
            "Neo4JStorage": Neo4JStorage,
            "NetworkXStorage": NetworkXStorage,
        }

    def insert(self, string_or_strings):
        """
        同步方式插入文档
        将字符串或字符串列表插入到系统中进行处理
        
        参数:
            string_or_strings: 单个字符串或字符串列表，表示要处理的文档内容
        """
        loop = always_get_an_event_loop()
        return loop.run_until_complete(self.ainsert(string_or_strings))

    async def ainsert(self, string_or_strings):
        """
        异步方式插入文档
        处理文档内容，包括分块、实体提取和向量化存储
        
        参数:
            string_or_strings: 单个字符串或字符串列表，表示要处理的文档内容
        """
        try:
            # 确保输入是列表形式
            if isinstance(string_or_strings, str):
                string_or_strings = [string_or_strings]

            # 为每个文档生成唯一ID并创建文档字典
            new_docs = {
                compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
                for c in string_or_strings
            }

            # 过滤掉已存在的文档
            _add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
            new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}

            if not len(new_docs):
                logger.warning("All docs are already in the storage")
                return

            logger.info(f"[New Docs] inserting {len(new_docs)} docs")

            # 处理文档分块
            inserting_chunks = {}
            for doc_key, doc in new_docs.items():
                # 对每个文档进行分块，并为每个块生成唯一ID
                chunks = {
                    compute_mdhash_id(dp["content"], prefix="chunk-"): {
                        **dp,
                        "full_doc_id": doc_key,
                    }
                    for dp in chunking_by_token_size(
                        doc["content"],
                        overlap_token_size=self.chunk_overlap_token_size,
                        max_token_size=self.chunk_token_size,
                        tiktoken_model=self.tiktoken_model_name,
                    )
                }
                inserting_chunks.update(chunks)

            # 过滤掉已存在的文本块
            _add_chunk_keys = await self.text_chunks.filter_keys(
                list(inserting_chunks.keys())
            )
            inserting_chunks = {
                k: v for k, v in inserting_chunks.items() if k in _add_chunk_keys
            }

            if not len(inserting_chunks):
                logger.warning("All chunks are already in the storage")
                return

            logger.info(f"[New Chunks] inserting {len(inserting_chunks)} chunks")

            # 将新的文本块插入向量数据库
            await self.chunks_vdb.upsert(inserting_chunks)

            # 提取实体和关系并更新知识图谱
            logger.info("[Entity Extraction]...")
            maybe_new_kg = await extract_entities(
                inserting_chunks,
                knowledge_graph_inst=self.chunk_entity_relation_graph,
                entity_vdb=self.entities_vdb,
                relationships_vdb=self.relationships_vdb,
                global_config=asdict(self),
            )

            if maybe_new_kg is None:
                logger.warning("No new entities and relationships found")
                return

            self.chunk_entity_relation_graph = maybe_new_kg

            # 更新文档和文本块存储
            await self.full_docs.upsert(new_docs)
            await self.text_chunks.upsert(inserting_chunks)
        finally:
            # 完成插入后执行清理工作
            await self._insert_done()

    async def _insert_done(self):
        """
        插入操作完成后的回调函数
        负责更新所有存储实例的索引并保存状态
        """
        tasks = []
        # 遍历所有需要执行回调的存储实例
        for storage_inst in [
            self.full_docs,              # 完整文档存储
            self.text_chunks,            # 文本块存储
            self.llm_response_cache,     # LLM响应缓存
            self.entities_vdb,           # 实体向量数据库
            self.relationships_vdb,      # 关系向量数据库
            self.chunks_vdb,             # 文本块向量<E59091><E9878F><EFBFBD>据库
            self.chunk_entity_relation_graph,  # 实体关系图
        ]:
            if storage_inst is None:
                continue
            # 将每个存储实例的回调任务添加到任务列表
            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
        # 并发执行所有回调任务
        await asyncio.gather(*tasks)

    def query(self, query: str, param: QueryParam = QueryParam()):
        """
        同步方式执行查询
        
        参数:
            query: 查询文本
            param: 查询参数配置
        返回值:
            查询结果
        """
        loop = always_get_an_event_loop()
        return loop.run_until_complete(self.aquery(query, param))

    async def aquery(self, query: str, param: QueryParam = QueryParam()):
        """
        异步方式执行查询
        支持多种查询模式：本地查询、全局查询、混合查询和简单查询
        
        参数:
            query: 查询文本
            param: 查询参数配置
        返回值:
            查询结果
        """
        # 根据查询模式选择相应的查询方法
        if param.mode == "local":
            # 本地查询：主要基于局部上下文
            response = await local_query(
                query,
                self.chunk_entity_relation_graph,
                self.entities_vdb,
                self.relationships_vdb,
                self.text_chunks,
                param,
                asdict(self),
            )
        elif param.mode == "global":
            # 全局查询：考虑整个知识图谱
            response = await global_query(
                query,
                self.chunk_entity_relation_graph,
                self.entities_vdb,
                self.relationships_vdb,
                self.text_chunks,
                param,
                asdict(self),
            )
        elif param.mode == "hybrid":
            # 混合查询：结合局部和全局信息
            response = await hybrid_query(
                query,
                self.chunk_entity_relation_graph,
                self.entities_vdb,
                self.relationships_vdb,
                self.text_chunks,
                param,
                asdict(self),
            )
        elif param.mode == "naive":
            # 简单查询：直接基于文本相似度
            response = await naive_query(
                query,
                self.chunks_vdb,
                self.text_chunks,
                param,
                asdict(self),
            )
        else:
            raise ValueError(f"Unknown mode {param.mode}")
            
        # 执行查询完成后的清理工作
        await self._query_done()
        return response

    async def _query_done(self):
        """
        查询操作完成后的回调函数
        主要用于更新LLM响应缓存的状态
        """
        tasks = []
        # 目前只需要处理LLM响应缓存的回调
        for storage_inst in [self.llm_response_cache]:
            if storage_inst is None:
                continue
            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
        # 并发执行所有回调任务
        await asyncio.gather(*tasks)

    def delete_by_entity(self, entity_name: str):
        """
        同步方式删除指定实体
        
        参数:
            entity_name: 要删除的实体名称
        """
        loop = always_get_an_event_loop()
        return loop.run_until_complete(self.adelete_by_entity(entity_name))

    async def adelete_by_entity(self, entity_name: str):
        """
        异步方式删除指定实体及其相关的所有信息
        
        参数:
            entity_name: 要删除的实体名称
        """
        # 标准化实体名称（转为大写并添加引号）
        entity_name = f'"{entity_name.upper()}"'
        try:
            # 依次删除实体在各个存储中的数据：
            # 1. 从实体向量数据库中删除
            await self.entities_vdb.delete_entity(entity_name)
            # 2. 从关系向量数据库中删除相关关系
            await self.relationships_vdb.delete_relation(entity_name)
            # 3. 从知识图谱中删除节点
            await self.chunk_entity_relation_graph.delete_node(entity_name)
            
            # 记录删除成功的日志
            logger.info(
                f"Entity '{entity_name}' and its relationships have been deleted."
            )
            # 执行删除完成后的清理工作
            await self._delete_by_entity_done()
        except Exception as e:
            # 记录删除过程中的错误
            logger.error(f"Error while deleting entity '{entity_name}': {e}")

    async def _delete_by_entity_done(self):
        """
        实体删除操作完成后的回调函数
        负责更新所有相关存储实例的状态
        """
        tasks = []
        # 遍历需要更新的存储实例：
        # - 实体向量数据库
        # - 关系向量数据库
        # - 实体关系图
        for storage_inst in [
            self.entities_vdb,
            self.relationships_vdb,
            self.chunk_entity_relation_graph,
        ]:
            if storage_inst is None:
                continue
            # 将每个存储实例的回调添加到任务列表
            tasks.append(cast(StorageNameSpace, storage_inst).index_done_callback())
        # 并发执行所有回调任务
        await asyncio.gather(*tasks)