Chainlit集成LlamaIndex和Chromadb实现RAG增强生成对话AI应用

前言

本文主要讲解如何使用LlamaIndex和Chromadb向量数据库实现RAG应用，并使用Chainlit快速搭建一个前端对话网页，实现RAG聊天问答增强的应用。文章中还讲解了LlamaIndex 的CallbackManager回调，实现案例是使用TokenCountingHandler，实现tokens使用计算回调应用。方便知道自己的tokens使用量。

LlamaIndex 回调

LlamaIndex 提供回调来帮助调试、跟踪和追溯库的内部工作。使用回调管理器，可以根据需要添加尽可能多的回调。

除了记录与事件相关的数据之外，您还可以跟踪每个事件的持续时间和发生的次数。

此外，还会记录事件的跟踪图，回调可以随意使用这些数据。例如，LlamaDebugHandler默认情况下，大多数操作后都会打印事件的跟踪。

回调事件类型虽然每个回调可能不会利用每种事件类型，但可以跟踪以下事件：

CHUNKING-> 文本分割前后的日志。
NODE_PARSING-> 文档及其解析成的节点的日志。
EMBEDDING-> 记录嵌入的文本数量。
LLM-> LLM 调用的模板和响应的日志。
QUERY-> 跟踪每个查询的开始和结束。
RETRIEVE-> 为查询检索到的节点的日志。
SYNTHESIZE-> 记录合成调用的结果。
TREE-> 生成摘要和摘要级别的日志。
SUB_QUESTION-> 记录生成的子问题和答案。

您可以实现自己的回调来跟踪这些事件，或者使用现有的回调。

Chroma简介

Chroma 是一个开源的向量数据库，专门设计用于处理大规模的向量数据，如文本嵌入、图像特征等。它提供了高效的向量搜索能力，使得应用程序能够在海量数据中快速找到与查询向量最相似的数据点。Chroma 的主要特点和优势如下：

主要特点

高性能：
- Chroma 使用先进的索引技术来加速向量搜索，即使面对非常大的数据集也能保持高效的查询速度。
- 支持多种索引方法，可以根据具体的应用场景选择最适合的索引策略。
易用性：
- 提供简单直观的API，使得开发者可以轻松地集成Chroma到自己的应用中。
- 支持多种编程语言的客户端库，方便不同技术栈的开发团队使用。
灵活性：
- Chroma 可以运行在本地或云端，支持分布式部署，能够根据需求扩展集群规模。
- 支持多种数据格式和向量表示方法，适应不同的数据处理需求。
可扩展性：
- 具有良好的水平扩展能力，可以通过增加更多的节点来提高处理能力和查询吞吐量。
- 支持数据分区和复制，确保高可用性和容错性。
社区支持：
- Chroma 拥有一个活跃的开源社区，不断贡献新的特性和改进。
- 提供详细的文档和示例代码，帮助开发者快速上手。

应用场景

Chroma 在多个领域都有广泛的应用，包括但不限于：

推荐系统：通过搜索用户行为向量，找到相似的用户或物品，从而实现个性化推荐。
搜索引擎：利用文本嵌入技术，快速检索与查询语义最接近的文档或网页。
内容过滤：检测和过滤潜在的有害内容，如垃圾邮件、恶意评论等。
图像识别：通过比较图像特征向量，实现相似图片的搜索和分类。
自然语言处理：支持文本相似度计算、情感分析等多种NLP任务。

Chroma 的设计目标是成为处理大规模向量数据的最佳选择之一，它不仅提供了强大的功能，而且易于使用，非常适合现代数据密集型应用的需求。

快速上手

创建一个文件，例如“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-vector-stores-chroma

执行以下命令安装依赖：

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
DashScope模型服务灵积文档地址 https://help.aliyun.com/zh/dashscope/

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

import os
import timeimport chainlit as cl
import chromadb
import pandas as pd
import tiktoken
from llama_index.core import (Settings,VectorStoreIndex,StorageContext, )
from llama_index.core.callbacks import TokenCountingHandler, CallbackManager
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import TextNode
from llama_index.embeddings.dashscope import DashScopeEmbedding, DashScopeTextEmbeddingModels, \DashScopeTextEmbeddingType
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from llama_index.vector_stores.chroma import ChromaVectorStoreSettings.llm = DashScope(model_name=DashScopeGenerationModels.QWEN_TURBO, api_key=os.environ["DASHSCOPE_API_KEY"], max_tokens=512,system_prompt="总是用中文回答"
)
Settings.embed_model = DashScopeEmbedding(model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
)
Settings.node_parser = SentenceSplitter(chunk_size=128, chunk_overlap=20)
Settings.num_output = 512
Settings.context_window = 6000encoding = tiktoken.get_encoding('cl100k_base')token_counter = TokenCountingHandler(tokenizer=encoding.encode,verbose=False,  # set to true to see usage printed to the console
)Settings.callback_manager = CallbackManager([token_counter])@cl.cache
def get_vector_store_index():db = chromadb.PersistentClient(path="./chroma_db")chroma_collection = db.get_or_create_collection("quickstart")vector_store = ChromaVectorStore(chroma_collection=chroma_collection, persist_dir="./chroma_db")vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store, embed_model=Settings.embed_model)print('total_embedding_token_count', token_counter.total_embedding_token_count)# reset the counts at your discretion!token_counter.reset_counts()for doc in vector_index.docstore.docs.values():print(doc)return vector_index@cl.cache
def get_vector_store_index2():def csv_to_nodes(csv_path):# 读取CSV文件df = pd.read_csv(csv_path)# 将每一行转换为Node对象text_nodes = []for (key, value) in df.values:value = str(value) if pd.notna(value) else ''text_nodes.append(TextNode(text=f"{key} {value}"))return text_nodes# 示例使用file_path = './data_file/石家庄医专客服知识库.csv'  # 替换为你的CSV文件路径nodes = csv_to_nodes(file_path)db = chromadb.PersistentClient(path="./chroma_db")chroma_collection = db.get_or_create_collection("quickstart")vector_store = ChromaVectorStore(chroma_collection=chroma_collection, persist_dir="./chroma_db")storage_context = StorageContext.from_defaults(vector_store=vector_store)vector_index = VectorStoreIndex(nodes, storage_context=storage_context, show_progress=True)return vector_indexvector_store_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()query_engine = vector_store_index.as_query_engine(similarity_top_k=5, streaming=True)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} 秒")print("Embedding Tokens: ",token_counter.total_embedding_token_count,"\n","LLM Prompt Tokens: ",token_counter.prompt_llm_token_count,"\n","LLM Completion Tokens: ",token_counter.completion_llm_token_count,"\n","Total LLM Token Count: ",token_counter.total_llm_token_count,)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()

这里我使用的国内阿里云的DashScope sdk 服务。
首次启动使用vector_store_index = get_vector_store_index2()之后改为vector_store_index = get_vector_store_index(),get_vector_store_index 是从chroma_db获取数据索引的，get_vector_store_index2 是将文本数据插入到chroma_db中。
代码还有一些不完善的地方，比如异常的处理还不完善，部署生产的时候记得完善
还可以使用tiktoken.encoding_for_model("gpt-3.5-turbo").encode 获取token编码器，但是llamaindex中只有open ai的模型和编码器的映射关系，所以这个方法只适用OPENAI。


MODEL_TO_ENCODING: dict[str, str] = {# chat"gpt-4o": "o200k_base","gpt-4": "cl100k_base","gpt-3.5-turbo": "cl100k_base","gpt-3.5": "cl100k_base",  # Common shorthand"gpt-35-turbo": "cl100k_base",  # Azure deployment name# base"davinci-002": "cl100k_base","babbage-002": "cl100k_base",# embeddings"text-embedding-ada-002": "cl100k_base","text-embedding-3-small": "cl100k_base","text-embedding-3-large": "cl100k_base",# DEPRECATED MODELS# text (DEPRECATED)"text-davinci-003": "p50k_base","text-davinci-002": "p50k_base","text-davinci-001": "r50k_base","text-curie-001": "r50k_base","text-babbage-001": "r50k_base","text-ada-001": "r50k_base","davinci": "r50k_base","curie": "r50k_base","babbage": "r50k_base","ada": "r50k_base",# code (DEPRECATED)"code-davinci-002": "p50k_base","code-davinci-001": "p50k_base","code-cushman-002": "p50k_base","code-cushman-001": "p50k_base","davinci-codex": "p50k_base","cushman-codex": "p50k_base",# edit (DEPRECATED)"text-davinci-edit-001": "p50k_edit","code-davinci-edit-001": "p50k_edit",# old embeddings (DEPRECATED)"text-similarity-davinci-001": "r50k_base","text-similarity-curie-001": "r50k_base","text-similarity-babbage-001": "r50k_base","text-similarity-ada-001": "r50k_base","text-search-davinci-doc-001": "r50k_base","text-search-curie-doc-001": "r50k_base","text-search-babbage-doc-001": "r50k_base","text-search-ada-doc-001": "r50k_base","code-search-babbage-code-001": "r50k_base","code-search-ada-code-001": "r50k_base",# open source"gpt2": "gpt2","gpt-2": "gpt2",  # Maintains consistency with gpt-4
}

代码解读

这段代码是一个使用Python编写的基于向量检索的对话系统。它结合了多个库和技术，包括chainlit用于构建对话界面，chromadb作为向量数据库存储，以及llama_index（一个用于构建语言模型应用的框架）来处理文档解析、索引创建和查询等任务。下面是对代码的主要部分进行的详细解释：

1. 导入必要的库

首先导入了代码中将要用到的各种库和模块。

2. 设置环境

使用Settings类配置了语言模型（LLM）和嵌入模型（embedding model），这里选择了阿里云的Qwen Turbo模型作为LLM，并设置了API密钥和最大输出令牌数。
嵌入模型使用的是DashScope提供的文本嵌入模型。
配置了节点解析器SentenceSplitter，用于将文档分割成适合向量化的小片段。
初始化了一个令牌计数处理器TokenCountingHandler，用于统计模型处理过程中的令牌使用情况。

3. 向量存储索引

定义了两个函数get_vector_store_index和get_vector_store_index2，这两个函数负责从不同的数据源创建向量存储索引。

get_vector_store_index函数直接从已有的Chroma数据库加载数据，创建向量存储索引。
get_vector_store_index2函数则从CSV文件读取数据，将其转换为TextNode对象列表，再通过这些节点创建向量存储索引。这适用于需要从头开始处理和索引新数据的情况。