1. 项目概述
最近在研究如何用Python实现基于Chroma的RAG(Retrieval-Augmented Generation)检索系统,这个技术组合在实际应用中表现相当出色。RAG的核心思想是通过检索相关文档片段来增强大语言模型的生成能力,而Chroma作为轻量级向量数据库,完美契合了这个需求场景。
我实现的这个系统主要包含两个核心模块:embeding.py负责文本向量化和检索,function.py处理文档分块。整个流程可以概括为:文档分块→向量化存储→问题向量化→相似度检索→生成回答。这种架构特别适合知识库问答、文档摘要等场景,实测下来回答质量比单纯用大语言模型要准确得多。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 核心组件解析
2.1 Chroma向量数据库
Chroma是我选择的核心存储组件,主要考虑以下几点:
- 轻量级:单文件持久化(./chroma.db),部署简单
- Python原生支持:API设计非常符合Python开发者的习惯
- 高性能:实测在千万级向量下仍能保持毫秒级检索
初始化客户端和集合的代码非常简洁:
python复制chromadb_client = chromadb.PersistentClient("./chroma.db")
chromadb_collection = chromadb_client.get_or_create_collection("question")
注意:同一个集合(collection)中的向量维度必须一致,这里我们使用gemini-embedding-exp-03-07模型的输出维度
2.2 文本嵌入模型
选用Google的Gemini Embedding模型主要基于以下考量:
- 支持区分文档嵌入和查询嵌入(task_type参数)
- 在检索任务上表现优异
- 免费额度足够个人和小型项目使用
嵌入函数的实现关键点:
python复制def embed(text: str, store: bool) -> list[float]:
result = google_client.models.embed_content(
model=EMBEDDING_MODEL,
contents=text,
config={
"task_type": "RETRIEVAL_DOCUMENT" if store else "RETRIEVAL_QUERY"
}
)
return result.embeddings[0].values
实操心得:RETRIEVAL_DOCUMENT和RETRIEVAL_QUERY两种任务类型会产生不同的嵌入向量,这种区分能显著提升检索相关性
3. 完整实现流程
3.1 文档预处理
function.py中的分块逻辑值得详细说明:
- 按双换行符(\n\n)进行初步分割
- 保留Markdown标题结构
- 将标题与后续内容关联
python复制def get_chunks() -> list[str]:
content = read_data()
chunks = content.split('\n\n')
result = []
header = ""
for c in chunks:
if c.startswith("#"):
header += f"{c}\n"
else:
result.append(f"{header}{c}")
header = ""
return result
避坑指南:简单的按段落分块会导致标题与内容分离,这种处理方式保持了文档的语义完整性
3.2 向量库构建
create_db()函数完成了从原始文本到向量库的转换:
python复制def create_db() -> None:
for idx, c in enumerate(function.get_chunks()):
embedding = embed(c, store=True)
chromadb_collection.upsert(
ids=str(idx),
documents=c,
embeddings=embedding
)
参数说明:
- ids: 使用简单序号作为唯一标识
- documents: 存储原始文本用于后续展示
- embeddings: 向量化后的稠密向量
3.3 查询处理流程
query_db函数实现了核心检索逻辑:
python复制def query_db(question: str) -> list[str]:
question_embedding = embed(question, store=False)
result = chromadb_collection.query(
query_embeddings=question_embedding,
n_results=5
)
return result["documents"][0]
关键参数n_results控制返回的相关文档数量,需要根据应用场景调整:
- 知识问答:3-5个片段足够
- 文档摘要:可能需要10-15个片段
4. 问答生成实现
主程序中的prompt构建技巧值得关注:
python复制prompt = "Please answer user's question according to context\n"
prompt += f"Question: {question}\n"
prompt += "Context:\n"
for c in chunks:
prompt += f"{c}\n"
prompt += "-------------\n"
这种模板设计:
- 明确指示模型基于上下文回答
- 清晰分隔问题和上下文
- 多个上下文片段用分隔线区分
生成调用非常简单:
python复制result = google_client.models.generate_content(
model=LLM_MODEL,
contents=prompt
)
5. 性能优化技巧
5.1 批量处理文档
当前实现是逐条处理文档,对于大规模数据建议:
python复制# 批量嵌入
batch_size = 32
chunks = function.get_chunks()
for i in range(0, len(chunks), batch_size):
batch = chunks[i:i+batch_size]
embeddings = [embed(c, True) for c in batch]
chromadb_collection.upsert(
ids=[str(i) for i in range(i, i+len(batch))],
documents=batch,
embeddings=embeddings
)
5.2 混合检索策略
单纯向量检索有时会漏掉关键词匹配的文档,可以结合BM25:
python复制from rank_bm25 import BM25Okapi
# 建立BM25索引
tokenized_chunks = [c.split() for c in chunks]
bm25 = BM25Okapi(tokenized_chunks)
def hybrid_query(question: str, alpha=0.5):
# 向量检索
vec_results = query_db(question)
# 关键词检索
tokenized_q = question.split()
bm25_scores = bm25.get_scores(tokenized_q)
top_k = np.argsort(bm25_scores)[-5:][::-1]
keyword_results = [chunks[i] for i in top_k]
# 混合去重
return list(dict.fromkeys(vec_results + keyword_results))
5.3 缓存机制
频繁查询时可以缓存嵌入结果:
python复制from functools import lru_cache
@lru_cache(maxsize=1000)
def cached_embed(text: str, store: bool) -> list[float]:
return embed(text, store)
6. 常见问题排查
6.1 检索结果不相关
可能原因及解决方案:
- 嵌入模型不匹配:确保查询和文档使用相同的模型
- 分块大小不当:建议200-500字/块
- 缺少元数据过滤:可以为文档添加metadata进行筛选
6.2 生成答案质量差
改进方法:
- 优化prompt模板
- 增加相关上下文数量
- 添加指令"若上下文不包含答案,请回答'我不知道'"
6.3 处理长文档技巧
对于超长文档:
- 分层分块:先按章节分,再按段落分
- 添加位置元数据:
python复制chromadb_collection.upsert(
ids=str(idx),
documents=c,
embeddings=embedding,
metadatas={"section": "3.2", "doc_id": "manual_v2"}
)
7. 扩展应用场景
7.1 多语言支持
只需更换嵌入模型:
python复制EMBEDDING_MODEL = "text-embedding-multilingual-001"
7.2 对话历史管理
在prompt中添加对话上下文:
python复制prompt = f"Previous conversation:\n{history}\n\n"
prompt += "Please answer user's question according to context\n"
7.3 领域知识微调
可以训练领域特定的嵌入模型:
- 使用sentence-transformers库
- 在领域数据上继续训练
- 替换默认的embed()函数实现
我在实际项目中发现,Chroma+RAG的组合特别适合以下场景:
- 企业内部知识库问答
- 产品文档智能助手
- 法律/医疗等专业领域咨询
最后分享一个实用技巧:定期用代表性问题测试系统,建立评估基准,这样可以量化检索质量的改进效果。
