从零学RAG0x02向量数据库

0 阅读4分钟

向量数据库

前面我们讲到文档索引化有几步:分块 ->> 向量化 ->> 存储。分块和向量化都讲过了,那么存储是怎么存储呢?我们传统软件开发的数据存储依赖后台数据库,其实RAG知识库也是存储在向量数据库。

概念

在人工智能时代,向量数据库已成为数据管理和AI模型不可或缺的一部分。向量数据库是一种专门设计用来存储和查询向量嵌入数据的数据库。

常见的向量数据库,可见下图:

image.png

Chroma

Chroma是一个开源的、嵌入(向量)数据库。

特性轻量级、易用、API简洁、完全开源。它支持内存和持久化模式,非常适合快速原型开发和中小规模生产应用。专门为AI时代设计,用于高效存储、检索由大模型生成的“向量”。

为什么是他

  • API设计极其简洁,几行代码就能搭建起一个可用的向量检索服务。这对于需要快速验证想法、构建MVP的团队来说,试错成本极低。相比需要独立部署、配置复杂的Milvus或Weaviate,Chroma的**“一键启动”**特性吸引力巨大。

  • 与LangChain、LlamaIndex等主流RAG框架的集成几乎是“开箱即用”。

  • "够用且简单" 在RAG应用初期,知识库的规模通常在百万级文档以下。Chroma的内存计算模式在此规模下性能表现足够好,同时支持持久化到磁盘。

  • 完全开源,零成本启动

Chroma基础

  1. 创建 client client = chromadb.Client()
    client = chromadb.PersistentClient(path="路径")

  2. collection 的使用 # 列出所有的 collection 列表 client.list_collections() # 创建: client.create_collection(name="collection名称", embedding_function=向量化函数)

  3. 获取collection:

collection=client.get_collection(name="my_collection", embedding_function=emb_fn)
# 没有就创建,有则获取:client.get_or_create_collection(name="collection名称")

4. 删除

client.delete_collection(name="my_collection")

RAG+LLM

Coding

现在,我们就可以将RAG+LLM结合起来写一个简单的基于RAG知识库的LLM问答程序。

1. RAG知识库准备

我们准备一个txt文档作为RAG知识库

# 1.加载文档,读取数据
with open('../Data/某Doc.txt', 'r', encoding='utf-8') as f:
    content = f.read()
print(len(content))

2. 向量数据库类VectorDBConnector

封装一个向量数据库类VectorDBConnector,负责和和数据库交互。以及分块和向量化功能。

a.初始化

```
# 初始化,传入集合名称,和向量化函数名
def __init__(self, collection_name):
# 当前配置中,数据保存在内存中,如果需要持久化到磁盘,需使用 PersistentClient创建客户端
chroma_client = chromadb.Client(Settings(allow_reset=True))

# 持久化到磁盘
# chroma_client = chromadb.PersistentClient(path="./chroma_data")
# 为了演示,实际不需要每次 reset()
# chroma_client.reset()

# 创建一个 collection
self.collection = chroma_client.get_or_create_collection(name=collection_name)

# 连接大模型的客户端
self.client = get_normal_client()
```

b. 向量化函数

 def get_embeddings(self, texts, model=ALI_TONGYI_EMBEDDING_V4):
 '''封装 OpenAI 的 Embedding 模型接口'''
 data = self.client.embeddings.create(input=texts, model=model).data
 return [x.embedding for x in data]

 # get_embeddings函数的变体版,因为各个模型对一次能处理的文本条数有限制且每个平台不一致,新增一个batch_size参数用以控制。
 def get_embeddings_batch(self,texts, model=ALI_TONGYI_EMBEDDING_V4, batch_size=10):
 all_embeddings = []
 for i in range(0, len(texts), batch_size):
     batch_text = texts[i:i + batch_size]
     data = self.client.embeddings.create(input=batch_text, model=model).data
     all_embeddings.extend([x.embedding for x in data])
 return all_embeddings

c. 向数据库添加文档和向量

 # 添加文档与向量
 def add_documents(self, documents):
 '''向 collection 中添加文档与向量'''
 embeddings = self.get_embeddings_batch(documents)

 self.collection.add(
     embeddings=embeddings,  # 每个文档的向量
     documents=documents,  # 文档的原文
     ids=[f"id{i}" for i in range(len(documents))]  # 每个文档的 id
 )
 print("self.collection.count():", self.collection.count())

d. 检索向量数据库

 # 检索向量数据库
 def search(self, query, top_k):
    # query是用户的查询,
    # top_k指查出top_k个相似高的记录
    
 results = self.collection.query(
     query_embeddings=self.get_embeddings_batch([query]),
     n_results=top_k
 )
 return results

3. add_documents

文档切分 & 向量数据库中添加文档

plitter = RecursiveCharacterTextSplitter(
   chunk_size=500,  # 分割长度
   chunk_overlap=50,  # 重叠长度 /重叠窗口大小
   separators=["\n\n", "\n", "。", "?", ",", ""],
)
texts = splitter.split_text(content)

# 向 向量数据库 中添加文档
vector_db.add_documents(texts)

4. 检索”问题A“

user_query = '问题A?'
results = vector_db.search(user_query, 5)

5. 接入LLM

a. 重新构建Prompt

 prompt = f"""
 你是一个问答机器人。
 你的任务是根据下述给定的已知信息回答用户问题。
 确保你的回复完全依据下述已知信息。不要编造答案。
 如果下述已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。

 已知信息:
 {contents}

 ----
 用户问:
 {user_query}

 请用中文回答用户问题。
 """

b. LLM调用封装

 def get_completion(prompt, model=ALI_TONGYI_TURBO_MODEL):
 '''封装 openai 接口'''
 messages = [{"role": "user", "content": prompt}]
 client = get_normal_client()
 response = client.chat.completions.create(
     model=model,
     messages=messages,
     temperature=0,  # 模型输出的随机性,0 表示随机性最小
 )
 return response.choices[0].message.content

c. 调用LLM

 response = get_completion(prompt)
 print(response)

RAG流程回顾

  • RAG知识库构建

    • /Data/某Doc.txt文档分块,向量化;
    • add_documents添加到向量数据库Vect
  • LLM

    1. 原始Prompt输入”问题A“

    2. ”问题A“ 向量化为 Vect(A)

    3. Vect(A)到向量数据库Vect检索search到topK相关块

    4. 优化提示词为”已知topK相关块文档,问问题A

    5. LLM基于优化后的Prompt给出答案