跳到主要内容
使用 Python 和 ChromaDB 构建简易 RAG 应用 | 极客日志
Python AI 算法
使用 Python 和 ChromaDB 构建简易 RAG 应用 检索增强生成(RAG)的基本原理及实现流程。通过加载私有文档、向量化存储至向量数据库、检索相关片段并结合大模型生成回答,解决了大模型知识滞后和幻觉问题。内容涵盖 PDF 解析、ChromaDB 搭建、OpenAI Embedding 调用、Prompt 模板构建及完整代码封装。同时分析了文本分块策略、向量匹配准确性等关键挑战,并提供了相应的解决方案与优化建议,帮助开发者快速上手 RAG 应用开发。
城市逃兵 发布于 2025/2/6 更新于 2026/4/25 7 浏览0. 什么是 RAG
大语言模型(LLM)虽然强大,但也存在局限性:
LLM 的知识不是实时的,训练数据截止后无法获取新信息。
LLM 可能不知道私有的领域或业务知识。
RAG(Retrieval Augmented Generation,检索增强生成)顾名思义:通过检索的方法来增强生成模型的能力。你可以把这个过程想象成开卷考试。让 LLM 先翻书(检索知识库),再回答问题(生成内容)。
1. RAG 基本流程
[图:RAG 基本流程图]
看图就很容易理解 RAG 的流程了:
私有知识通过切分、向量化保存到向量数据库中,供后续使用。
用户提问时,将用户提问用同样的方式向量化,然后去向量数据库中检索。
检索出相似度最高的 k 个切分段落。
将检索结果和用户的提问放到 Prompt 模板中,组装成一个完整的 Prompt。
组装好的 Prompt 给大模型,让大模型生成回答。
理想状态下,大模型是完全依赖检索出的文档片段进行组织答案的。
简化一下,可以看出 RAG 有两大过程:
加载文档,生成向量数据库(离线阶段)。
查询向量数据库,询问大模型得到答案(在线阶段)。
下面我们一步步拆解,深入了解下 RAG 的流程和实现 RAG 所需的基本模块。
2. 向量数据库的生成
2.1 文档加载与分块
首先加载我们私有的知识库。这里以加载 PDF 文件为例。Python 提供了加载 PDF 的一些库,这里使用 pdfminer.six。
pip install pdfminer.six
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
class PDFFileLoader ():
def __init__ (self, file ) -> None :
self .paragraphs = self .extract_text_from_pdf(file, page_numbers=[0 , 3 ])
i = 1
for para in self .paragraphs[:3 ]:
print (f"========= 第{i} 段 ==========" )
print (para + "\n" )
i +=
( ):
.paragraphs
( ):
paragraphs = []
buffer =
full_text =
i, page_layout (extract_pages(filename)):
page_numbers i page_numbers:
element page_layout:
(element, LTTextContainer):
full_text += element.get_text() +
lines = full_text.split( )
text lines:
buffer = text.replace( , )
buffer:
paragraphs.append(buffer)
buffer =
row_count =
buffer:
paragraphs.append(buffer)
paragraphs
PDFFileLoader( )
1
def
getParagraphs
self
return
self
def
extract_text_from_pdf
self, filename, page_numbers=None
'''从 PDF 文件中(按指定页码)提取文字'''
''
''
for
in
enumerate
if
is
not
None
and
not
in
continue
for
in
if
isinstance
'\n'
'。\n'
for
in
'\n'
' '
if
''
0
if
return
"D:\GitHub\LEARN_LLM\RAG\如何向 ChatGPT 提问以获得高质量答案:提示技巧工程完全指南.pdf"
(1)我们首先定义了一个 PDFFileLoader 的类,接收一个 PDF 文件路径。然后类内部调用 extract_text_from_pdf 去解析 PDF 文件并分段。
(2)extract_text_from_pdf 中前半部分代码是利用 extract_pages 按页提取出 PDF 文件中的文字,然后组装成 full_text。
(3)extract_text_from_pdf 中后半部分代码是将 full_text 进行段落划分。
说明 :因为每个 PDF 提取出来的文字格式可能不同,有的每一行后面都带有"\n\n",有的不带有"\n\n",有的每一行中的单词都粘在一起…,各种各样,所以 PDF 文字划分和段落分割的算法都无法做到完美适应所有 PDF。本文重点不在这一部分,所以粗暴地根据"。\n"划分了段落。实际应用中这里你应该按照你的 PDF 文件去进行调试和分割,段落划分这几行代码不能直接用。
可以简单看下我为什么能如此粗暴的划分段落:通过 extract_pages 提取出来的文本如下:
'如何向 ChatGPT 提问以获得高质量答案:提示\n技巧工程完全指南\n\n介绍\n\n我很高兴欢迎您阅读我的最新书籍《The Art of Asking ChatGPT for High-Quality Answers: A complete \n\nGuide to Prompt Engineering Techniques》。本书是一本全面指南,介绍了各种提示技术,用于从\n\nChatGPT 中生成高质量的答案。\n\n我们将探讨如何使用不同的提示工程技术来实现不同的目标。ChatGPT 是一款最先进的语言模型,能够生成\n\n类似人类的文本。然而,理解如何正确地向 ChatGPT 提问以获得我们所需的高质量输出非常重要。而这正是\n本书的目的。\n\n无论您是普通人、研究人员、开发人员,还是只是想在自己的领域中将 ChatGPT 作为个人助手的人,本书都\n是为您编写的。我使用简单易懂的语言,提供实用的解释,并在每个提示技术中提供了示例和提示公式。通\n\n过本书,您将学习如何使用提示工程技术来控制 ChatGPT 的输出,并生成符合您特定需求的文本。\n\n在整本书中,我们还提供了如何结合不同的提示技术以实现更具体结果的示例。我希望您能像我写作时一\n\n样,享受阅读本书并从中获得知识。\n\n \n\n
与原文对比,大体上按"。\n"来分割能与实际段落比较接近,所以本例我就先这样干了。这实际是不能用于实际项目的:
2.2 创建向量数据库 本文以 chromadb 向量数据库为例进行实操。
2.2.1 创建过程 (1)创建一个向量数据库类。该类 add_documents 函数用来添加数据,它需要三个参数:
import chromadb
from chromadb.config import Settings
class MyVectorDBConnector :
def __init__ (self, collection_name, embedding_fn ):
chroma_client = chromadb.Client(Settings(allow_reset=True ))
chroma_client.reset()
self .collection = chroma_client.get_or_create_collection(name=collection_name)
self .embedding_fn = embedding_fn
def add_documents (self, documents ):
'''向 collection 中添加文档与向量'''
self .collection.add(
embeddings=self .embedding_fn(documents),
documents=documents,
ids=[f"id{i} " for i in range (len (documents))]
)
def search (self, query, top_n ):
'''检索向量数据库'''
results = self .collection.query(
query_embeddings=self .embedding_fn([query]),
n_results=top_n
)
return results
(2)文档的向量怎么来?可以通过 OpenAI 的 embeddings 接口计算得到:
from openai import OpenAI
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())
client = OpenAI()
def get_embeddings (texts, model="text-embedding-3-small" ):
'''封装 OpenAI 的 Embedding 模型接口'''
data = client.embeddings.create(input =texts, model=model).data
return [x.embedding for x in data]
注意 :在实际生产环境中,考虑到成本和隐私,也可以使用本地部署的 Embedding 模型(如 BGE-M3、m3e-base 等),配合 Sentence Transformers 库使用。
vector_db = MyVectorDBConnector("demo" , get_embeddings)
vector_db.add_documents(pdf_loader.getParagraphs())
user_query = "什么是角色提示?"
results = vector_db.search(user_query, 3 )
for para in results['documents' ][0 ]:
print (para + "\n\n" )
2.2.2 运行结果 (1)通过 OpenAI 的 embeddings 接口计算得到的文本向量
2.2.3 踩坑
2.2.3.1 坑一:NoneType object is not iterable
不知道这种情况为什么会导致 NoneType 的错误,可能是 OpenAI 向量化时对特殊字符进行了去除?
2.2.3.2 坑二:Number of embeddings 9 must match number of ids 10
原因:可以看下下面的代码,上面的错误指的是 embeddings 是 9 个值,而 ids 有 10 个值。这是因为在解决坑一时,将里面最后那个空的文档分块去掉了,没去生成 embeddings。
self .collection.add(
embeddings=self .embedding_fn(documents),
documents=documents,
ids=[f"id{i} " for i in range (len (documents))]
)
解决方法:保证 documents 和 embeddings 的数组大小长度一致。
以上两个坑总体的解决方案代码,看下里面修改的部分(注释部分),在段落分割部分就把异常的分块去掉,从源头上保证 documents 的正常以及后面 documents 和 embeddings 数组大小一致:
lines = full_text.split('。\n' )
for text in lines:
buffer = text.strip(' ' ).replace('\n' , ' ' ).replace('[' , '' ).replace(']' , '' )
if len (buffer) < 10 :
continue
if buffer:
paragraphs.append(buffer)
buffer = ''
row_count = 0
if buffer and len (buffer) > 10 :
paragraphs.append(buffer)
return paragraphs
注意:文档分块不一定是按段落分。常见的分块策略包括固定字符数分块(Fixed Size Chunking)、递归字符分块(Recursive Character Splitting)以及语义分块(Semantic Chunking)。
3. Prompt 模板 上面我们已经拿到了检索回来的相关文档。下面我们写一个 Prompt 模板用来组装这些文档以及用户的提问。
def build_prompt (prompt_template, **kwargs ):
'''将 Prompt 模板赋值'''
prompt = prompt_template
for k, v in kwargs.items():
if isinstance (v,str ):
val = v
elif isinstance (v, list ) and all (isinstance (elem, str ) for elem in v):
val = '\n' .join(v)
else :
val = str (v)
prompt = prompt.replace(f"__{k.upper()} __" ,val)
return prompt
prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。
确保你的回复完全依据下述已知信息。不要编造答案。
如果下述已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
已知信息:
__INFO__
用户问:
__QUERY__
请用中文回答用户问题。
"""
注意以上最重要的提示词,要求大模型完全按照给定的文本回答问题:
你的任务是根据下述给定的已知信息回答用户问题。确保你的回复完全依据下述已知信息。不要编造答案。如果下述已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
4. 使用大模型得到答案
4.1 封装 OpenAI 接口 def get_completion (prompt, model="gpt-3.5-turbo-1106" ):
'''封装 openai 接口'''
messages = [{"role" : "user" , "content" : prompt}]
response = client.chat.completions.create(
model=model,
messages=messages,
temperature=0 ,
)
return response.choices[0 ].message.content
4.2 组装 Prompt prompt = build_prompt(prompt_template, info=results['documents' ][0 ], query=user_query)
print (prompt)
4.3 使用大模型得到答案 response = get_completion(prompt)
print (response)
5. 总结 至此,我们已经实现了 RAG 的基本流程。总结下流程:
(5)使用得到的 k 个文本块和用户提问组装 Prompt 模板
5.1 封装 RAG 我们将 RAG 流程封装一下,createVectorDB 完成离线部分,创建出向量数据库和灌入数据。chat 完成在线部分。
class RAG_Bot :
def __init__ (self, n_results=2 ):
self .llm_api = get_completion
self .n_results = n_results
def createVectorDB (self, file ):
print (file)
pdf_loader = PDFFileLoader(file)
self .vector_db = MyVectorDBConnector("demo" , get_embeddings)
self .vector_db.add_documents(pdf_loader.getParagraphs())
def chat (self, user_query ):
search_results = self .vector_db.search(user_query,self .n_results)
prompt = build_prompt(prompt_template, info=search_results['documents' ][0 ], query=user_query)
response = self .llm_api(prompt)
return response
rag_bot = RAG_Bot()
rag_bot.createVectorDB("D:\GitHub\LEARN_LLM\RAG\如何向 ChatGPT 提问以获得高质量答案:提示技巧工程完全指南.pdf" )
response = rag_bot.chat("什么是角色提示?" )
print ("response=====================>" )
print (response)
5.2 完整代码 from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer
class PDFFileLoader ():
def __init__ (self, file ) -> None :
self .paragraphs = self .extract_text_from_pdf(file, page_numbers=[0 ,3 ])
i = 1
for para in self .paragraphs:
print (f"========= 第{i} 段 ==========" )
print (para+"\n" )
i += 1
def getParagraphs (self ):
return self .paragraphs
def extract_text_from_pdf (self, filename, page_numbers=None ):
'''从 PDF 文件中(按指定页码)提取文字'''
paragraphs = []
buffer = ''
full_text = ''
for i, page_layout in enumerate (extract_pages(filename)):
if page_numbers is not None and i not in page_numbers:
continue
for element in page_layout:
if isinstance (element, LTTextContainer):
full_text += element.get_text() + '\n'
lines = full_text.split('。\n' )
for text in lines:
buffer = text.strip(' ' ).replace('\n' , ' ' ).replace('[' , '' ).replace(']' , '' )
if len (buffer) < 10 :
continue
if buffer:
paragraphs.append(buffer)
buffer = ''
row_count = 0
if buffer and len (buffer) > 10 :
paragraphs.append(buffer)
return paragraphs
from openai import OpenAI
import os
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())
client = OpenAI()
def get_embeddings (texts, model="text-embedding-3-small" ):
'''封装 OpenAI 的 Embedding 模型接口'''
data = client.embeddings.create(input =texts, model=model).data
return [x.embedding for x in data]
import chromadb
from chromadb.config import Settings
class MyVectorDBConnector :
def __init__ (self, collection_name, embedding_fn ):
chroma_client = chromadb.Client(Settings(allow_reset=True ))
chroma_client.reset()
self .collection = chroma_client.get_or_create_collection(name=collection_name)
self .embedding_fn = embedding_fn
def add_documents (self, documents ):
'''向 collection 中添加文档与向量'''
self .collection.add(
embeddings=self .embedding_fn(documents),
documents=documents,
ids=[f"id{i} " for i in range (len (documents))]
)
def search (self, query, top_n ):
'''检索向量数据库'''
results = self .collection.query(
query_embeddings=self .embedding_fn([query]),
n_results=top_n
)
return results
def build_prompt (prompt_template, **kwargs ):
'''将 Prompt 模板赋值'''
prompt = prompt_template
for k, v in kwargs.items():
if isinstance (v,str ):
val = v
elif isinstance (v, list ) and all (isinstance (elem, str ) for elem in v):
val = '\n' .join(v)
else :
val = str (v)
prompt = prompt.replace(f"__{k.upper()} __" ,val)
return prompt
prompt_template = """
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。
确保你的回复完全依据下述已知信息。不要编造答案。
如果下述已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
已知信息:
__INFO__
用户问:
__QUERY__
请用中文回答用户问题。
"""
def get_completion (prompt, model="gpt-3.5-turbo-1106" ):
'''封装 openai 接口'''
messages = [{"role" : "user" , "content" : prompt}]
response = client.chat.completions.create(
model=model,
messages=messages,
temperature=0 ,
)
return response.choices[0 ].message.content
class RAG_Bot :
def __init__ (self, n_results=2 ):
self .llm_api = get_completion
self .n_results = n_results
def createVectorDB (self, file ):
print (file)
pdf_loader = PDFFileLoader(file)
self .vector_db = MyVectorDBConnector("demo" , get_embeddings)
self .vector_db.add_documents(pdf_loader.getParagraphs())
def chat (self, user_query ):
search_results = self .vector_db.search(user_query,self .n_results)
prompt = build_prompt(prompt_template, info=search_results['documents' ][0 ], query=user_query)
print ("prompt===================>" )
print (prompt)
response = self .llm_api(prompt)
return response
rag_bot = RAG_Bot()
rag_bot.createVectorDB("D:\GitHub\LEARN_LLM\RAG\如何向 ChatGPT 提问以获得高质量答案:提示技巧工程完全指南.pdf" )
response = rag_bot.chat("什么是角色提示?" )
print ("response=====================>" )
print (response)
6. 进阶优化与挑战 RAG 是一个增强大模型垂直领域能力和减少幻觉的通用方法论,所以了解其原理和流程对实现出效果较好的大模型应用非常有用。
但是上面也可以看到,它也限制了大模型使用其自身的知识库去回答问题,只能够用给定的文本回复问题。这就导致这个 RAG 应用的通用性大大降低。
另外,从 RAG 流程中也可以看到要想实现的效果好,也是困难重重:
文本分割的粒度太小,查找到的参考文本较少,上下文缺失。
文本颗粒度太大,参考文本太多,消耗 token,同时也会带入更多的干扰信息,导致大模型出现幻觉的概率增加。
建议在实际项目中采用递归字符分块(Recursive Character Splitting),设置合适的 chunk_size 和 chunk_overlap(重叠窗口),通常 overlap 设置为 chunk_size 的 10%-20% 有助于保持上下文连贯性。
(2)有些问题的回答是需要依赖上下文的,怎样将上下文所在的文本块都找出来也不容易。
可以通过多跳检索(Multi-hop Retrieval)或者引入重排序(Re-ranking)机制来解决。例如,先用向量检索召回 Top 50,再用 Cross-Encoder 模型进行精排,选取 Top 5 输入 LLM。
(3)召回正确性 :召回文档的相关性也对结果比较重要。查找出的文档虽然与用户提问的向量值比较相似,但某些时候,最相似的并不一定是与问题答案相关的。
可以尝试混合搜索(Hybrid Search),结合关键词匹配(BM25)和向量相似度,提高召回的准确率。
选择参数量更大、指令遵循能力更强的模型通常能获得更好的效果,但成本也会相应增加。
目前针对以上各个困难都有非常多的研究,还在快速发展阶段,未形成一套通用、效果好的方法论。
后续可以针对这部分进行深入探索和学习,关注和整理当下最新的 RAG 调优方法,例如 LangChain 的高级组件、LlamaIndex 的索引策略等。
相关免费在线工具 加密/解密文本 使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
RSA密钥对生成器 生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
Mermaid 预览与可视化编辑 基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
随机西班牙地址生成器 随机生成西班牙地址(支持马德里、加泰罗尼亚、安达卢西亚、瓦伦西亚筛选),支持数量快捷选择、显示全部与下载。 在线工具,随机西班牙地址生成器在线工具,online
Gemini 图片去水印 基于开源反向 Alpha 混合算法去除 Gemini/Nano Banana 图片水印,支持批量处理与下载。 在线工具,Gemini 图片去水印在线工具,online
curl 转代码 解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online