Chainlit集成LlamaIndex实现知识库高级检索（句子窗口节点检索）

检索原理

句子窗口检索原理
通常在执行基础的RAG检索时我们会将文档按指定的块大小(chunk_size)进行切割，然后进行embedding的向量化处理后存入向量数据库中，在检索时我们会计算用户问题(question) 与文档块的相似度，并选取K个最相似的文档(context)，并将其和用户问题一起发送给LLM, 并最终由LLM来生成最终的回复(response)。那么context的质量将直接影响到response的质量，然而context的质量往往取决于文档块的大小即chunk_size, 当chunk_size较小时它与question的匹配度越高，但此时context的信息量就会相对较少，这样也会导致最终的response质量变差，而当chunk_size较大时虽然context的信息量较大，但是context与question的匹配度就会降低，这也会导致最终的response质量变差，这就是基本RAG架构的弊端所在。
而句子-窗口检索, 该方法的主要思想是首先将文档切割成更小的文档块（句子级别的块，根据结束语标点符号切割）, 当匹配到问题后，将该句子块周围的文档内容（上下两册的句子）作为context输出给LLM，当前句子的前后句子数量由参数window_size来确定。“句子-窗口检索”方法中我们将不再按chunk_size来切割文档，而是按完整的句子来切割文档即每一个句子切割成一个文档，然而如何识别出文本中的句子呢？在LlamaIndx中采样的是通过句尾的标点符号如句号(.), 问号(?), 感叹号(!)等来识别句子。

LlamaIndex官方地址 https://docs.llamaindex.ai/en/stable/

快速上手

创建一个文件，例如“chainlit_chat”

mkdir chainlit_chat

进入 chainlit_chat文件夹下，执行命令创建python 虚拟环境空间(需要提前安装好python sdk。 Chainlit 需要python>=3.8。,具体操作，由于文章长度问题就不在叙述，自行百度)，命令如下：

python -m venv .venv

这一步是避免python第三方库冲突，省事版可以跳过
.venv是创建的虚拟空间文件夹可以自定义

接下来激活你创建虚拟空间，命令如下：

#linux or mac
source .venv/bin/activate
#windows
.venv\Scripts\activate

在项目根目录下创建`requirements.txt`，内容如下：

chainlit
llama-index-core
llama-index-llms-dashscope
llama-index-embeddings-dashscope

执行以下命令安装依赖：

import os
import timeimport chainlit as cl
from llama_index.core import (Settings,VectorStoreIndex,SimpleDirectoryReader, StorageContext, load_index_from_storage, Document, )
from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels, \DashScopeTextEmbeddingType
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModelsSettings.llm = DashScope(model_name=DashScopeGenerationModels.QWEN_MAX, api_key=os.environ["DASHSCOPE_API_KEY"]
)
Settings.embed_model = DashScopeEmbedding(model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
)@cl.cache
def vector_store_index():storage_dir = "./storage_window"if os.path.exists(storage_dir):# rebuild storage contextstorage_context = StorageContext.from_defaults(persist_dir=storage_dir)# load indexsentence_index = load_index_from_storage(storage_context)else:documents = SimpleDirectoryReader("./data_file").load_data(show_progress=True)print(f"documents: {len(documents)}")for d in documents:d.text = d.text.replace('。', '. ')d.text = d.text.replace('！', '! ')d.text = d.text.replace('？', '? ')text_splitter = SentenceSplitter(chunk_size=512, chunk_overlap=20)base_nodes = text_splitter.get_nodes_from_documents(documents)print(f"base_nodes: {len(base_nodes)}")docs = []for node in base_nodes:doc = Document(text=node.text)docs.append(doc)node_parser = SentenceWindowNodeParser.from_defaults(window_size=3,window_metadata_key="window",original_text_metadata_key="original_text")nodes = node_parser.get_nodes_from_documents(docs)for node in nodes:print("------------------------------------")print(node.text)print(f"nodes: {len(nodes)}")sentence_index = VectorStoreIndex(nodes)sentence_index.storage_context.persist(persist_dir=storage_dir)return sentence_indexvector_index = vector_store_index()@cl.on_chat_start
async def start():await cl.Message(author="Assistant", content="你好! 我是泰山AI智能助手. 有什么可以帮助你的吗?").send()@cl.on_message
async def main(message: cl.Message):start_time = time.time()query_engine = vector_index.as_query_engine(streaming=True,similarity_top_k=5,node_postprocessors=[MetadataReplacementPostProcessor(target_metadata_key="window")])msg = cl.Message(content="", author="Assistant")res = await query_engine.aquery(message.content)async for token in res.response_gen:await msg.stream_token(token)print(f"代码执行时间: {time.time() - start_time} 秒")source_names = []for idx, node_with_score in enumerate(res.source_nodes):node = node_with_score.nodesource_name = f"source_{idx}"source_names.append(source_name)msg.elements.append(cl.Text(content=node.get_text(), name=source_name, display="side"))await msg.stream_token(f"\n\n **数据来源**: {', '.join(source_names)}")await msg.send()

代码中的persist_dir=storage_dir 不设置的默认是 ./storage.
代码中的 SentenceSplitter(chunk_size=512, chunk_overlap=20) ,chunk_size是将长文档分割的文本块的大小，chunk_overlap 是和上下文本块的重合文本的大小。
代码中先用句子分割器，将文本分割大小差不多的文本块。在用句子窗口节点解析器，分割成单个句子。然后将句子的集合存入向量数据库中。（为什么要多一层分割的原因？因为有些事markdown表格数据，没有结束标点符号，分割的的块太大，很多向量模型不处理单个文本块超过2048个tokens的长度的数据）
中文的句子结束的标点符号要转换为英文的，不然llamaindex中的句子窗口节点解析器(SentenceWindowNodeParser)解析不了。

代码解读

这段代码使用了chainlit和llama_index两个Python库来创建一个基于文档的问答系统。下面是对代码段的解释：

导入必要的模块：
- os 和 time 是Python标准库的一部分，分别用于操作系统相关的功能和计时。
- chainlit 是一个用于快速构建交互式AI应用的库。
- llama_index 是一个框架，用于构建索引并进行文档检索。
配置llama_index的核心设置：
- 设置了使用的LLM（大语言模型）为DashScope的Qwen Turbo版本，并通过环境变量获取API密钥。
- 设置了嵌入模型（Embedding Model）为DashScope的文本嵌入模型，并指定了模型类型。
- 使用SentenceSplitter来分割文本节点，定义了块大小和重叠。
- 定义了输出长度和上下文窗口大小。
缓存函数get_vector_store_index()：
- 这个函数负责加载或创建一个向量存储索引。如果存储目录存在，则从该目录加载已有的索引；否则，从指定的数据文件夹读取文档并创建新的索引。
使用chainlit装饰器定义事件处理函数：
- @cl.on_chat_start 在聊天开始时发送欢迎消息。
- @cl.on_message 在接收到用户消息时触发，使用向量索引来查询相关性最高的文档，并将结果流式传输给用户。同时，显示每个答案片段的来源。
主逻辑部分：
- 创建一个流式查询引擎，设置相似度搜索的前k个结果。
- 当接收到消息时，使用查询引擎异步查询并流式传输响应到用户。
- 计算执行时间，并记录下每个源文档的名字以便后续引用。
- 将每个源文档的内容作为元素附加到消息中，并在最后告知用户数据来源。

这个程序提供了一个基于向量存储索引的问答系统的基本框架，可以用于从大量的文档中提取信息以回答用户的问题。

postprocessor组件

要实现最终的检索我们还需要创建query engine组件，但是在query engine组件中需要设置一个postprocessor组件作为其参数，而postprocessor组件可以由若干个子组件组合在一起，下面我们首先来简单介绍一下postprocessor子组件：Replacement组件，该组件的作用是用来选择(由target_metadata_key参数确定)将哪些context发送给llm, 也就是说Replacement组件会从检索到的context中挑选指定的内容发送给llm，所以它具有选择context的功能.

另外postprocessor还有一个叫rerank的子组件,它的作用是对检索到的上下文进行从新排序，从而得到一个精度更高的检索结果，最后Replacement组件会将rerank组件的排序结果发送给llm, 不过这里需要说明一下的是rerank是可选组件，它不是必须的，rerank组件的作用仅仅是为了提高检索的精度。

#创建Replacement组件
postproc = MetadataReplacementPostProcessor(target_metadata_key="window"
)#创建rerank组件
# 参考: https://huggingface.co/BAAI/bge-reranker-base
rerank = SentenceTransformerRerank(top_n=2, model="BAAI/bge-reranker-base"
)#创建查询引擎
sentence_window_engine = sentence_index.as_query_engine(similarity_top_k=6, node_postprocessors=[postproc, rerank]
)

代码中的BAAI/bge-reranker-base需要运行的在本地的重排模型，这个可选的。

这里创建的Replacement组件中我们设置了target_metadata_key参数为"window", 它的作用是当执行检索操作时会将context中的元数据的“窗口”数据发送给llm。而rerank组件中的top_n=2的作用是对检索到的多个context进行重新排序并选取精度最高前2个context。这里所谓的精度是指相似度计算的精度，所以可以认为经过rerank模型的重新排序后会得到和question相关度更高的context。