使用Chroma作为向量库,OpenAI作为 embedding 与生成模型。
1. 安装依赖
pip install chromadb openai tiktoken2. 设置 OpenAI Key
import os os.environ["OPENAI_API_KEY"] = "your-api-key"3. 完整RAG流程
import chromadb from chromadb.utils import embedding_functions from openai import OpenAI # ----- 初始化 ----- openai_client = OpenAI() # 使用OpenAI的embedding函数 openai_ef = embedding_functions.OpenAIEmbeddingFunction( api_key=os.environ["OPENAI_API_KEY"], model_name="text-embedding-ada-002" ) # 创建本地Chroma客户端(持久化到./chroma_db目录) chroma_client = chromadb.PersistentClient(path="./chroma_db") # 创建或获取collection collection_name = "my_knowledge" try: chroma_client.delete_collection(collection_name) except: pass collection = chroma_client.create_collection( name=collection_name, embedding_function=openai_ef ) # ----- 1. 准备知识库文档 ----- documents = [ "OpenAI 的 GPT-4 是一个强大的大语言模型,支持多模态和长上下文。", "Chroma 是一个开源的向量数据库,专门用于AI应用和embedding检索。", "RAG 是检索增强生成,能够减少大模型幻觉并提供最新知识。", "Python 是数据科学和AI领域最常用的编程语言。" ] # 添加文档到Chroma collection.add( documents=documents, ids=[f"doc_{i}" for i in range(len(documents))] ) print(f"知识库已添加 {collection.count()} 条文档") # ----- 2. 检索函数 ----- def retrieve(query, top_k=2): results = collection.query( query_texts=[query], n_results=top_k ) return results['documents'][0] # ----- 3. 生成函数(基于检索结果) ----- def rag_generate(query): # 检索 retrieved_docs = retrieve(query) # 构建prompt context = "\n\n".join(retrieved_docs) prompt = f"""基于以下参考信息回答问题。如果参考信息中没有答案,请说不知道。 参考信息: {context} 问题:{query} 回答:""" # 调用OpenAI生成 response = openai_client.chat.completions.create( model="gpt-3.5-turbo", messages=[ {"role": "system", "content": "你是一个基于知识库的问答助手。"}, {"role": "user", "content": prompt} ], temperature=0 ) return response.choices[0].message.content # ----- 4. 测试 ----- if __name__ == "__main__": question = "什么是RAG?" answer = rag_generate(question) print(f"问题: {question}") print(f"回答: {answer}") print("\n--- 另一个问题 ---") question2 = "Chroma是什么?" answer2 = rag_generate(question2) print(f"问题: {question2}") print(f"回答: {answer2}")4. 输出示例
知识库已添加 4 条文档 问题: 什么是RAG? 回答: RAG是检索增强生成,能够减少大模型幻觉并提供最新知识。 问题: Chroma是什么? 回答: Chroma是一个开源的向量数据库,专门用于AI应用和embedding检索。关键点说明
Embedding:
text-embedding-ada-002将文本转为向量,Chroma自动完成检索:默认使用余弦相似度,
top_k=2返回最相关的两个文档生成:将检索到的文档拼接到 prompt 中,由 GPT 总结回答
持久化:
PersistentClient使知识库保存到磁盘,下次启动可复用
这个流程已经是一个完整的 RAG 原型,你可以通过替换documents为自己的数据进行扩展。
不是。标准的RAG(如前面Chroma示例)是一次性检索,不是增量检索。
标准RAG和增量检索两者的核心区别
| 特性 | 标准RAG | 增量检索 |
|---|---|---|
| 检索次数 | 1次(query时) | 多次(迭代式) |
| 上下文构建 | 一次拼装 | 逐步累积 |
| 典型场景 | 问答、摘要 | 多跳推理、长文生成 |
标准RAG流程(一次检索)
def standard_rag(query): # 只检索一次 docs = vector_db.search(query, top_k=3) # 一次生成 return llm.generate(query, context=docs) # 示例:问"Claude和GPT哪个好?" # 一次检索就拿到两种模型对比的文档增量检索(Iterative RAG / Self-Ask)
def incremental_rag(query): context = [] # 第1轮:检索初步信息 docs1 = vector_db.search(query, top_k=2) context.extend(docs1) # 第2轮:基于已有信息,发现缺失 missing = llm.identify_gaps(query, context) docs2 = vector_db.search(missing, top_k=2) context.extend(docs2) # 第3轮:可能需要再细化 if need_more_detail(context): sub_query = llm.extract_subquery(context) docs3 = vector_db.search(sub_query, top_k=1) context.extend(docs3) return llm.generate(query, context)实际案例对比
场景:查询"2024年AI融资最多的三家公司"
标准RAG(1次):
检索: "2024年AI融资最多的三家公司" 返回: ["OpenAI融资66亿", "Anthropic融资40亿", "Inflection融资13亿"] ✅ 直接满足增量检索(3轮):
# 第1轮 检索: "2024年AI融资最多的三家公司" 返回: ["OpenAI融资金额", "AI融资趋势报告"] # 第2轮(发现不完整) 检索: "2024年融资金额排名第四的公司" 返回: ["Cohere融资5亿"] # 第3轮(验证数据) 检索: "CBInsights 2024 AI融资排行榜" 返回: [完整排名表格]什么时候需要增量检索?
| 问题类型 | 标准RAG | 增量RAG |
|---|---|---|
| "什么是RAG?" | ✅ 一次够 | ❌ 过度设计 |
| "A和B的区别" | ✅ 同时检索 | ❌ 浪费 |
| "某电影主角的出生地" | ❌ 缺失信息 | ✅ 先找电影→再找主演→再找出生地 |
| "对比GPT-4、Claude、Gemini的性能" | ❌ 可能不完整 | ✅ 分别检索每个 |