记录下学习agent应用开发的第六天(完整开发基于RAG框架的聊天机器人)
(各位访问大哥能否留一些宝贵建议给小弟,没人的话就当个人日常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时一定要严谨且明确,避免提示词之间有冲突。
第六天就到这吧。
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐
所有评论(0)