少量 Python代码,从＂能跑＂到＂能打＂ 三个版本搞懂 RAG

pip install openai chromadb

import osfrom openai import OpenAI# DeepSeek API（用于对话生成）deepseek_key = os.getenv("DEEPSEEK_API_KEY")if not deepseek_key:    raise ValueError("请设置环境变量 DEEPSEEK_API_KEY")# 硅基流动 API（用于 embedding）siliconflow_key = os.getenv("SILICONFLOW_API_KEY")if not siliconflow_key:    raise ValueError("请设置环境变量 SILICONFLOW_API_KEY")# 对话模型客户端chat_client = OpenAI(base_url="https://api.deepseek.com", api_key=deepseek_key)# Embedding 模型客户端embedding_client = OpenAI(base_url="https://api.siliconflow.cn/v1", api_key=siliconflow_key)

 export DEEPSEEK_API_KEY="your_key_here" export SILICONFLOW_API_KEY="your_key_here"

documents = [    {        "id": "doc1",        "title": "请假制度",        "content": """公司请假制度（2024年修订版）一、年假规定1. 入职满1年不满10年的员工，每年享有5天带薪年假。2. 入职满10年不满20年的员工，每年享有10天带薪年假。3. 入职满20年以上的员工，每年享有15天带薪年假。4. 年假需提前3个工作日在OA系统中提交申请，直属上级审批通过后生效。5. 年假可以分次使用，但每次不少于半天。6. 当年未使用的年假可以顺延至次年第一季度，逾期作废。二、病假规定1. 病假需提供正规医院出具的诊断证明和病假条。2. 3天以内的病假由直属上级审批。3. 3天以上的病假需由部门总监审批。4. 病假期间工资按基本工资的80%发放。5. 连续病假超过30天的，按公司长期病假政策处理。三、事假规定1. 事假为无薪假期，按日扣除工资。2. 事假每次不超过3天，需提前2个工作日申请。3. 全年事假累计不超过15天。4. 事假审批流程与病假相同。"""    },    {        "id": "doc2",        "title": "报销制度",        "content": """公司费用报销制度一、差旅报销标准1. 交通费：飞机经济舱、高铁二等座实报实销...2. 住宿费：一线城市（北上广深）每晚不超过500元...（完整内容见代码仓库）"""    },    {        "id": "doc3",        "title": "技术规范",        "content": """前端技术开发规范一、技术栈要求1. 框架统一使用 React 18+...2. 状态管理使用 Zustand...（完整内容见代码仓库）"""    },    {        "id": "doc4",        "title": "新人入职指南",        "content": """新员工入职指南一、入职第一天1. 上午 9:00 到前台报到...（完整内容见代码仓库）"""    },    {        "id": "doc5",        "title": "绩效考核制度",        "content": """绩效考核管理办法一、考核周期1. 公司采用季度考核制...（完整内容见代码仓库）"""    }]

def simple_chunk(text, chunk_size=200):    """按固定大小切块"""    chunks = []    for i in range(0, len(text), chunk_size):        chunk = text[i:i + chunk_size].strip()        if chunk:            chunks.append(chunk)    return chunks

import chromadbchroma_client = chromadb.Client()collection = chroma_client.create_collection(name="company_docs_v1")def get_embedding(text):    """调用硅基流动 Embedding 接口"""    response = embedding_client.embeddings.create(        model="netease-youdao/bce-embedding-base_v1",        input=text    )    return response.data[0].embeddingdef build_index_v1(documents):    """V1：简单切块 + 入库"""    print("正在建立索引...")    for doc in documents:        chunks = simple_chunk(doc["content"], chunk_size=200)        for i, chunk in enumerate(chunks):            chunk_id = f"{doc['id']}_chunk_{i}"            embedding = get_embedding(chunk)            collection.add(                ids=[chunk_id],                embeddings=[embedding],                documents=[chunk],                metadatas=[{"source": doc["title"]}]            )        print(f"  ✅ {doc['title']}：切成 {len(chunks)} 块")    print(f"索引建立完成，共 {collection.count()} 个文档块")

def search_v1(query, top_k=3):    """在向量库中搜索最相关的文档块"""    query_embedding = get_embedding(query)    results = collection.query(        query_embeddings=[query_embedding],        n_results=top_k    )    return results["documents"][0], results["metadatas"][0]def rag_query_v1(question):    """V1 完整的 RAG 问答流程"""    print(f"\n{'='*60}")    print(f"  V1 朴素 RAG")    print(f"  问题: {question}")    print(f"{'='*60}")    # 第一步：检索    chunks, metadatas = search_v1(question, top_k=3)    print(f"\n📚 检索到 {len(chunks)} 段相关内容：")    for i, (chunk, meta) in enumerate(zip(chunks, metadatas)):        print(f"\n  [{i+1}] 来源：{meta['source']}")        print(f"      {chunk[:80]}...")    # 第二步：把检索到的内容塞进 prompt    context = "\n\n---\n\n".join(chunks)    prompt = f"""请根据以下参考资料回答用户的问题。如果参考资料中没有相关信息，请如实告知。参考资料：{context}用户问题：{question}请直接回答，不要编造参考资料中没有的信息。"""    # 第三步：让大模型生成回答    response = chat_client.chat.completions.create(        model="deepseek-chat",        messages=[{"role": "user", "content": prompt}]    )    answer = response.choices[0].message.content    print(f"\n✅ 回答：\n{answer}")    return answer

# 先建索引（只需要跑一次）build_index_v1(documents)# 测试几个问题rag_query_v1("入职满5年有几天年假？")rag_query_v1("出差住酒店的报销标准是多少？")rag_query_v1("前端项目用什么构建工具？")

块1: "公司请假制度（2024年修订版）一、年假规定 1. 入职满1年不满10年的     员工...4. 年假需提前3个工作日在OA系统中提交申"    ← 注意这里块2: "请，直属上级审批通过后生效。5. 年假可以分次使用..."    ← 断裂了！

def overlap_chunk(text, chunk_size=200, overlap=50):    """带重叠的分块：相邻块之间有 overlap 个字符的重叠"""    chunks = []    start = 0    while start < len(text):        end = start + chunk_size        chunk = text[start:end].strip()        if chunk:            chunks.append(chunk)        start += chunk_size - overlap  # 关键：每次只前进 (chunk_size - overlap)    return chunks

collection_v2 = chroma_client.create_collection(    name="company_docs_v2",    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度)

def build_index_v2(documents):    """V2：重叠切块 + 入库"""    print("正在建立 V2 索引...")    for doc in documents:        chunks = overlap_chunk(doc["content"], chunk_size=200, overlap=50)        for i, chunk in enumerate(chunks):            chunk_id = f"{doc['id']}_v2_chunk_{i}"            embedding = get_embedding(chunk)            collection_v2.add(                ids=[chunk_id],                embeddings=[embedding],                documents=[chunk],                metadatas=[{"source": doc["title"]}]            )        print(f"  ✅ {doc['title']}：切成 {len(chunks)} 块（带重叠）")    print(f"V2 索引建立完成，共 {collection_v2.count()} 个文档块")def search_v2(query, top_k=5, threshold=0.6):    """带相似度过滤的检索"""    query_embedding = get_embedding(query)    results = collection_v2.query(        query_embeddings=[query_embedding],        n_results=top_k,        include=["documents", "metadatas", "distances"]    )    filtered_chunks = []    filtered_metas = []    for doc, meta, distance in zip(        results["documents"][0],        results["metadatas"][0],        results["distances"][0]    ):        similarity = 1 - distance  # ChromaDB 余弦距离转相似度        if similarity >= threshold:            filtered_chunks.append(doc)            filtered_metas.append({**meta, "similarity": round(similarity, 3)})    return filtered_chunks, filtered_metas

def rag_query_v2(question):    """V2 改进版 RAG：重叠分块 + 相似度过滤"""    print(f"\n{'='*60}")    print(f"  V2 改进版 RAG")    print(f"  问题: {question}")    print(f"{'='*60}")    chunks, metadatas = search_v2(question, top_k=5, threshold=0.3)    if not chunks:        print("\n❌ 未找到足够相关的文档内容")        return "抱歉，知识库中没有找到与您的问题相关的内容。"    print(f"\n📚 检索到 {len(chunks)} 段相关内容（已过滤低相关度）：")    for i, (chunk, meta) in enumerate(zip(chunks, metadatas)):        print(f"\n  [{i+1}] 来源：{meta['source']}（相似度：{meta['similarity']}）")        print(f"      {chunk[:80]}...")    context = "\n\n---\n\n".join(chunks)    prompt = f"""请根据以下参考资料回答用户的问题。要求：1. 只基于参考资料中的信息回答，不要编造2. 如果参考资料不足以回答问题，如实告知3. 回答要准确、简洁、有条理参考资料：{context}用户问题：{question}"""    response = chat_client.chat.completions.create(        model="deepseek-chat",        messages=[{"role": "user", "content": prompt}]    )    answer = response.choices[0].message.content    print(f"\n✅ 回答：\n{answer}")    return answer

build_index_v2(documents)rag_query_v2("病假工资怎么算？超过30天怎么办？")

def rerank_with_llm(query, chunks, top_n=3):    """用大模型对检索结果做重排序"""    if len(chunks) <= top_n:        return list(range(len(chunks)))    chunks_text = ""    for i, chunk in enumerate(chunks):        chunks_text += f"\n[段落{i+1}]\n{chunk[:150]}...\n"    prompt = f"""你是一个文档相关性评估专家。用户的问题是："{query}"以下是从知识库中检索到的多个段落，请根据与用户问题的相关性进行排序。{chunks_text}请按相关性从高到低排列段落编号，只输出编号列表，格式如：1,3,2,5,4不要输出任何解释，只输出逗号分隔的编号。"""    response = chat_client.chat.completions.create(        model="deepseek-chat",        messages=[{"role": "user", "content": prompt}],        temperature=0    )    try:        order_str = response.choices[0].message.content.strip()        order = [int(x.strip()) - 1 for x in order_str.split(",")]        return order[:top_n]    except:        return list(range(top_n))def rag_query_v3(question):    """V3 完整版 RAG：重叠分块 + 相似度过滤 + Rerank"""    print(f"\n{'='*60}")    print(f"  V3 Rerank 版 RAG")    print(f"  问题: {question}")    print(f"{'='*60}")    # 第一步：粗筛——向量检索，多召回一些    chunks, metadatas = search_v2(question, top_k=6, threshold=0.25)    if not chunks:        return "抱歉，知识库中没有找到相关内容。"    print(f"\n📚 粗筛：向量检索召回 {len(chunks)} 段")    for i, (chunk, meta) in enumerate(zip(chunks, metadatas)):        print(f"  [{i+1}] {meta['source']}（相似度：{meta['similarity']}）- {chunk[:50]}...")    # 第二步：精排——Rerank 重排序    print(f"\n🔄 精排：Rerank 重排序中...")    reranked_indices = rerank_with_llm(question, chunks, top_n=3)    reranked_chunks = [chunks[i] for i in reranked_indices]    reranked_metas = [metadatas[i] for i in reranked_indices]    print(f"\n📚 精排后 Top {len(reranked_chunks)} 段：")    for i, (chunk, meta) in enumerate(zip(reranked_chunks, reranked_metas)):        print(f"  [{i+1}] {meta['source']}（相似度：{meta['similarity']}）- {chunk[:50]}...")    # 第三步：生成回答    context = "\n\n---\n\n".join(reranked_chunks)    prompt = f"""请根据以下参考资料回答用户的问题。要求：1. 严格基于参考资料回答，不编造信息2. 如果信息不足，如实告知3. 如果涉及数字、金额、时间等关键信息，务必准确引用4. 回答要结构清晰，重点突出参考资料：{context}用户问题：{question}"""    response = chat_client.chat.completions.create(        model="deepseek-chat",        messages=[{"role": "user", "content": prompt}]    )    answer = response.choices[0].message.content    print(f"\n✅ 回答：\n{answer}")    return answer

ragqueryv3( "我是新入职的员工，试用期工资打几折？转正后怎么评定绩效？拿到S评分有什么奖励？" )