(各位访问大哥能否留一些宝贵建议给小弟,没人的话就当个人日常blog)我把前几天学的内容整合在一起做了一个基于RAG框架的聊天机器人,包含保存历史会话、滑窗裁剪、流式输出等功能。核心逻辑是:把用户输入的文本和知识库chunk通过Embedding模型转成向量并计算比对,然后找出和用户文本最相关的几个知识库chunk(R),再把这些chunk和用户的文本拼成提示词prompt(A)一起发给接入的模型,模型思考后再根据system prompt进行回答(G)。
完整代码:
""" day6: RAG 聊天机器人 """ import json import os from pathlib import Path import chromadb # 向量数据库 import tiktoken # Token 计数器:把文字转成 token 数 from sentence_transformers import SentenceTransformer # Embedding 模型:文字 → 向量 from openai import OpenAI from dotenv import load_dotenv # ============================================================ # 第一部分:基础设施初始化 # ============================================================ # 加载 API Key 和接口地址 root = Path(__file__).parent.parent # 从 day6/ 回到项目根目录 load_dotenv(root / ".env") # 大模型客户端:负责"生成回答" client = OpenAI( api_key=os.getenv("DEEPSEEK_API_KEY"), base_url=os.getenv("DEEPSEEK_BASE_URL"), ) # Embedding 模型:负责"文字 → 向量" emb_model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2") # Token 计数器:负责"算总长" encoder = tiktoken.get_encoding("cl100k_base") # 上下文窗口安全阈值:128K 是模型上限,留 28K 给回复用 MAX_TOKENS = 100_000 # ============================================================ # 第二部分:ChromaDB 向量数据库 # 负责"存储和检索知识库向量" # ============================================================ chroma_client = chromadb.PersistentClient( path=str(Path(__file__).parent / "chroma_db") # 向量数据存到磁盘 ) collection = chroma_client.get_or_create_collection(name="knowledge_base") # 知识库文本文件路径 KNOWLEDGE_FILE = Path(__file__).parent / "python_knowledge.txt" # ============================================================ # 第三部分:JSON 文件持久化 # 负责"关机后恢复聊天记录" # ============================================================ HISTORY_FILE = Path(__file__).parent / "chat_history.json" # ============================================================ # 第四部分:函数定义 # ============================================================ # --- 4.1 知识库加载 --- def load_knowledge(): """ 读知识库文本 → 切段落 → 转 Embedding 向量 → 存入 ChromaDB 只在第一次运行时执行,后续启动跳过(因为向量已存在) """ if collection.count() > 0: print(f"[知识库已加载,共 {collection.count()} 个段落]") return with open(KNOWLEDGE_FILE, "r", encoding="utf-8") as f: text = f.read() chunks = [s.strip() + "。" for s in text.split("。") if s.strip()] embeddings = emb_model.encode(chunks).tolist() collection.add( ids=[f"chunk_{i}" for i in range(len(chunks))], embeddings=embeddings, documents=chunks, ) print(f"[知识库已加载,共 {len(chunks)} 个段落]") # --- 4.2 聊天记录持久化 --- def load_history() -> list: """启动时从 JSON 文件恢复之前的聊天记录""" if HISTORY_FILE.exists(): with open(HISTORY_FILE, "r", encoding="utf-8") as f: return json.load(f) return [] def save_history(history: list): """每次对话完成后,把全部聊天记录覆盖写入 JSON 文件""" with open(HISTORY_FILE, "w", encoding="utf-8") as f: json.dump(history, f, ensure_ascii=False, indent=2) # --- 4.3 Token 计数与裁剪 --- def count_tokens(messages: list) -> int: """遍历 messages 列表中每条消息,统计总 token 数""" total = 0 for msg in messages: total += len(encoder.encode(msg["content"])) return total def trim_history(history: list) -> list: """ 滑动窗口裁剪:如果聊天记录超过 MAX_TOKENS, 从最前面开始删除最老的对话对(user + assistant), 直到总 token 数降到安全线以下 """ if not history: return history while count_tokens(history) > MAX_TOKENS and len(history) >= 2: history = history[2:] # 每次跳 2 条 = 删一整轮问答 return history # --- 4.4 向量检索 --- def search(query: str, top_k: int = 3) -> list: """ 用户问题 → Embedding 模型转成向量 → ChromaDB 检索 → 返回相似度最高的 top_k 个段落 """ query_embedding = emb_model.encode([query]).tolist()[0] results = collection.query( query_embeddings=[query_embedding], n_results=top_k ) return results["documents"][0] # --- 4.5 组装 Prompt --- def build_messages(query: str, context_chunks: list, history: list) -> list: """ 拼接最终发给大模型的 messages: 结构 = [system(含检索结果)] + [历史对话] + [当前问题] 这样大模型能同时看到: - 知识库的相关内容(回答依据) - 之前的聊天记录(理解追问和指代) - 用户当前的问题 """ context = "\n".join([f"- {c}" for c in context_chunks]) system_msg = { "role": "system", "content": f"""你是小码,基于知识库和对话历史回答问题的 AI 助手。 规则: 1. 只根据提供的资料回答问题,资料里没有就诚实说"资料中没有相关信息" 2. 回答时引用资料中的具体内容 3. 参考对话历史理解用户的追问和指代 4. 用中文回答,简洁清晰 当前资料: {context}""" } return [system_msg] + history + [{"role": "user", "content": query}] # --- 4.6 流式输出 --- def chat_stream(messages: list) -> str: """ 流式调用大模型:AI 边生成边逐字显示(打字机效果) 返回完整回复文本,供存入历史记录和 JSON 文件 """ stream = client.chat.completions.create( model="deepseek-v3.2", messages=messages, temperature=0.5, max_tokens=800, stream=True, ) full = "" for chunk in stream: content = chunk.choices[0].delta.content if content: print(content, end="", flush=True) # 逐字打印,不换行 full += content # 同时拼成完整回复 print() return full # ============================================================ # 第五部分:主程序 — 串联全部功能 # ============================================================ # 1. 加载知识库 load_knowledge() # 2. 从 JSON 恢复聊天记录 conversation_history = load_history() if conversation_history: print(f"[恢复了 {len(conversation_history) // 2} 轮历史对话]") # 3. 启动时裁剪一次:防止历史记录太长 old_len = len(conversation_history) conversation_history = trim_history(conversation_history) if len(conversation_history) < old_len: print(f"[自动裁剪了 {(old_len - len(conversation_history)) // 2} 轮旧对话]") # 4. 显示 token 用量 token_count = count_tokens(conversation_history) print(f"[当前上下文: {token_count} / {MAX_TOKENS} tokens]") print("=" * 50) print(" 小码 RAG 全功能版 (quit 退出, !clear 清空记忆)") print("=" * 50) # 5. 对话循环 while True: question = input("\n你: ").strip() if not question: continue if question.lower() in ("quit", "exit", "q"): print("小码: 再见!") break if question == "!clear": conversation_history = [] save_history(conversation_history) # 清空内存后同步清空 JSON 文件 print("[记忆已清空]") continue # R:向量检索 — 从知识库中找相关段落 results = search(question) print(f"[检索到 {len(results)} 个相关段落]") # A:拼接 Prompt — 知识库 + 历史对话 + 当前问题 messages = build_messages(question, results, conversation_history) # 发送前裁剪 messages = trim_history(messages) # G:生成回答 — 流式输出 print("小码: ", end="", flush=True) try: answer = chat_stream(messages) # 存入对话历史:内存 list + JSON 文件 conversation_history.append({"role": "user", "content": question}) conversation_history.append({"role": "assistant", "content": answer}) save_history(conversation_history) except Exception as e: import traceback print(f"\n出错了: {e}") traceback.print_exc()知识库内容:
Python 是一门解释型、面向对象的高级编程语言。 Python 由 Guido van Rossum 在 1989 年圣诞节期间开始设计,1991 年首次发布。 Python 的名字来源于 Guido 喜爱的 BBC 电视节目《蒙提·派森的飞行马戏团》。 Python 的设计哲学强调代码可读性,使用缩进来表示代码块,而不是花括号。 Python 2.0 于 2000 年发布,引入了列表推导式、垃圾回收机制。 Python 3.0 于 2008 年发布,是一次不向后兼容的大版本更新。 Python 3.0 的主要变化包括:print 变为函数、整数除法返回浮点数、Unicode 成为默认字符串类型。 Python 2.7 是 Python 2 系列的最后一个版本,官方支持于 2020 年 1 月 1 日终止。 pip 是 Python 的包管理工具,用于安装和管理第三方库。 Python 的标准库非常丰富,涵盖文件操作、网络通信、正则表达式、JSON 处理等。 Python 最流行的 Web 框架包括 Django、Flask 和 FastAPI。 Python 在数据科学领域的主要库包括 NumPy、Pandas 和 Matplotlib。 Python 在机器学习领域的主要库包括 Scikit-learn、TensorFlow 和 PyTorch。 Python 的变量不需要声明类型,是动态类型语言。 Python 支持多种编程范式,包括面向对象、函数式和过程式编程。 Python 使用 GIL(全局解释器锁)来保证线程安全,但这限制了多线程的并行性能。 Python 3.12 引入了子解释器,为绕过 GIL 提供了新的可能性。运行效果演示:
向机器人提问跟知识库有关的内容能得到清晰的回答。
发送跟知识库无关的问题就会触发system prompt的限制1:只根据提供的资料回答问题,资料里没有就诚实说"资料中没有相关信息"。
以下存在特殊情况:我在相邻的两轮会话中分别告诉机器人我的名字和今天天气,再次询问时虽然都根据知识库和system prompt回答“没有关于***的相关信息”,但是下面图一的回答中有提到“小明”,而图二的回答没有提到“下雨”。
对于上面这两种情况,我的感觉是模型出现了幻觉,但opencode给我的解释是:模型把"我"识别为一个指代词,而不是一个需要查资料的事实,所以不会去知识库里搜"我"这个词,而是查看历史内容看"我"指代什么,符合限制3:参考对话历史理解用户的追问和指代。"今天天气怎么样"这句话里没有任何指代词,模型判定这是一个独立的新问题,不需要翻历史(“今天下雨了”),于是就按照限制1:只根据提供的资料回答问题,资料里没有就诚实说"资料中没有相关信息"正常回答。总的来说就是system prompt写的不够明确,并不是模型出现了幻觉。
总结:今天结合前几天学的知识完整开发了基于RAG框架的聊天机器人,算是完成度比较高的小项目。在开发过程中我不仅回顾了之前学过的内容,还对RAG框架有了更深入的理解。然后是在写system prompt时一定要严谨且明确,避免提示词之间有冲突。
第六天就到这吧。