Chainlit集成LlamaIndex实现知识库高级检索（从小到大递归检索器）

泰山AI

已于 2024-09-24 09:44:48 修改

阅读量1k

点赞数 18

分类专栏：微软GraphRAG 文章标签：前端 javascript 开发语言 chainlit llm llamaindex python

于 2024-09-23 19:29:44 首次发布

本文链接：https://blog.csdn.net/weixin_40986713/article/details/142466276

版权

微软GraphRAG 专栏收录该内容

7 篇文章 2 订阅

订阅专栏

检索原理

从小到大的检索是指我们在切割文档时可以同时设置多个不同的chunk_size的颗粒度，比如我们可以同时设置chunk_size为128，256，512即按这三个不同的颗粒度对同时对所有文档都切割一遍。利用LlamaIndex中的RecursiveRetriever递归检索器实现对不同颗粒度的文本块检索。

递归检索的概念是我们不仅探索最直接相关的节点，还探索节点关系到额外的检索器/查询引擎并执行它们。例如，一个节点可以表示一个结构化表格的简洁摘要，并链接到该结构化表格之上的SQL/Pandas查询引擎。那么如果这个节点被检索出来，我们也希望查询底层的查询引擎以获得答案。

这对于具有层次关系的文档尤其有用。在这个例子中，我们浏览一篇关于亿万富翁的维基百科文章（以PDF形式），它包含文本和各种嵌入的结构化表格。我们首先为每个表格创建一个Pandas查询引擎，同时用一个节点（存储了一个指向查询引擎的链接）来表示每个表格；这个节点与其他节点一起存储在一个向量存储中，我们称之为IndexNode。

RecursiveRetriever检索器的优缺点

LlamaIndex 是一个用于创建索引并从文档中检索信息的框架。在 LlamaIndex 中，RecursiveRetriever 是一种用于从复杂的数据结构中递归地提取信息的方法。这种检索器主要用于处理嵌套数据结构的情况，例如文档中有多个子文档，而这些子文档又可能包含更深层次的子文档。

优点

深度检索：RecursiveRetriever 能够遍历整个嵌套数据结构，确保不会遗漏任何相关信息。
灵活性：对于具有层次或嵌套结构的信息，如多级目录或复杂的文档集合，RecursiveRetriever 提供了很好的灵活性。
适应性：对于那些需要在不同层级上进行搜索的场景，如知识图谱或树状结构数据，RecursiveRetriever 可以提供有效的解决方案。

缺点

性能问题：由于需要遍历整个嵌套结构，因此在处理大规模数据集时可能会遇到性能瓶颈。随着数据量的增长，检索速度可能会变慢。
资源消耗：递归操作可能会导致较高的内存使用，尤其是在处理深度嵌套的数据时。
复杂度增加：实现和维护递归逻辑可能会增加代码的复杂性，这可能会导致更高的开发和维护成本。
过拟合风险：如果递归层次过深，可能会导致检索结果过于具体化，忽略了更广泛的相关信息。

在使用 RecursiveRetriever 时，应该考虑到上述优点和缺点，并根据实际应用场景来决定是否采用这种方法。此外，还可以考虑结合其他检索技术来优化检索效果和提高效率。

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
llama-index-retrievers-bm25~=0.3.0

执行以下命令安装依赖：

pip install -r .\requirements.txt

安装后，项目根目录下会多出.chainlit 和.files文件夹和chainlit.md文件

代码创建

只使用通义千问的DashScope模型服务灵积的接口

在项目根目录下创建`.env`环境变量，配置如下：

DASHSCOPE_API_KEY="sk-api_key"

DASHSCOPE_API_KEY 是阿里dashscope的服务的APIkey，代码中使用DashScope的sdk实现，所以不需要配置base_url。默认就是阿里的base_url。
阿里模型接口地址 https://dashscope.console.aliyun.com/model

在项目根目录下创建app.py文件，代码如下：

import os
import time

import chainlit as cl
from llama_index.core import (
    Settings,
    VectorStoreIndex,
    SimpleDirectoryReader, StorageContext, load_index_from_storage, )
from llama_index.core.node_parser import SimpleNodeParser, SentenceSplitter
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.schema import IndexNode
from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels, \
    DashScopeTextEmbeddingType
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels

Settings.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 get_vector_store_index():
    storage_dir = "./storage_recursion"
    all_nodes = []
    if os.path.exists(storage_dir):
        # rebuild storage context
        storage_context = StorageContext.from_defaults(persist_dir=storage_dir)
        # load index
        vector_store_index = load_index_from_storage(storage_context)
        for node in storage_context.docstore.docs.values():
            all_nodes.append(node)
        print(f"11 all_nodes: {len(storage_context.docstore.docs)}")
        return vector_store_index
    else:
        documents = SimpleDirectoryReader("./data_file").load_data(show_progress=True)
        print(f"documents: {len(documents)}")
        node_parser = SentenceSplitter.from_defaults(chunk_size=512, chunk_overlap=20)
        base_nodes = node_parser.get_nodes_from_documents(documents)
        print(f"base_nodes: {len(base_nodes)}")
        sub_chunk_sizes = [128]
        sub_node_parsers = [
            SentenceSplitter.from_defaults(chunk_size=size, chunk_overlap=(int(size / 10))) for size in sub_chunk_sizes
        ]
        for base_node in base_nodes:
            for sub_node_parser in sub_node_parsers:
                sub_nodes = sub_node_parser.get_nodes_from_documents([base_node])
                sub_inodes = [
                    IndexNode.from_text_node(sn, base_node.node_id) for sn in sub_nodes
                ]
                all_nodes.extend(sub_inodes)
            # 添加父节点文档
            original_node = IndexNode.from_text_node(base_node, base_node.node_id)
            all_nodes.append(original_node)
        print(f"all_nodes: {len(all_nodes)}")

        vector_store_index = VectorStoreIndex(
            all_nodes
        )

        vector_store_index.storage_context.persist(persist_dir=storage_dir)
        return vector_store_index


vector_index = get_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()
    msg = cl.Message(content="", author="Assistant")
    vector_retriever = vector_index.as_retriever(similarity_top_k=10)
    all_ids = vector_index.docstore.docs
    node_ids = []
    for ids in all_ids:
        print(ids)
        node_ids.append(ids)
    all_nodes = vector_index.docstore.get_nodes(node_ids=node_ids)
    print(f"all_nodes: {len(all_nodes)}")
    all_nodes_dict = {n.node_id: n for n in all_nodes}
    recursive_retriever = RecursiveRetriever(
        "vector",
        retriever_dict={"vector": vector_retriever},
        node_dict=all_nodes_dict,
        verbose=True,
    )
    query_engine = RetrieverQueryEngine.from_args(
        retriever=recursive_retriever, streaming=True
    )
    res = await query_engine.aquery(message.content)
    print('res', type(res), res)
    async for token in res.async_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.node
        source_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.
代码中chunk_size是将长文档分割的文本块的大小，chunk_overlap 是和上下文本块的重合文本的大小。
如何想流式输出，请将代码中的print('res', type(res), res)注释掉，异步响应时，打印res，会变成同步。
本代码展示出从向量文档库里获取所有节点的方法

代码解读

这段代码展示了如何使用 LlamaIndex 框架来构建一个基于向量存储的索引，并通过 Chainlit 创建一个聊天应用。下面是对代码的逐行解读：

导入必要的模块：

import os
import time
import chainlit as cl
from llama_index.core import (
    Settings,
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage,
)
from llama_index.core.node_parser import SimpleNodeParser, SentenceSplitter
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.schema import IndexNode
from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels, \
    DashScopeTextEmbeddingType
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels

设置 LLM 和 Embedding 模型：

Settings.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,
)

这里配置了大模型（LLM）和嵌入模型（Embedding Model），使用的是来自 DashScope 的模型。

定义函数 get_vector_store_index：
```
@cl.cache
def get_vector_store_index():
    ...
```
此函数负责获取或构建向量存储索引。如果存储目录存在，则加载已有的索引；否则，从指定目录读取文档，解析节点，并构建新的索引。

定义 start 函数：

@cl.on_chat_start
async def start():
    await cl.Message(
        author="Assistant", content="你好! 我是泰山AI智能助手. 有什么可以帮助你的吗?"
    ).send()

当聊天开始时发送欢迎消息。

定义 main 函数：
```
@cl.on_message
async def main(message: cl.Message):
    ...
```
此函数处理用户输入的消息。它首先初始化一个向量检索器，然后构建一个递归检索器，最后使用这个检索器来查询用户输入，并将结果流式传输给用户。
主逻辑：
- 加载索引或创建新的索引。
- 使用向量检索器（vector_retriever）和递归检索器（recursive_retriever）。
- 构建查询引擎（query_engine），并使用异步查询方法 aquery 处理用户输入的消息。
- 将结果通过 stream_token 方法逐字流式传输给用户。
- 记录并显示数据来源。