概述
本文介绍LightRAG文档插入的实现逻辑。通过本文的分析，可以对文档插入的实现过程有一个比较清晰的理解。可以根据本文的分析，进一步分析文档插入的细节。
文档查询的基本使用
下面的代码是使用LightRAG的最简单的流程。打开一个文本文件，插入到LightRAG中，然后对文档进行对话。
WORKING_DIR = "./dickens"
if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
# 构建LightRAG对象
rag = LightRAG(
working_dir=WORKING_DIR,
llm_model_func=gpt_4o_mini_complete, # Use gpt_4o_mini_complete LLM model
# llm_model_func=gpt_4o_complete # Optionally, use a stronger model
)
# 打开文档，并插入到LightRAG中
with open("./dickens/book.txt", "r", encoding="utf-8") as f:
rag.insert(f.read())
# 执行原始的本地查询
print(
rag.query("What are the top themes in this story?", param=QueryParam(mode="naive"))
)
文档插入过程实现分析
文档插入的实现函数如下：
class LightRAG:
def insert(self, string_or_strings):
loop = always_get_an_event_loop()
return loop.run_until_complete(self.ainsert(string_or_strings))
ainsert()函数实现分析
该函数总体实现逻辑：
（1）为每个文档生成唯一ID并创建文档字典，移除首尾空白
（2）检查哪些文档是最新的，只保留最新的文档，若没有最新的文档，则告警并返回
（3）对文件内容进行分块，并检查哪些分块是最新的，若没有最新的分块，则直接返回；
（4）将分块保存到分块存储字典中，并将分块向量化后保存到向量数据库中
（5）在所有分块中提取实体和关系，若没有提取到最新的实体和关系，则直接返回。
（6）将实体和实体的关系保存到图数据库中。
实现代码如下
async def ainsert(self, string_or_strings):
# 定义异步插入方法，接收字符串或字符串列表参数
update_storage = False
# 初始化存储更新标志为False
try:
if isinstance(string_or_strings, str):
string_or_strings = [string_or_strings]
# 如果输入是单个字符串，转换为列表形式
new_docs = {
compute_mdhash_id(c.strip(), prefix="doc-"): {"content": c.strip()}
for c in string_or_strings
}
# 为每个文档生成唯一ID并创建文档字典，移除首尾空白
_add_doc_keys = await self.full_docs.filter_keys(list(new_docs.keys()))
# 检查哪些文档ID是新的(未存储过的)
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
# 如果没有新文档，记录警告并返回
update_storage = True
logger.info(f"[New Docs] inserting {len(new_docs)} docs")
# 设置更新标志，记录新文档数量
inserting_chunks = {}
# 初始化分块字典
for doc_key, doc in tqdm_async(new_docs.items(), desc="Chunking documents", unit="doc"):
# 遍历每个文档，显示进度条
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,
)
}
# 对文档内容进行分块，为每个分块生成ID，并关联原始文档ID
inserting_chunks.update(chunks)
# 将当前文档的分块添加到总分块字典中
_add_chunk_keys = await self.text_chunks.filter_keys(list(inserting_chunks.keys()))
# 检查哪些分块ID是新的
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:
if update_storage:
await self._insert_done()
# 如果进行了存储更新，执行清理工作
小结
本文分析了LightRAG的文档插入过程的总体实现过程。通过本文的分析可知，LightRAG首先会对文档进行切割并计算嵌入向量，除此之外，还会从文档中提取实体和实体之间的关系。
粤ICP备14082021号
免费POC，
零成本试错
