第四阶段：第一个 RAG（Chroma + LLM）

上一阶段我们学会了用向量找相关文档，但每次都要重算向量。本阶段把向量存进向量数据库 Chroma，再把检索结果交给 LLM 生成答案——这就是完整的 RAG（检索增强生成）。

一、流程

1. 文档 → 向量：用本地 bge-m3 把文档转成向量；
2. 存入 Chroma：向量和原文一起持久化保存到磁盘；
3. 检索 top 3：把问题转成向量，找出最相近的 3 条文档；
4. 生成回答：把“资料 + 问题”一起交给 LLM，让它只根据资料回答。

pip install chromadb    # 安装 Chroma 向量数据库

二、04_rag.py 完整代码

import os
import chromadb                   # 向量数据库
from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()  # 读取 .env 里的 LLM 配置

# 对话模型：中转平台上的 LLM，负责“生成回答”
llm = OpenAI(api_key=os.getenv("LLM_API_KEY"),
             base_url=os.getenv("LLM_BASE_URL"))

# Embedding 模型：本地 LM Studio 上的 bge-m3，负责“文字转向量”
emb = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")


def embed(texts):
    """把一组文字转成向量，返回普通 Python 列表（Chroma 需要列表格式）"""
    resp = emb.embeddings.create(model="bge-m3", input=texts)
    return [item.embedding for item in resp.data]


# 创建持久化客户端：数据保存在当前目录的 chroma_db 文件夹，程序重启后仍在
client = chromadb.PersistentClient(path="./chroma_db")

# 获取或创建一个“集合”（类似数据库里的表）
# hnsw:space 设为 cosine，表示用余弦距离衡量相似度
collection = client.get_or_create_collection(
    "company_docs", metadata={"hnsw:space": "cosine"})

# 公司知识库：5 条规章制度
docs = [
    "公司年假为每年 15 天，入职满一年后生效",
    "报销需在费用发生后 30 天内提交发票",
    "办公室 Wi-Fi 密码每月更新一次，可找行政获取",
    "员工生日当月可领取 200 元购物卡",
    "加班需提前在 OA 系统中申请，否则不计加班费",
]

# 只有集合为空时才写入，避免每次运行都重复添加
if collection.count() == 0:
    collection.add(
        ids=[f"doc{i}" for i in range(len(docs))],  # 每条文档需要唯一 id
        documents=docs,                             # 原文
        embeddings=embed(docs),                     # 对应的向量
    )


def ask(question):
    """RAG 主流程：检索 → 拼 Prompt → 调用 LLM"""
    # 1. 检索：把问题转成向量，找最相近的 3 条
    results = collection.query(query_embeddings=embed([question]), n_results=3)
    retrieved = results["documents"][0]   # 第一个（也是唯一一个）问题的检索结果

    print("【检索到】")
    for d in retrieved:
        print(" -", d)

    # 2. 拼 Prompt：把资料和问题放在一起，并加上防编造约束
    context = "\n".join(retrieved)
    prompt = f"""请只根据下面的资料回答问题。
如果资料中没有答案，就回答“资料中没有相关信息”，不要编造。

资料：
{context}

问题：{question}"""

    # 3. 调用 LLM 生成回答；temperature=0 让回答稳定、少发挥
    resp = llm.chat.completions.create(
        model=os.getenv("LLM_MODEL"),
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    print("【回答】", resp.choices[0].message.content)


# 命令行循环提问，输入 exit 退出
while True:
    q = input("\n问题：")
    if q == "exit":
        break
    ask(q)

三、实测结果（真实运行数据）

问题	检索到（top 3）	回答
我入职半年能休年假吗	年假 / 生日 / 加班	不能。资料显示：年假入职满一年后生效。你入职半年，尚未满一年。
不提前申请加班会怎样？	加班 / 年假 / 报销	不提前在 OA 系统中申请加班，则不计加班费。
公司几点下班？	年假 / 加班 / 报销	资料中没有相关信息。

第一个问题里，LLM 结合资料做了推理（半年 < 一年 → 不能），这是单纯检索做不到的。

四、要点

1. 向量库总会返回“最接近”的几条，哪怕完全无关。问“几点下班”时，它照样返回了年假、加班、报销。所以 Prompt 里的防编造约束（“没有就说资料中没有相关信息”）非常关键，否则 LLM 可能根据无关资料瞎编。
2. temperature=0：问答类任务要稳定、忠实于资料，不需要创意。

五、提交代码

chroma_db/ 是本地生成的数据库文件，不要传到 GitHub：

echo chroma_db/ >> .gitignore    # 把数据库目录加入忽略列表
git add .
git commit -m "04 第一个 RAG"
git push

下一节预告

用 LlamaIndex 读取 PDF / Word 文档、自动切分，构建一个回答时注明出处的企业知识库。

👉 下一阶段：LlamaIndex 企业知识库