【案例】物流行业信息问答系统

"""入库前文档体检：先看资料值不值得入库，再谈检索。

RAG 项目最贵的一类返工，是文档已经全量灌完、上线问了几十句才发现
资料本身有问题——扫描件没文字、表格被拉成一行、同一条制度有三个版本。
这些在入库前花十分钟就能查出来。

六项体检，逐项给出结论和处理建议。纯标准库，直接运行。
"""
import os
import re
import sys

MIN_CHARS = 50            # 一份文档少于这么多字，基本可以判定抽取失败
LONG_LINE = 200           # 单行超过这个长度，多半是表格被拍平了
DUP_WINDOW = 30           # 判定重复段落时比较的前缀长度


def read(path):
    with open(path, encoding="utf-8", errors="replace") as f:
        return f.read()


def check_extractable(text):
    """① 抽不抽得出文字。扫描件走到这一步会是空的或者几个零散字符。"""
    n = len(text.strip())
    ok = n >= MIN_CHARS
    return ok, "抽出 %d 字" % n, "少于 %d 字：多半是扫描件，要先过 OCR" % MIN_CHARS


def check_garbled(text):
    """② 有没有编码事故。U+FFFD 是解码失败留下的指纹。"""
    n = text.count("\ufffd")
    return n == 0, "替换字符 %d 个" % n, "有乱码：换编码重新抽取，别把乱码灌进库"


def check_long_lines(text):
    """③ 有没有被拍平的表格。

    表格转成纯文本后经常变成一行几百字，切块时要么被拦腰截断，
    要么整块超长——这类内容应该单独抽成结构化数据，不走向量检索。
    """
    lines = text.splitlines()
    bad = [i for i, ln in enumerate(lines, 1) if len(ln) > LONG_LINE]
    return not bad, "超长行 %d 条" % len(bad), \
        "表格被拍平了：结构化字段建议入数据库，别指望语义检索去查数值"


def check_duplicates(text):
    """④ 有没有大段重复。同一条制度复制到多份文档里非常常见。"""
    paras = [p.strip() for p in re.split(r"\n\s*\n", text) if len(p.strip()) > DUP_WINDOW]
    seen, dup = set(), 0
    for p in paras:
        key = p[:DUP_WINDOW]
        if key in seen:
            dup += 1
        seen.add(key)
    return dup == 0, "重复段落 %d 处" % dup, \
        "有重复：入库前去重，否则同一句话会在提示词里出现好几遍"

def check_structure(text):
    """⑤ 有没有可用的切分锚点。

    一份通篇没有空行也没有标点的文本，任何切分器都只能按字数硬切。
    """
    blank = text.count("\n\n")
    cn_stop = len(re.findall(r"[。；！？]", text))
    ok = blank > 0 or cn_stop > 3
    return ok, "空行 %d 处，中文句末标点 %d 个" % (blank, cn_stop), \
        "没有可用锚点：切分只能按字数硬切，块会拦腰断句"


def check_dates(text):
    """⑥ 有没有时效性内容。

    带日期的内容会过期，而向量库不会自己知道这件事。
    这类资料必须配一个更新机制，否则半年后系统会一本正经地讲旧规定。
    """
    dates = re.findall(r"\d{4}[-/年]\d{1,2}[-/月]\d{1,2}", text)
    return True, "日期 %d 处" % len(dates), ""


CHECKS = [
    ("① 能否抽出文字", check_extractable),
    ("② 有无编码乱码", check_garbled),
    ("③ 有无拍平表格", check_long_lines),
    ("④ 有无重复段落", check_duplicates),
    ("⑤ 有无切分锚点", check_structure),
    ("⑥ 有无时效内容", check_dates),
]


def main():
    path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "wuliu_kb.txt")
    text = read(path)

    print("体检对象：%s（%d 字节）\n" % (os.path.basename(path), os.path.getsize(path)))
    print("%-18s %-8s %-24s" % ("检查项", "结论", "实测"))
    print("-" * 70)

    failed = []
    for name, fn in CHECKS:
        ok, fact, advice = fn(text)
        print("%-18s %-8s %-24s" % (name, "通过" if ok else "注意", fact))
        if not ok and advice:
            failed.append((name, advice))

    print()
    if failed:
        print("需要处理的项：")
        for name, advice in failed:
            print("  %s → %s" % (name, advice))
    else:
        print("六项全部通过，这份资料可以进入切分环节。")

    dates = re.findall(r"\d{4}[-/年]\d{1,2}[-/月]\d{1,2}", text)
    if dates:
        print()
        print("⚠️ 文档里有 %d 处日期（如 %s）。" % (len(dates), dates[0]))
        print("   这类内容会过期，而向量库不会自己知道。上线时要一并定好：")
        print("   谁负责更新、多久更新一次、更新后哪些块需要重新入库。")


if __name__ == "__main__":
    main()

物流公司：速达物流
公司总部：北京市
业务范围：国际快递、仓储管理

货物追踪：
货物编号：ABC123456
发货日期：2023-01-15
当前位置：上海分拨中心
预计到达日期：2023-01-20

运输方式：
运输公司：快运通
运输方式：陆运
出发地：广州
目的地：重庆
预计运输时间：3天

仓储信息：
仓库名称：东方仓储中心
仓库位置：深圳市
存储货物类型：电子产品
存储条件：常温仓储
当前库存量：1000件

"""建库脚本：把物流文档变成一个可检索的向量目录。

只在资料更新时跑，跑完在磁盘上留下一个目录，问答脚本直接读它。
四步：加载 → 切分 → 向量化 → 落盘。

依赖：langchain、langchain-community、faiss-cpu、pymupdf，
Embedding 走本地 Ollama，不需要任何云端密钥。
"""
import os

from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 路径与模型都从环境变量取，换文档、换模型不用改代码
PDF_PATH = os.environ.get("RAG_PDF", "物流信息.pdf")
INDEX_DIR = os.environ.get("RAG_INDEX", "./faiss/wuliu")
EMBED_MODEL = os.environ.get("RAG_EMBED_MODEL", "mxbai-embed-large")

# 中文标点必须进分隔符列表，否则递归切分会一路降级到按字数硬切
SEPARATORS = ["\n\n", "\n", "。", "；", "，", " ", ""]

CHUNK_SIZE = int(os.environ.get("RAG_CHUNK_SIZE", "300"))
CHUNK_OVERLAP = int(os.environ.get("RAG_CHUNK_OVERLAP", "50"))


def load(pdf_path):
    """第一步：加载。PyMuPDFLoader 按页返回 Document，页码在 metadata 里。"""
    loader = PyMuPDFLoader(pdf_path)
    docs = loader.load()
    print("加载完成：%d 页" % len(docs))
    total = sum(len(d.page_content) for d in docs)
    print("正文合计 %d 字" % total)
    if total == 0:
        raise SystemExit(
            "抽出来 0 个字：这份 PDF 多半是扫描件，文字是图片。\n"
            "这种情况要先过 OCR，直接进切分器只会得到一个空库。")
    return docs


def split(docs):
    """第二步：切分。块大小与重叠都影响后面每一步，改之前先想清楚。"""
    splitter = RecursiveCharacterTextSplitter(
        separators=SEPARATORS,
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
    )
    chunks = splitter.split_documents(docs)
    print("切分完成：%d 块，平均 %.0f 字"
          % (len(chunks), sum(len(c.page_content) for c in chunks) / max(len(chunks), 1)))
    for i, c in enumerate(chunks[:3], 1):
        preview = c.page_content.replace("\n", " ")[:50]
        print("  样例块 %d（%d 字）：%s…" % (i, len(c.page_content), preview))
    return chunks


def build(chunks):
    """第三、四步：逐块向量化并落盘。

    ⚠️ 这里用的 Embedding 模型，问答时必须原样再用一次。
    换模型等于换了坐标系，旧库里的向量全部作废，必须整库重建。
    """
    embeddings = OllamaEmbeddings(model=EMBED_MODEL)
    db = FAISS.from_documents(chunks, embeddings)
    os.makedirs(os.path.dirname(INDEX_DIR) or ".", exist_ok=True)
    db.save_local(INDEX_DIR)
    print("已写入 %s（Embedding 模型：%s）" % (INDEX_DIR, EMBED_MODEL))
    return db


def main():
    docs = load(PDF_PATH)
    chunks = split(docs)
    db = build(chunks)

    # 建完立刻自检一次：能不能把一个已知答案原样找回来
    probe = "出发地是哪里"
    hits = db.similarity_search(probe, k=2)
    print("-" * 56)
    print("建库自检，问「%s」取回 %d 块：" % (probe, len(hits)))
    for i, d in enumerate(hits, 1):
        print("  [%d] %s" % (i, d.page_content.replace("\n", " ")[:60]))
    print("取回的内容对不上题，说明切分或模型有问题，别急着往下做问答。")


if __name__ == "__main__":
    main()

"""问答脚本：读磁盘上的向量目录，检索 top-k，拼提示词，交给本地模型。

每次提问都跑，不重建库。和建库脚本唯一的硬约束是：
Embedding 模型必须一模一样，否则检索出来的东西毫无意义。
"""
import os
import time

from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms import Ollama
from langchain_community.vectorstores import FAISS
from langchain.prompts import PromptTemplate

INDEX_DIR = os.environ.get("RAG_INDEX", "./faiss/wuliu")
EMBED_MODEL = os.environ.get("RAG_EMBED_MODEL", "mxbai-embed-large")
LLM_MODEL = os.environ.get("RAG_LLM_MODEL", "qwen2.5:7b")
TOP_K = int(os.environ.get("RAG_TOP_K", "3"))
SCORE_MAX = float(os.environ.get("RAG_SCORE_MAX", "1.2"))

PROMPT_TEMPLATE = """你是物流信息查询助手。请只根据【已知信息】回答【问题】。

规则：
1. 只使用【已知信息】里出现过的内容，不允许补充任何外部知识。
2. 【已知信息】不足以回答时，回复「资料中未提及」，不要推测。
3. 每个结论后面用方括号标出所依据的编号，例如 [1]。

【已知信息】
{context}

【问题】
{question}"""

PROMPT = PromptTemplate(input_variables=["context", "question"],
                        template=PROMPT_TEMPLATE)


def load_db():
    """读盘。

    allow_dangerous_deserialization=True 是因为 FAISS 的本地索引用 pickle
    存元数据，反序列化会执行代码。只对自己生成的目录开这个开关，
    从外部拿到的索引文件不要直接加载。
    """
    embeddings = OllamaEmbeddings(model=EMBED_MODEL)
    return FAISS.load_local(INDEX_DIR, embeddings,
                            allow_dangerous_deserialization=True)


def retrieve(db, question, k=TOP_K):
    """带分数检索。

    FAISS 默认返回的是 L2 距离，越小越近——和余弦相似度方向相反，
    所以这里的阈值是「大于就丢掉」，别照着相似度的写法写反了。
    """
    pairs = db.similarity_search_with_score(question, k=k)
    kept = [(doc, score) for doc, score in pairs if score <= SCORE_MAX]
    return kept


def build_context(pairs):
    lines = []
    for i, (doc, score) in enumerate(pairs, 1):
        page = doc.metadata.get("page")
        src = doc.metadata.get("source", "未知来源")
        where = "%s 第 %s 页" % (os.path.basename(str(src)), page + 1) if page is not None else src
        text = doc.page_content.replace("\n\n", "\n").strip()
        lines.append("[%d] （来源：%s，距离 %.3f）%s" % (i, where, score, text))
    return "\n".join(lines)


def answer(db, question):
    pairs = retrieve(db, question)
    if not pairs:
        return "资料中未提及。", []

    prompt_text = PROMPT.format(context=build_context(pairs), question=question)
    model = Ollama(model=LLM_MODEL)
    return model.invoke(prompt_text), pairs


def main():
    db = load_db()
    questions = [
        "我的快递出发地是哪？预计几天的时间到达？",
        "货物现在在哪个分拨中心？",
        "深圳的仓库存的是什么类型的货？",
        "公司的年假制度是怎么规定的？",
    ]
    for q in questions:
        start = time.time()
        text, pairs = answer(db, q)
        print("=" * 60)
        print("问：%s" % q)
        print("检索到 %d 块，耗时 %.2f 秒" % (len(pairs), time.time() - start))
        for i, (doc, score) in enumerate(pairs, 1):
            print("  [%d] 距离 %.3f  %s"
                  % (i, score, doc.page_content.replace("\n", " ")[:46]))
        print("答：%s" % text)


if __name__ == "__main__":
    main()

"""Web 问答界面：把问答链包成一个能对外用的页面。

依赖 streamlit。启动：
    streamlit run web_qa.py

界面只做三件事：收问题、显示答案、把检索到的原文摆出来让人自己核对。
第三件最重要——没有出处的知识库问答是不可验收的。
"""
import os

import streamlit as st
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.llms import Ollama
from langchain_community.vectorstores import FAISS
from langchain.prompts import PromptTemplate

INDEX_DIR = os.environ.get("RAG_INDEX", "./faiss/wuliu")
EMBED_MODEL = os.environ.get("RAG_EMBED_MODEL", "mxbai-embed-large")
LLM_MODEL = os.environ.get("RAG_LLM_MODEL", "qwen2.5:7b")

PROMPT = PromptTemplate(
    input_variables=["context", "question"],
    template="""你是物流信息查询助手。请只根据【已知信息】回答【问题】。
【已知信息】里没有的内容，回复「资料中未提及」，不要推测。
每个结论后面用方括号标出所依据的编号。

【已知信息】
{context}

【问题】
{question}""")

st.set_page_config(page_title="物流行业信息咨询系统", layout="wide")
st.title("物流行业信息咨询系统")


@st.cache_resource(show_spinner="正在加载向量库…")
def load_db():
    """缓存住，别每问一句就重新加载一次索引。

    不加缓存时，每次交互都会重新读盘并重建 Embedding 客户端，
    几百毫秒的操作被放大成每轮都等——这是这类页面最常见的卡顿来源。
    """
    embeddings = OllamaEmbeddings(model=EMBED_MODEL)
    return FAISS.load_local(INDEX_DIR, embeddings,
                            allow_dangerous_deserialization=True)


@st.cache_resource
def load_llm():
    return Ollama(model=LLM_MODEL)


def build_context(pairs):
    lines = []
    for i, (doc, score) in enumerate(pairs, 1):
        page = doc.metadata.get("page")
        src = os.path.basename(str(doc.metadata.get("source", "未知")))
        where = "%s 第 %s 页" % (src, page + 1) if page is not None else src
        lines.append("[%d] （%s）%s" % (i, where, doc.page_content.strip()))
    return "\n".join(lines)


with st.sidebar:
    st.subheader("检索参数")
    top_k = st.slider("取回多少块（top-k）", 1, 10, 3)
    score_max = st.slider("距离上限（越小越严）", 0.1, 2.0, 1.2, 0.1)
    st.caption("先固定阈值调 k，再固定 k 调阈值；两个一起动就说不清是谁起的作用。")

if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

if prompt := st.chat_input("请输入你的问题："):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        db = load_db()
        pairs = db.similarity_search_with_score(prompt, k=top_k)
        pairs = [(d, s) for d, s in pairs if s <= score_max]

        if not pairs:
            reply = "资料中未提及。"
            st.markdown(reply)
        else:
            text = PROMPT.format(context=build_context(pairs), question=prompt)
            reply = load_llm().invoke(text)
            st.markdown(reply)
            # 把检索到的原文摊开，答案对不对由人当场核对
            with st.expander("查看检索到的原文（%d 块）" % len(pairs)):
                for i, (doc, score) in enumerate(pairs, 1):
                    st.markdown("**[%d]** 距离 `%.3f`" % (i, score))
                    st.text(doc.page_content.strip())

        st.session_state.messages.append({"role": "assistant", "content": reply})

"""召回评测：给这套物流问答系统一个可验收的数字。

不做评测的 RAG 项目，调参全靠感觉——改完一个参数问两句觉得"好像好点了"，
然后就上线了。这个脚本把验收拆成可重复的两步：

    第一层  检索对不对：命中率 / MRR，能自动算
    第二层  答案对不对：只能人工核，脚本负责把核对表生成出来

纯标准库，直接运行，不依赖向量库和模型。
换成真实系统时，把 retrieve() 换成对 FAISS 的调用即可。
"""
import math
import re
from collections import Counter

# 知识库：与建库脚本切出来的块对应
CHUNKS = [
    "物流公司：速达物流。公司总部：北京市。业务范围：国际快递、仓储管理。",
    "货物编号：ABC123456。发货日期：2023-01-15。当前位置：上海分拨中心。",
    "预计到达日期：2023-01-20。",
    "运输公司：快运通。运输方式：陆运。出发地：广州。目的地：重庆。预计运输时间：3天。",
    "仓库名称：东方仓储中心。仓库位置：深圳市。存储货物类型：电子产品。",
    "存储条件：常温仓储。当前库存量：1000件。",
]

# 测试集：问题 → 允许命中的块下标集合。
# 这张表是整个评测的地基，必须由懂业务的人来标，不能让模型自己生成。
TESTSET = [
    ("我的快递出发地是哪？", {3}),
    ("这批货预计几天能到？", {3}),
    ("货物现在到哪个分拨中心了？", {1}),
    ("预计到达日期是哪天？", {2}),
    ("仓库在哪个城市？", {4}),
    ("仓库里存的是什么货？", {4}),
    ("当前库存还有多少？", {5}),
    ("公司总部在哪？", {0}),
]

# 这些问题库里根本没有，系统必须答"资料中未提及"
OUT_OF_SCOPE = [
    "公司的年假制度是怎么规定的？",
    "运费一公斤多少钱？",
]


def tokenize(t):
    return re.findall(r"[a-z0-9]+", t.lower()) + re.findall(r"[\u4e00-\u9fff]", t)


def embed(t):
    return Counter(tokenize(t))


def cosine(a, b):
    common = set(a) & set(b)
    dot = sum(a[k] * b[k] for k in common)
    na = math.sqrt(sum(v * v for v in a.values()))
    nb = math.sqrt(sum(v * v for v in b.values()))
    return dot / (na * nb) if na and nb else 0.0


INDEX = [(i, c, embed(c)) for i, c in enumerate(CHUNKS)]


def retrieve(question, k):
    """换成真实系统时，这里替换成 db.similarity_search_with_score(question, k)。"""
    qv = embed(question)
    scored = [(cosine(qv, v), i) for i, _c, v in INDEX]
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[:k]


def evaluate(k):
    """命中率：标准块在不在 top-k 里。
       MRR：标准块排第几名，取名次倒数再平均——排第 1 得 1 分，排第 3 得 0.33 分。
       命中率只看"进没进来"，MRR 还看"排得靠不靠前"，两个一起看才完整。
    """
    hit, rr_sum, detail = 0, 0.0, []
    for question, gold in TESTSET:
        got = retrieve(question, k)
        idxs = [i for _s, i in got]
        rank = next((r for r, i in enumerate(idxs, 1) if i in gold), None)
        if rank:
            hit += 1
            rr_sum += 1.0 / rank
        detail.append((question, gold, idxs, rank))
    n = len(TESTSET)
    return hit / n, rr_sum / n, detail


def main():
    print("知识库 %d 块，测试集 %d 条，范围外问题 %d 条\n"
          % (len(CHUNKS), len(TESTSET), len(OUT_OF_SCOPE)))

    print("=== 第一层：检索对不对（自动算） ===\n")
    print("%-4s %-10s %-10s %s" % ("k", "命中率", "MRR", "没命中的问题"))
    print("-" * 64)
    best_k, best_mrr = None, -1
    for k in (1, 2, 3, 4, 5):
        rate, mrr, detail = evaluate(k)
        misses = [q for q, _g, _i, rank in detail if rank is None]
        if mrr > best_mrr:
            best_k, best_mrr = k, mrr
        print("%-4d %-10.3f %-10.3f %s"
              % (k, rate, mrr, "、".join(misses) if misses else "无"))

    print("\nMRR 最高的是 k=%d（%.3f）。" % (best_k, best_mrr))
    print("注意 MRR 会随 k 增大而饱和：命中率还在涨、MRR 已经不动时，")
    print("多出来的那几条都排在很后面，基本只是在给提示词加噪声。")

    print("\n=== 逐条明细（k=3） ===\n")
    _rate, _mrr, detail = evaluate(3)
    for question, gold, idxs, rank in detail:
        flag = "第%d名" % rank if rank else "未命中"
        print("  %-18s 期望 %-6s 实取 %-12s %s"
              % (question, sorted(gold), idxs, flag))

    print("\n=== 范围外问题：必须被挡住 ===\n")
    for question in OUT_OF_SCOPE:
        top = retrieve(question, 1)
        score = top[0][0] if top else 0.0
        print("  %-22s 最高分 %.4f" % (question, score))
    print("  这几条的最高分要明显低于上面那批，阈值才卡得住。")
    print("  分数区分不开时，别去调阈值，先回头看切分和向量化。")

    print("\n=== 第二层：答案对不对（人工核对表） ===\n")
    print("把下面这张表打印出来，逐行人工判定，不要让模型自己给自己打分：\n")
    print("  %-20s %-8s %-8s %-8s %s" % ("问题", "有编造", "有漏答", "标了出处", "判定"))
    print("  " + "-" * 60)
    for question, _gold in TESTSET[:4]:
        print("  %-20s %-8s %-8s %-8s %s" % (question, "□", "□", "□", "□通过 □不通过"))
    print("\n  三列全部为否、且标了出处，才算这一条通过。")


if __name__ == "__main__":
    main()

"""行业知识库骨架：换个行业只改配置，不改流程。

这份模板把「物流」这三个字全部收进了配置区。换成法务、医疗、制造，
改的是 DOMAIN 里的四项内容，下面的建库与问答流程一行都不用动。

用法
    1. 改 DOMAIN 与 PATHS
    2. 把 TESTSET 换成自己行业的验收问题（这一步不能省）
    3. python3 rag_industry_skeleton.py build    建库
    4. python3 rag_industry_skeleton.py eval     跑验收
    5. python3 rag_industry_skeleton.py ask "问题"
"""
import os
import sys

# ------------------------------------------------------------ 配置区
DOMAIN = {
    # TODO: 换行业时只改这四项
    "name": "物流",
    "role": "物流信息查询助手",
    "doc_types": ["运单", "仓储", "线路", "时效"],
    "refuse": "资料中未提及。",
}

PATHS = {
    "docs": os.environ.get("RAG_DOC_DIR", "./docs"),      # TODO: 资料目录
    "index": os.environ.get("RAG_INDEX", "./faiss/kb"),   # TODO: 索引落盘目录
}

MODELS = {
    # ⚠️ 建库与问答必须用同一个 embed，换了就要整库重建
    "embed": os.environ.get("RAG_EMBED_MODEL", "mxbai-embed-large"),
    "llm": os.environ.get("RAG_LLM_MODEL", "qwen2.5:7b"),
}

PARAMS = {
    "chunk_size": 300,
    "chunk_overlap": 50,
    "top_k": 3,
    "score_max": 1.2,     # FAISS 默认返回距离，越小越近
}

# 中文标点必须在列表里，否则递归切分会一路降级到按字数硬切
SEPARATORS = ["\n\n", "\n", "。", "；", "，", " ", ""]

PROMPT_TEMPLATE = """你是{role}。请只根据【已知信息】回答【问题】。

规则：
1. 只使用【已知信息】里出现过的内容，不允许补充任何外部知识。
2. 【已知信息】不足以回答时，回复「{refuse}」，不要推测。
3. 每个结论后面用方括号标出所依据的编号。

【已知信息】
{{context}}

【问题】
{{question}}""".format(role=DOMAIN["role"], refuse=DOMAIN["refuse"])

# TODO: 换行业时把这张表整张换掉。没有这张表就没有验收，只有感觉。
TESTSET = [
    ("我的快递出发地是哪？", "线路文档中的出发地字段"),
    ("这批货预计几天能到？", "线路文档中的运输时长"),
    ("货物现在到哪个分拨中心了？", "运单文档中的当前位置"),
    ("仓库里存的是什么货？", "仓储文档中的存储类型"),
]


# ------------------------------------------------------------ 建库
def build():
    """加载 → 切分 → 向量化 → 落盘。"""
    from langchain_community.document_loaders import (
        PyMuPDFLoader, TextLoader)
    from langchain_community.embeddings import OllamaEmbeddings
    from langchain_community.vectorstores import FAISS
    from langchain.text_splitter import RecursiveCharacterTextSplitter

    docs = []
    for name in sorted(os.listdir(PATHS["docs"])):
        path = os.path.join(PATHS["docs"], name)
        if name.lower().endswith(".pdf"):
            docs += PyMuPDFLoader(path).load()
        elif name.lower().endswith((".txt", ".md")):
            docs += TextLoader(path, encoding="utf-8").load()
    print("加载 %d 个文档片段" % len(docs))

    splitter = RecursiveCharacterTextSplitter(
        separators=SEPARATORS,
        chunk_size=PARAMS["chunk_size"],
        chunk_overlap=PARAMS["chunk_overlap"])
    chunks = splitter.split_documents(docs)
    print("切成 %d 块" % len(chunks))

    db = FAISS.from_documents(chunks, OllamaEmbeddings(model=MODELS["embed"]))
    db.save_local(PATHS["index"])
    print("索引已写入 %s（embed=%s）" % (PATHS["index"], MODELS["embed"]))


# ------------------------------------------------------------ 问答
def load_db():
    from langchain_community.embeddings import OllamaEmbeddings
    from langchain_community.vectorstores import FAISS
    return FAISS.load_local(PATHS["index"],
                            OllamaEmbeddings(model=MODELS["embed"]),
                            allow_dangerous_deserialization=True)


def retrieve(db, question):
    pairs = db.similarity_search_with_score(question, k=PARAMS["top_k"])
    return [(d, s) for d, s in pairs if s <= PARAMS["score_max"]]


def ask(question):
    from langchain_community.llms import Ollama

    db = load_db()
    pairs = retrieve(db, question)
    if not pairs:
        print(DOMAIN["refuse"])
        return

    context = "\n".join(
        "[%d] （%s）%s" % (i, os.path.basename(str(d.metadata.get("source", "未知"))),
                          d.page_content.strip())
        for i, (d, _s) in enumerate(pairs, 1))
    prompt = PROMPT_TEMPLATE.format(context=context, question=question)
    print(Ollama(model=MODELS["llm"]).invoke(prompt))
    print("\n检索到的块：")
    for i, (d, s) in enumerate(pairs, 1):
        print("  [%d] 距离 %.3f  %s" % (i, s, d.page_content.replace("\n", " ")[:50]))


# ------------------------------------------------------------ 验收
def evaluate():
    """跑测试集，人工核对命中的是不是「应该命中的那一块」。"""
    db = load_db()
    print("%-22s %-24s %s" % ("问题", "期望命中", "实际取回"))
    print("-" * 78)
    for question, expect in TESTSET:
        pairs = retrieve(db, question)
        got = pairs[0][0].page_content.replace("\n", " ")[:28] if pairs else "（空）"
        print("%-22s %-24s %s" % (question, expect, got))
    print("\n逐行判断实际取回的是不是期望的那一块，算出命中率再决定调什么。")


def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "help"
    if cmd == "build":
        build()
    elif cmd == "eval":
        evaluate()
    elif cmd == "ask":
        ask(sys.argv[2] if len(sys.argv) > 2 else "TODO: 写一个测试问题")
    else:
        print(__doc__)


if __name__ == "__main__":
    main()