第7章: 检索做准做快：RAG Retrieval 召回与重排实战

import chromadb

# 目标：本节重点讲“query 怎么查”。embedding 用 Chroma 默认的就够了（更合理也更简单）。
# 注意：默认 embedding 可能会在第一次运行时下载模型/依赖，所以首次执行可能较慢。

# ✅ 关键改动：用持久化库 + 只初始化一次数据
# 否则你后面的 3️⃣/4️⃣/… 节再复用 demo_retrieval 时，经常会因为“数据只在内存里”而查到空。
DB_PATH = "./chroma_demo_retrieval"

def ensure_demo_data(collection):
    """只在 collection 为空时写入演示数据（可重复运行、可跨脚本复用）"""
    try:
        if collection.count() > 0:
            return
    except Exception:
        pass

    documents = [
        # account
        "重置密码：进入设置-账号安全-选择重置方式",
        "忘记密码：通过手机号验证码或邮箱验证码找回",
        "登录邮箱修改：先验证原邮箱，再绑定新邮箱",
        "手机号变更：需要人脸/短信二次验证，防止盗号",
        "账号异常登录：先检查异地登录记录，再修改密码",

        # security
        "开启两步验证：在账号安全中打开 2FA，提高安全性",
        "安全建议：使用强密码，避免重复使用同一个密码",

        # refund
        "退款流程：提交申请-审核-原路退回",
        "退款到账时间：通常 1-3 个工作日，视支付渠道而定",
        "退款失败：检查支付账号状态，或联系人工客服处理",

        # order/billing
        "订单查询：在订单中心按时间筛选，支持导出账单",
        "发票开具：填写抬头与税号，电子发票发送到邮箱",
        "扣费异常：先核对账单明细，再提交工单核查",
        "订阅取消：到订阅管理里取消，取消后本周期仍可使用",

        # errors
        "错误码 401：通常是 token 过期或未登录",
        "错误码 403：通常是权限不足或鉴权失败",
        "错误码 429：请求过快被限流，建议退避重试",
        "网络超时：检查代理/网络，或提高 timeout 并重试",

        # misc
        "上传失败：检查文件大小限制与格式要求",
        "账号注销：注销前请确认余额与未完成订单",
    ]

    metadatas = [
        {"doc_id": "acc_reset_pwd", "category": "account"},
        {"doc_id": "acc_forgot_pwd", "category": "account"},
        {"doc_id": "acc_change_email", "category": "account"},
        {"doc_id": "acc_change_phone", "category": "account"},
        {"doc_id": "acc_risk_login", "category": "account"},

        {"doc_id": "sec_2fa", "category": "security"},
        {"doc_id": "sec_password", "category": "security"},

        {"doc_id": "refund_process", "category": "refund"},
        {"doc_id": "refund_time", "category": "refund"},
        {"doc_id": "refund_fail", "category": "refund"},

        {"doc_id": "order_query", "category": "order"},
        {"doc_id": "billing_invoice", "category": "billing"},
        {"doc_id": "billing_charge", "category": "billing"},
        {"doc_id": "sub_cancel", "category": "billing"},

        {"doc_id": "err_401", "category": "error"},
        {"doc_id": "err_403", "category": "error"},
        {"doc_id": "err_429", "category": "error"},
        {"doc_id": "err_timeout", "category": "error"},

        {"doc_id": "upload_fail", "category": "misc"},
        {"doc_id": "acc_close", "category": "account"},
    ]

    ids = [str(i) for i in range(1, len(documents) + 1)]
    collection.add(documents=documents, metadatas=metadatas, ids=ids)


# 1) 初始化：创建持久化 client/collection，并确保有演示数据
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_or_create_collection("demo_retrieval")
ensure_demo_data(collection)

# 2) 核心：Top-K 检索（n_results 就是 K）
query = "用户如何重置密码？"
res = collection.query(
    query_texts=[query],
    n_results=3,
    include=["documents", "metadatas", "distances"],
)

# 3) 返回结构是“按 query 分组”的二维列表，所以这里用 [0]
docs = res["documents"][0]
metas = res["metadatas"][0]
dists = res["distances"][0]

for i, (text, meta, dist) in enumerate(zip(docs, metas, dists), start=1):
    score = 1 - dist  # Chroma 默认距离越小越相关，转为相似度分数（0~1）
    print(i, meta, "distance=", dist, "score=", f"{score:.3f}")
    print(text)

import chromadb

DB_PATH = "./chroma_demo_retrieval"

def ensure_demo_data(collection):
    try:
        if collection.count() > 0:
            return
    except Exception:
        pass
    raise RuntimeError("demo_retrieval 为空：请先运行 2️⃣ Baseline 的初始化示例（会写入演示数据）")


client = chromadb.PersistentClient(path=DB_PATH)

# 承接上一节：复用同一个持久化库里的 demo_retrieval 数据
collection = client.get_or_create_collection("demo_retrieval")
ensure_demo_data(collection)

# 现实情况：Top-K 返回的是“最像的 K 条”，但后几条经常匹配度不够（也会被塞进 prompt）
query = "如何申请退款？"
res = collection.query(
    query_texts=[query],
    n_results=4,
    include=["documents", "metadatas", "distances"],
)

docs = res["documents"][0]
metas = res["metadatas"][0]
dists = res["distances"][0]

print("=== 原始 Top-K（里面可能混入低匹配内容）===")
for t, m, d in zip(docs, metas, dists):
    print(m, "distance=", d)
    print(t)

# 核心：对返回结果做阈值过滤（distance 越小越相关）
threshold_dist = 0.45  # 示例阈值：需要你根据真实数据调
filtered = [(t, m, d) for t, m, d in zip(docs, metas, dists) if d is not None and d <= threshold_dist]

print("\n=== 阈值过滤后（只保留足够相关的证据）===")
if not filtered:
    print("证据不足：触发澄清/转人工/或直接声明‘知识库证据不足’")
else:
    for t, m, d in filtered:
        print(m, "distance=", d)
        print(t)

import chromadb

DB_PATH = "./chroma_demo_retrieval"

def ensure_demo_data(collection):
    try:
        if collection.count() > 0:
            return
    except Exception:
        pass
    raise RuntimeError("demo_retrieval 为空：请先运行 2️⃣ Baseline 的初始化示例（会写入演示数据）")


client = chromadb.PersistentClient(path=DB_PATH)

# 承接上面：复用同一个持久化库里的 demo_retrieval 数据
collection = client.get_or_create_collection("demo_retrieval")
ensure_demo_data(collection)

# 核心：where 只让“满足条件”的数据参与检索（先过滤再相似度排序）
query = "如何修改登录邮箱？"
res = collection.query(
    query_texts=[query],
    n_results=3,
    where={"category": "account"},
    include=["documents", "metadatas"],
)

print(res["documents"][0])
print(res["metadatas"][0])

你是“知识库检索查询改写器”。

目标：把用户问题改写成适合检索的查询语句。

要求：
1) 只输出 JSON 数组，不要解释。
2) 输出 3 条以内，每条不超过 25 个字。
3) 每条尽量包含：关键实体/错误提示/目标动作。
4) 不要编造不存在的产品名/错误码。

用户问题：{user_question}

输出示例：
["账号异常 登录失败 如何处理", "登录提示账号异常 排查步骤", "错误码/提示语 账号异常 解决"]

import chromadb
import json
import os
import dashscope
from dashscope import Generation

# 目标：把“口语问题”改写成更像“搜索词”的 query，再用多 query 检索扩大召回覆盖

DB_PATH = "./chroma_demo_retrieval"

# 通义千问（DashScope）
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
if not DASHSCOPE_API_KEY:
    raise RuntimeError("环境变量 DASHSCOPE_API_KEY 未设置")
dashscope.api_key = DASHSCOPE_API_KEY


def ensure_demo_data(collection):
    try:
        if collection.count() > 0:
            return
    except Exception:
        pass
    raise RuntimeError("demo_retrieval 为空：请先运行 2️⃣ Baseline 的初始化示例（会写入演示数据）")


def rewrite_to_queries(user_question: str, max_queries: int = 3):
    """用 LLM 输出 JSON 数组（多查询），失败则降级为原问题"""
    prompt = f"""你是“知识库检索查询改写器”。

目标：把用户问题改写成适合检索的查询语句。

要求：
1) 只输出 JSON 数组，不要解释。
2) 输出 {max_queries} 条以内，每条不超过 25 个字。
3) 每条尽量包含：关键实体/错误提示/目标动作。
4) 不要编造不存在的产品名/错误码。

用户问题：{user_question}

输出示例：
["账号异常 登录失败 如何处理", "登录提示账号异常 排查步骤", "错误码/提示语 账号异常 解决"]
"""

    try:
        response = Generation.call(
            model='qwen-turbo',
            messages=[
                {'role': 'system', 'content': '你是一个helpful助手'},
                {'role': 'user', 'content': prompt}
            ],
            temperature=0.7,
            max_tokens=1000
        )

        content = (response.output.text or "").strip()

        queries = json.loads(content)
        if isinstance(queries, list) and queries:
            queries = [q.strip() for q in queries if isinstance(q, str) and q.strip()]
            return queries[:max_queries] if queries else [user_question]
        return [user_question]
    except Exception:
        return [user_question]


def multi_query_search(collection, queries, top_k_each: int = 4):
    """对每个 query 检索 Top-K，再按 doc_id 去重（保留最小 distance）"""
    best_by_doc_id = {}  # doc_id -> (doc, meta, dist, hit_query)

    for q in queries:
        res = collection.query(
            query_texts=[q],
            n_results=top_k_each,
            include=["documents", "metadatas", "distances"],
        )

        docs = res.get("documents", [[]])[0]
        metas = res.get("metadatas", [[]])[0]
        dists = res.get("distances", [[]])[0]

        for doc, meta, dist in zip(docs, metas, dists):
            doc_id = (meta or {}).get("doc_id") or ""
            if not doc_id:
                continue

            prev = best_by_doc_id.get(doc_id)
            if prev is None or (dist is not None and dist < prev[2]):
                best_by_doc_id[doc_id] = (doc, meta, dist, q)

    merged = list(best_by_doc_id.values())
    merged.sort(key=lambda x: (x[2] is None, x[2]))
    return merged


# 1) 连接持久化库（承接第 2 节）
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_or_create_collection("demo_retrieval")
ensure_demo_data(collection)

# 2) 原始口语问题
user_question = "我登不上去了，提示账号异常"

# 3) Rewrite → 多查询
queries = rewrite_to_queries(user_question, max_queries=3)
print("=== Rewrite 结果（多查询）===")
for i, q in enumerate(queries, start=1):
    print(i, q)

# 4) 多查询检索 → 合并去重
hits = multi_query_search(collection, queries, top_k_each=4)

print("\n=== 合并去重后（按 doc_id）===")
for i, (doc, meta, dist, hit_q) in enumerate(hits, start=1):
    score = None if dist is None else (1 - dist)
    if score is None:
        print(i, meta, f"hit_query={hit_q}")
    else:
        print(i, meta, f"distance={dist:.3f}", f"score={score:.3f}", f"hit_query={hit_q}")
    print(doc)

import chromadb

client = chromadb.Client()
try:
    client.delete_collection("demo_batch")
except Exception:
    pass

collection = client.get_or_create_collection("demo_batch")
collection.add(
    documents=[
        "退款流程：提交申请-审核-原路退回",
        "退款时效：不同支付渠道到账时间不同",
        "重置密码：进入设置-账号安全-选择重置方式",
        "修改邮箱：先验证原邮箱，再绑定新邮箱",
    ],
    metadatas=[
        {"doc_id": "refund_1", "category": "refund"},
        {"doc_id": "refund_2", "category": "refund"},
        {"doc_id": "acc_1", "category": "account"},
        {"doc_id": "acc_2", "category": "account"},
    ],
    ids=["1", "2", "3", "4"],
)

queries = ["怎么退款？", "怎么重置密码？"]
res = collection.query(
    query_texts=queries,
    n_results=2,
    include=["documents", "metadatas", "distances"],
)

# 返回是“按 query 分组”的二维列表：res["documents"][i] 对应第 i 个 query
for i, q in enumerate(queries):
    print("\n=== query:", q, "===")
    for text, meta, dist in zip(res["documents"][i], res["metadatas"][i], res["distances"][i]):
        print(meta, "distance=", dist)
        print(text)

import chromadb

client = chromadb.Client()
try:
    client.delete_collection("demo_where")
except Exception:
    pass

collection = client.get_or_create_collection("demo_where")
collection.add(
    documents=[
        "退款流程：提交申请-审核-原路退回",
        "退款时效：不同支付渠道到账时间不同",
        "重置密码：进入设置-账号安全-选择重置方式",
        "修改邮箱：先验证原邮箱，再绑定新邮箱",
    ],
    metadatas=[
        {"doc_id": "refund_1", "category": "refund", "version": 2},
        {"doc_id": "refund_2", "category": "refund", "version": 1},
        {"doc_id": "acc_1", "category": "account", "version": 2},
        {"doc_id": "acc_2", "category": "account", "version": 1},
    ],
    ids=["1", "2", "3", "4"],
)

query = "怎么退款？"

# 示例：只检索 version>=2 且 category 在 refund/account 里的文档
where = {
    "$and": [
        {"version": {"$gte": 2}},
        {"category": {"$in": ["refund", "account"]}},
    ]
}

res = collection.query(
    query_texts=[query],
    n_results=5,
    where=where,
    include=["documents", "metadatas", "distances"],
)

print(res["metadatas"][0])
print(res["documents"][0])

import chromadb

client = chromadb.Client()
try:
    client.delete_collection("demo_where_doc")
except Exception:
    pass

collection = client.get_or_create_collection("demo_where_doc")
collection.add(
    documents=[
        "错误码 403：通常是权限不足或鉴权失败",
        "错误码 404：资源不存在，检查 URL 或路由",
        "退款流程：提交申请-审核-原路退回",
    ],
    metadatas=[
        {"doc_id": "err_403"},
        {"doc_id": "err_404"},
        {"doc_id": "refund_1"},
    ],
    ids=["1", "2", "3"],
)

# 先用语义检索召回，再用 where_document 做“包含某关键词”的过滤
res = collection.query(
    query_texts=["怎么处理错误码？"],
    n_results=5,
    where_document={"$contains": "错误码"},
    include=["documents", "metadatas"],
)

print(res["metadatas"][0])
print(res["documents"][0])

import chromadb

client = chromadb.Client()
try:
    client.delete_collection("demo_maintain")
except Exception:
    pass

collection = client.get_or_create_collection("demo_maintain")
collection.add(
    documents=["退款流程：提交申请-审核-原路退回"],
    metadatas=[{"doc_id": "refund_1", "category": "refund", "version": 1}],
    ids=["refund_1"],
)

# 1) get：按条件取出数据（用于排错/回放/追溯）
rows = collection.get(where={"category": "refund"}, include=["documents", "metadatas"])
print(rows["ids"], rows["metadatas"], rows["documents"])

# 2) upsert：有就更新，没有就新增（适合“增量入库/修订版本”）
collection.upsert(
    ids=["refund_1", "refund_2"],
    documents=[
        "退款流程：提交申请-审核-原路退回（v2 修订版）",
        "退款时效：不同支付渠道到账时间不同",
    ],
    metadatas=[
        {"doc_id": "refund_1", "category": "refund", "version": 2},
        {"doc_id": "refund_2", "category": "refund", "version": 1},
    ],
)

# 3) delete：删除过期/违规/无权限的数据（可按 ids 或 where 条件）
collection.delete(where={"version": {"$lte": 1}})

rag_demo/
├── ingest.py              # 第6章：离线入库（把第6章“存储实战”的代码保存为这个文件）
├── query.py               # 第7章：在线检索 + 生成
├── my_rag_db/             # 第6章示例使用的持久化目录（db_path="./my_rag_db"）
└── requirements.txt

dashscope
chromadb
numpy

export DASHSCOPE_API_KEY="你的key"

import os
import chromadb
import dashscope
from dashscope import TextEmbedding, Generation


dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")


def format_context(docs, max_chars=2400):
    chunks = []
    used = 0
    for i, d in enumerate(docs, start=1):
        piece = f"[证据{i}]\n{d}\n"
        if used + len(piece) > max_chars:
            break
        chunks.append(piece)
        used += len(piece)
    return "\n".join(chunks)


def ask(question: str, db_path: str = "./my_rag_db", collection_name: str = "python_docs", top_k: int = 5):
    dashscope.api_key = os.getenv("DASHSCOPE_API_KEY", "")
    if not dashscope.api_key:
        raise RuntimeError("请先设置环境变量 DASHSCOPE_API_KEY，然后再运行")

    def embed_query(text: str):
        resp = TextEmbedding.call(model='text-embedding-v2', input=text)
        return resp.output['embeddings'][0]['embedding']

    client = chromadb.PersistentClient(path=db_path)
    collection = client.get_collection(collection_name)

    # 与第6章保持一致：第6章用 DashScope 的 text-embedding-v2 写入向量。
    # 这里显式计算 query embedding，并用 query_embeddings 检索，避免“embedding_function 不一致”导致 query_texts 不能用。
    q_emb = embed_query(question)
    results = collection.query(
        query_embeddings=[q_emb],
        n_results=top_k,
        include=["documents", "metadatas", "distances"],
    )

    docs = results.get("documents", [[]])[0]
    metadatas = results.get("metadatas", [[]])[0]
    distances = results.get("distances", [[]])[0]

    # distance 越小越相关（阈值需要你根据数据调）
    threshold = 1.1
    filtered = [(d, m, dist) for d, m, dist in zip(docs, metadatas, distances) if dist is not None and dist <= threshold]

    if not filtered:
        return {
            "answer": "我在知识库中没有检索到足够相关的证据。你可以补充关键词或描述更具体的问题。",
            "sources": []
        }

    context = format_context([d for d, _, _ in filtered])

    prompt = f"""你是一个严谨的知识库问答助手。请仅基于【证据】回答问题；
如果证据不足以回答，请明确说“证据不足”。

【证据】
{context}

【问题】
{question}

【回答】
"""

    resp = Generation.call(model="qwen-plus", prompt=prompt)
    answer = resp.output["text"] if resp and resp.output else "生成失败"

    sources = [{"metadata": m, "distance": dist} for _, m, dist in filtered]
    return {"answer": answer, "sources": sources}


if __name__ == "__main__":
    q = input("请输入问题：").strip()
    result = ask(q)
    print("\n=== 回答 ===\n")
    print(result["answer"])
    print("\n=== 证据来源（metadata/distance）===\n")
    for s in result["sources"]:
        print(s)

python ingest.py
python query.py