(各位访问大哥能否留一些宝贵建议给小弟,没人的话就当个人日常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时一定要严谨且明确,避免提示词之间有冲突。

        第六天就到这吧。

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐