向量化与切分：决定天花板的两步

"""手算余弦相似度，把「相似」这件事拆到只剩加减乘除。

三件事一次说清：
  1. cos = 点积 ÷ 两个模长，几何意义是夹角
  2. 向量归一化之后，点积就等于 cos —— 向量库常说的「内积 = 余弦」是这个意思
  3. 归一化之后，L2 距离与 cos 单调对应，排序结果完全一致

纯标准库，直接运行可看到全部中间量。
"""
import math

# 三个手写的三维向量，假装它们是三段话的 embedding
A = [0.60, 0.80, 0.00]   # 「快递从广州发出」
B = [0.30, 0.40, 0.00]   # 「包裹起运地是广州」——和 A 方向完全相同，只是长度一半
C = [0.00, 0.00, 1.00]   # 「仓库存的是电子产品」——和 A 完全正交


def dot(u, v):
    return sum(x * y for x, y in zip(u, v))


def norm(u):
    return math.sqrt(dot(u, u))


def cosine(u, v):
    return dot(u, v) / (norm(u) * norm(v))


def l2(u, v):
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(u, v)))


def normalize(u):
    n = norm(u)
    return [x / n for x in u]


def show(name, u):
    print("%s = [%s]   模长 |%s| = %.4f"
          % (name, ", ".join("%.4f" % x for x in u), name, norm(u)))


def main():
    print("=== 三个向量 ===")
    show("A", A)
    show("B", B)
    show("C", C)

    print("\n=== 逐步手算 cos(A, B) ===")
    print("点积 A·B = %.2f×%.2f + %.2f×%.2f + %.2f×%.2f = %.4f"
          % (A[0], B[0], A[1], B[1], A[2], B[2], dot(A, B)))
    print("模长 |A| = %.4f，|B| = %.4f" % (norm(A), norm(B)))
    print("cos = %.4f ÷ (%.4f × %.4f) = %.6f"
          % (dot(A, B), norm(A), norm(B), cosine(A, B)))
    print("→ A 和 B 方向完全一致，cos = 1，尽管 B 的长度只有 A 的一半。")

    print("\n=== 三组两两比较 ===")
    print("%-8s %10s %10s %10s" % ("组合", "点积", "余弦", "L2 距离"))
    for n1, v1, n2, v2 in [("A", A, "B", B), ("A", A, "C", C), ("B", B, "C", C)]:
        print("%-8s %10.4f %10.4f %10.4f"
              % ("%s-%s" % (n1, n2), dot(v1, v2), cosine(v1, v2), l2(v1, v2)))
    print("→ 点积把 A-B 算成 0.50，比 A 自己跟自己的 1.00 小，但它们明明是同一个方向。")
    print("  原始点积会被向量长度带偏，余弦不会。")

    print("\n=== 归一化之后，点积 == 余弦 ===")
    na, nb, nc = normalize(A), normalize(B), normalize(C)
    show("A'", na)
    show("B'", nb)
    for n1, v1, n2, v2 in [("A'", na, "B'", nb), ("A'", na, "C'", nc)]:
        d, c = dot(v1, v2), cosine(v1, v2)
        print("%s·%s = %.6f，cos = %.6f，相差 %.1e" % (n1, n2, d, c, abs(d - c)))
    assert abs(dot(na, nb) - cosine(A, B)) < 1e-12
    print("断言通过：归一化后用最快的点积算子，拿到的就是余弦。")

    print("\n=== 归一化之后，L2 和余弦排序一致 ===")
    print("恒等式：|u−v|² = 2 − 2·cos(u,v)   （u、v 都是单位向量）")
    for n1, v1, n2, v2 in [("A'", na, "B'", nb), ("A'", na, "C'", nc)]:
        left = l2(v1, v2) ** 2
        right = 2 - 2 * cosine(v1, v2)
        print("  %s vs %s：左 %.6f  右 %.6f  相差 %.1e" % (n1, n2, left, right, abs(left - right)))
        assert abs(left - right) < 1e-12
    print("断言通过：所以向量库里选 L2 还是内积，前提是向量有没有归一化。")
    print("没归一化就选内积，长文档会因为向量更长而系统性地排在前面——这是一类很隐蔽的召回偏差。")


if __name__ == "__main__":
    main()

"""切块实验：同一份资料，chunk_size 与 overlap 一变，能不能答全就变。

脚本用同一个问题、同一套相似度算法，只改切块参数，量两件事：
  A 库里到底存不存在一个「自身就能答全」的块   —— 这是切分决定的
  B 排在第一的那个块是不是这种块               —— 这是检索决定的
两件事必须分开量，否则召回不准时根本说不清该去调哪一边。

纯标准库，直接运行。
"""
import math
import re
from collections import Counter

DOC = (
    "速达物流公司总部设在北京市，主要业务范围包括国际快递和仓储管理两大板块。"
    "本次运单的货物编号为 ABC123456，发货日期是 2023 年 1 月 15 日，"
    "当前货物已经抵达上海分拨中心，预计到达日期为 2023 年 1 月 20 日。"
    "承运方为快运通，采用陆运方式，出发地是广州，目的地是重庆，预计运输时间三天。"
    "仓储方面，东方仓储中心位于深圳市，主要存储电子产品，采用常温仓储，当前库存量一千件。"
)

QUESTION = "这批货从哪里发出，大概几天能到？"


def fixed_split(text, size, overlap):
    """按固定字数硬切，块之间保留 overlap 个字符的重叠。"""
    if overlap >= size:
        raise ValueError("overlap 必须小于 chunk_size，否则窗口挪不动会死循环")
    out, start = [], 0
    while start < len(text):
        out.append(text[start:start + size])
        if start + size >= len(text):
            break
        start += size - overlap
    return out


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


def answers_fully(text):
    """一个块能不能独立答完这个问题：出发地 + 时长，缺一不可。"""
    return ("出发地是广州" in text) and ("运输时间三天" in text)


def rank(chunks, question):
    qv = embed(question)
    scored = [(cosine(qv, embed(c)), c) for c in chunks]
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored


def main():
    print("原文 %d 字，问题：%s\n" % (len(DOC), QUESTION))
    print("%-11s %-8s %-6s %-12s %-12s %-10s"
          % ("chunk_size", "overlap", "块数", "库里有完整块", "Top1 答得全", "Top1 分数"))
    print("-" * 66)

    for size, overlap in [(15, 0), (20, 5), (35, 0), (35, 7),
                          (50, 10), (80, 16), (120, 24), (300, 30)]:
        chunks = fixed_split(DOC, size, overlap)
        ranked = rank(chunks, QUESTION)
        exists = any(answers_fully(c) for c in chunks)
        top_ok = answers_fully(ranked[0][1])
        print("%-11d %-8d %-6d %-12s %-12s %-10.4f"
              % (size, overlap, len(chunks),
                 "有" if exists else "没有",
                 "答得全" if top_ok else "答不全",
                 ranked[0][0]))

    print()
    print("第一列到第四列是切分的锅，第五列是检索的锅。分开看：")
    print("  · 「库里有完整块 = 没有」 → 再怎么调 k、调阈值都没用，必须回去改切分。")
    print("  · 「有完整块但 Top1 答不全」 → 切分没问题，是检索排序的问题，该调 k 或换算法。")

    print()
    print("=" * 66)
    print("细看两组")
    for size, overlap in [(15, 0), (50, 10)]:
        chunks = fixed_split(DOC, size, overlap)
        ranked = rank(chunks, QUESTION)
        print("-" * 66)
        print("chunk_size=%d overlap=%d → %d 块" % (size, overlap, len(chunks)))
        full = [c for c in chunks if answers_fully(c)]
        print("  自身能答全的块：%d 个" % len(full))
        for score, text in ranked[:2]:
            mark = "✔" if answers_fully(text) else "✘"
            print("  %s %.4f 「%s」" % (mark, score, text))

    print("=" * 66)
    print("重叠到底救了什么：看刀口落在哪")
    for overlap in (0, 7):
        chunks = fixed_split(DOC, 35, overlap)
        # 找到「出发地是广州」这句话被切在哪两块之间
        idx = [i for i, c in enumerate(chunks) if "出发地" in c]
        print("  overlap=%-2d  含「出发地」的块共 %d 个：%s"
              % (overlap, len(idx), "、".join("第%d块" % (i + 1) for i in idx)))
        for i in idx:
            print("           第%d块：「%s」" % (i + 1, chunks[i]))
    print()
    print("结论：overlap 不是为了多存几份，是让被刀口切开的那句话")
    print("      在相邻两块里各留一份完整的，取到哪一块都读得通。")


if __name__ == "__main__":
    main()

"""中文切分的三个坑，每个都当场复现一遍。

坑一：默认分隔符是空行或空格，中文正文里根本没有 → 切不动
坑二：按字数切会把一句话拦腰斩断 → 块不成意思
坑三：chunk_size 的单位是字符还是 token，中英文不是一回事 → 估算失真

纯标准库，直接运行。
"""
import re

CN = (
    "速达物流公司总部设在北京市。本次运单货物编号为 ABC123456。"
    "当前货物已抵达上海分拨中心，预计到达日期为 2023 年 1 月 20 日。"
    "承运方为快运通，出发地是广州，目的地是重庆，预计运输时间三天。"
)

EN = ("Sudar Logistics is headquartered in Beijing. "
      "The shipment ABC123456 has arrived at the Shanghai hub. "
      "It departs from Guangzhou and arrives in Chongqing in three days.")


def split_by_separator(text, separator, size):
    """模拟「先按分隔符切，再按 size 聚合」这一类切分器的行为。"""
    pieces = text.split(separator) if separator else [text]
    out, buf = [], ""
    for p in pieces:
        candidate = (buf + separator + p) if buf else p
        if len(candidate) <= size:
            buf = candidate
        else:
            if buf:
                out.append(buf)
            buf = p
    if buf:
        out.append(buf)
    return out


def recursive_split(text, seps, size):
    """递归切分：按分隔符优先级依次尝试，切不动就降级到下一级。

    这正是 RecursiveCharacterTextSplitter 的思路，
    关键在于分隔符列表要为中文准备句号、分号、逗号这几级。
    """
    if len(text) <= size:
        return [text]
    for i, sep in enumerate(seps):
        if sep and sep in text:
            parts, out, buf = text.split(sep), [], ""
            for p in parts:
                cand = (buf + sep + p) if buf else p
                if len(cand) <= size:
                    buf = cand
                else:
                    if buf:
                        out.append(buf)
                    buf = p if len(p) <= size else ""
                    if not buf:
                        out += recursive_split(p, seps[i + 1:], size)
            if buf:
                out.append(buf)
            return out
    return [text[i:i + size] for i in range(0, len(text), size)]


def rough_tokens(text):
    """粗估 token 数：汉字按 1 个 token，英文按 4 个字符 1 个 token。

    只为说明量级差异，不是任何具体分词器的精确结果。
    """
    cjk = len(re.findall(r"[\u4e00-\u9fff]", text))
    rest = len(re.sub(r"[\u4e00-\u9fff]", "", text))
    return cjk + rest / 4


def main():
    print("=" * 62)
    print("坑一：默认按空行/空格切，中文切不动")
    print("=" * 62)
    print("原文 %d 字，里面有 %d 个空格、%d 个空行"
          % (len(CN), CN.count(" "), CN.count("\n\n")))
    for sep, name in [("\n\n", "空行"), (" ", "空格")]:
        got = split_by_separator(CN, sep, 50)
        print("  用「%s」切，size=50 → %d 块，最长一块 %d 字" % (name, len(got), max(len(g) for g in got)))
    print("  ↑ 用空行切：文本里一个空行都没有，整篇原样返回，远超 size。")
    print("    用空格切：确实切出了几块，但那 7 个空格全在英文单号和阿拉伯数字旁边，")
    print("    切点落在哪里完全是碰运气，跟句子边界没任何关系。")
    print("  英文对照：")
    got = split_by_separator(EN, " ", 50)
    print("    用「空格」切英文，size=50 → %d 块，最长 %d 字符" % (len(got), max(len(g) for g in got)))
    print("  ↑ 同一套默认参数，英文正常、中文失效。这就是照抄英文教程最容易踩的一脚。")

    print()
    print("=" * 62)
    print("坑二：按字数硬切，句子被拦腰斩断")
    print("=" * 62)
    hard = [CN[i:i + 30] for i in range(0, len(CN), 30)]
    for i, c in enumerate(hard[:4], 1):
        broken = not c.rstrip().endswith(("。", "，", "；"))
        print("  第%d块%s「%s」" % (i, "（断在半句）" if broken else "（收在标点）", c))
    print()
    print("  换成中文分隔符递归切，size=30：")
    good = recursive_split(CN, ["。", "；", "，", ""], 30)
    for i, c in enumerate(good[:5], 1):
        print("  第%d块「%s」" % (i, c))
    print("  ↑ 每块都收在标点上，块自身读得通。分隔符列表里必须有中文标点。")

    print()
    print("=" * 62)
    print("坑三：chunk_size 的单位，中英文差一倍还多")
    print("=" * 62)
    print("%-6s %-10s %-14s %-10s" % ("语言", "字符数", "粗估 token", "字符/token"))
    for name, text in [("中文", CN), ("英文", EN)]:
        ch, tk = len(text), rough_tokens(text)
        print("%-6s %-10d %-14.1f %-10.2f" % (name, ch, tk, ch / tk))
    print()
    print("  同样写 chunk_size=500：")
    print("    按字符算，中文一块约 500 token，英文一块约 125 token —— 差 4 倍。")
    print("    模型上下文与 Embedding 输入上限都是按 token 卡的，")
    print("    所以要先确认自己用的切分器数的是字符还是 token，再去定这个数。")


if __name__ == "__main__":
    main()

"""维度决定的三笔账：内存、磁盘、单次检索的计算量。

选 Embedding 模型时，「效果好不好」之外还有一个必须一起看的量：维度。
维度直接乘进内存、磁盘和检索耗时里，选之前先把账算出来。

表里的维度是各模型公开配置里的隐藏层宽度，不是估的。
纯标准库，直接运行。
"""

# (模型名, 维度, 单次可编码的最大长度, 说明)
MODELS = [
    ("bge-small-zh-v1.5", 512, 512, "中文小模型，维度最省"),
    ("m3e-base", 768, 512, "中文常用基线"),
    ("nomic-embed-text-v1.5", 768, 2048, "长文本友好"),
    ("bge-large-zh-v1.5", 1024, 512, "中文大模型"),
    ("mxbai-embed-large-v1", 1024, 512, "英文强，中文一般"),
    ("bge-m3", 1024, 8194, "多语言 + 超长输入"),
    ("text-embedding-3-small", 1536, 8191, "闭源 API，按 token 计费"),
    ("text-embedding-3-large", 3072, 8191, "闭源 API，维度最高"),
]

BYTES_PER_FLOAT32 = 4
BYTES_PER_INT8 = 1


def gib(n_bytes):
    return n_bytes / (1024 ** 3)


def mib(n_bytes):
    return n_bytes / (1024 ** 2)


def main():
    for n_chunks in (10_000, 1_000_000):
        print("=" * 74)
        print("库里有 %s 个 chunk 时，光是向量本身要占多少" % f"{n_chunks:,}")
        print("=" * 74)
        print("%-26s %-6s %-12s %-12s %-10s"
              % ("模型", "维度", "float32", "int8 量化", "省下"))
        print("-" * 74)
        for name, dim, _maxlen, _note in MODELS:
            f32 = n_chunks * dim * BYTES_PER_FLOAT32
            i8 = n_chunks * dim * BYTES_PER_INT8
            fmt = (lambda b: "%8.2f MB" % mib(b)) if n_chunks < 100_000 else (lambda b: "%8.2f GB" % gib(b))
            print("%-26s %-6d %-12s %-12s %-10s"
                  % (name, dim, fmt(f32), fmt(i8), "%.0f%%" % ((1 - i8 / f32) * 100)))
        print()

    print("=" * 74)
    print("检索一次要做多少次乘加：暴力扫描 = chunk 数 × 维度")
    print("=" * 74)
    print("%-26s %-6s %-18s %-18s" % ("模型", "维度", "1 万块", "100 万块"))
    print("-" * 74)
    for name, dim, _maxlen, _note in MODELS:
        print("%-26s %-6d %-18s %-18s"
              % (name, dim, f"{10_000 * dim:,} 次", f"{1_000_000 * dim:,} 次"))
    print()
    print("3072 维比 512 维贵 6 倍，内存、磁盘、计算三样一起贵 6 倍。")
    print("百万级别上暴力扫描已经不现实，所以向量库才要建近似索引")
    print("（HNSW、IVF 这类），用一点点召回率换回几十倍的速度。")

    print()
    print("=" * 74)
    print("还有一个容易忽略的上限：单块文本的最大长度")
    print("=" * 74)
    print("%-26s %-10s %-12s %s" % ("模型", "最大长度", "维度", "说明"))
    print("-" * 74)
    for name, dim, maxlen, note in MODELS:
        print("%-26s %-10d %-12d %s" % (name, maxlen, dim, note))
    print()
    print("chunk 超过这个长度不会报错，只会被悄悄截断——")
    print("后半段内容压根没进向量，检索自然找不回来，而且没有任何提示。")
    print("所以 chunk_size 的上限不只由模型上下文决定，也由 Embedding 的输入上限决定。")

    # 维度不是越高越好的反证：维度翻倍，成本确定翻倍，效果提升不保证
    print()
    print("选型顺序建议：先按语种和最大长度筛掉不合用的，")
    print("再在剩下的里面按维度挑最省的那个，最后才用自己的问题集实测效果。")


if __name__ == "__main__":
    main()

"""向量库选型：按自己的约束打分，而不是背一张「谁最好」的榜单。

六个候选都开源免费，差别不在功能多少，而在它们各自假设了什么部署形态。
把约束写成信号，让分数自己排出来；换了项目只改 need 字典。

star 数是排序的参考量之一，取自各仓库公开的星标数，
写死在这里只为可复现，真实选型时应当自己重新拉一次。
纯标准库，直接运行。
"""

# (名称, 星标数, 形态, 是否需要独立服务, 是否支持磁盘持久化, 是否支持元数据过滤, 是否支持分布式)
CANDIDATES = [
    ("FAISS",    40929, "库（进程内）",   False, True,  False, False),
    ("Chroma",   29325, "库 / 单机服务",  False, True,  True,  False),
    ("pgvector", 23062, "Postgres 扩展",  True,  True,  True,  False),
    ("Qdrant",   34660, "独立服务",       True,  True,  True,  True),
    ("Milvus",   46152, "独立服务",       True,  True,  True,  True),
    ("Weaviate", 16819, "独立服务",       True,  True,  True,  True),
]

FIELDS = ["name", "stars", "shape", "needs_server", "persist", "filter", "distributed"]


def to_dict(row):
    return dict(zip(FIELDS, row))


def score(cand, need):
    """按项目约束打分。每条约束要么加分要么直接出局。"""
    s, reasons = 0, []

    if need["no_extra_service"] and cand["needs_server"]:
        return None, ["要额外起服务，不满足「不加运维负担」"]
    if need["metadata_filter"] and not cand["filter"]:
        return None, ["不支持按元数据过滤，不满足「要按部门/时间筛」"]
    if need["distributed"] and not cand["distributed"]:
        return None, ["不支持分布式，撑不住目标数据量"]

    if not cand["needs_server"]:
        s += 3
        reasons.append("进程内即可，无额外运维 +3")
    if cand["filter"]:
        s += 2
        reasons.append("支持元数据过滤 +2")
    if cand["distributed"]:
        s += 2
        reasons.append("可水平扩展 +2")
    if need["reuse_existing_db"] and "Postgres" in cand["shape"]:
        s += 4
        reasons.append("直接复用现有 Postgres，数据与向量同库同事务 +4")

    # 使用量作为并列时的次级排序依据，权重刻意压低：
    # 星多只说明生态活跃，不代表更适合你这个项目。
    s += cand["stars"] / 20000
    reasons.append("生态活跃度 +%.2f" % (cand["stars"] / 20000))
    return s, reasons


SCENES = [
    ("单机小项目：几千个块，不想多起任何服务",
     {"no_extra_service": True, "metadata_filter": False,
      "distributed": False, "reuse_existing_db": False}),
    ("已有 Postgres 的业务系统，要按部门和时间过滤",
     {"no_extra_service": False, "metadata_filter": True,
      "distributed": False, "reuse_existing_db": True}),
    ("千万级向量，要水平扩展和在线扩容",
     {"no_extra_service": False, "metadata_filter": True,
      "distributed": True, "reuse_existing_db": False}),
]


def main():
    print("%-10s %-8s %-16s %-8s %-8s %-8s %-8s"
          % ("名称", "星标", "形态", "需起服务", "可持久化", "元数据过滤", "分布式"))
    print("-" * 74)
    for row in CANDIDATES:
        c = to_dict(row)
        print("%-10s %-8s %-16s %-8s %-8s %-8s %-8s"
              % (c["name"], f"{c['stars']:,}", c["shape"],
                 "是" if c["needs_server"] else "否",
                 "是" if c["persist"] else "否",
                 "是" if c["filter"] else "否",
                 "是" if c["distributed"] else "否"))

    for title, need in SCENES:
        print()
        print("=" * 74)
        print("场景：%s" % title)
        print("=" * 74)
        ranked, out = [], []
        for row in CANDIDATES:
            c = to_dict(row)
            s, reasons = score(c, need)
            (out if s is None else ranked).append((c["name"], s, reasons))
        ranked.sort(key=lambda x: x[1], reverse=True)

        for name, s, reasons in ranked:
            print("  %-10s %6.2f   %s" % (name, s, "；".join(reasons)))
        for name, _s, reasons in out:
            print("  %-10s %6s   出局：%s" % (name, "—", reasons[0]))
        if ranked:
            print("  → 首选：%s" % ranked[0][0])


if __name__ == "__main__":
    main()

"""离线建库骨架：加载 → 体检 → 切分 → 向量化 → 入库。

纯标准库即可跑通（用词袋顶替真实 Embedding），把 TODO 换成真实组件就是生产版本。
运行：
    python3 chunk_pipeline_skeleton.py
"""
import os
import re
import json
import math

# ---------------------------------------------------------------- 配置区
# TODO: 资料目录，先只放三五份做验证，别一上来全量灌
DOC_DIR = os.environ.get("KB_DOC_DIR", "./docs")

# TODO: 索引落盘位置
INDEX_PATH = os.environ.get("KB_INDEX", "./kb_index.json")

# TODO: 块大小。中文按字符算时 200~500 是常见区间，
#       且必须小于 Embedding 模型的最大输入（bge-large-zh-v1.5 是 512 token）
CHUNK_SIZE = int(os.environ.get("KB_CHUNK_SIZE", "300"))

# TODO: 重叠。一般取 chunk_size 的 10%~20%，必须小于 chunk_size
CHUNK_OVERLAP = int(os.environ.get("KB_CHUNK_OVERLAP", "50"))

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

# TODO: 换真实模型时改这里，建库与提问必须用同一个
EMBED_MODEL = os.environ.get("KB_EMBED_MODEL", "bag-of-words")


# ---------------------------------------------------------------- ① 加载
def load_docs(doc_dir):
    """返回 [(text, source)]。

    TODO: 换成真实加载器（PDF/Word/HTML）时，务必保住来源与页码——
    来源在这一步丢了，后面就再也标不出出处。
    """
    items = []
    if not os.path.isdir(doc_dir):
        # 没有目录时用一段内置样例，保证骨架能直接跑起来
        demo = ("速达物流公司总部设在北京市，业务范围包括国际快递与仓储管理。\n\n"
                "本次运单货物编号为 ABC123456，当前已抵达上海分拨中心，"
                "预计到达日期为 2023 年 1 月 20 日。\n\n"
                "承运方为快运通，采用陆运方式，出发地是广州，目的地是重庆，"
                "预计运输时间三天。\n\n"
                "东方仓储中心位于深圳市，存储电子产品，常温仓储，当前库存一千件。\n\n"
                "国际件报关所需材料包括商业发票、装箱单与贸易合同，"
                "缺任一项都会被海关退回重新申报。\n\n"
                "冷链货物需全程温控，交接时须记录温度并由双方签字确认。\n\n"
                "所有分拨中心均在每日 22:00 前完成当日件的装车发运。\n\n"
                "超重货物需提前一个工作日报备，并由发运网点安排专车，"
                "未报备的超重件不予插单。\n\n"
                "货物破损理赔按声明价值或实际损失较低者计算，"
                "理赔申请须在签收后七日内提出。")
        return [(demo, "内置样例")]
    for name in sorted(os.listdir(doc_dir)):
        path = os.path.join(doc_dir, name)
        if os.path.isfile(path) and name.endswith((".txt", ".md")):
            with open(path, encoding="utf-8") as f:
                items.append((f.read(), name))
    return items


# ---------------------------------------------------------------- ② 体检
def health_check(text, source):
    """入库前拦一道。返回 (是否通过, 问题列表)。"""
    problems = []
    if len(text.strip()) == 0:
        problems.append("抽出 0 个字，多半是扫描件，要先过 OCR")
    if text.count("\ufffd") > 0:
        problems.append("含 %d 个替换字符，编码有问题" % text.count("\ufffd"))
    longest = max((len(line) for line in text.splitlines()), default=0)
    if longest > CHUNK_SIZE * 4:
        problems.append("最长一行 %d 字，疑似表格被拍平" % longest)
    if "\n\n" not in text and not re.search(r"[。！？]", text):
        problems.append("既无空行也无中文句末标点，缺切分锚点")
    return (len(problems) == 0, problems)


# ---------------------------------------------------------------- ③ 切分
def split(text, size=CHUNK_SIZE, overlap=CHUNK_OVERLAP, seps=SEPARATORS):
    """递归切分：按分隔符从粗到细尝试，最后才按字数硬切。"""
    if overlap >= size:
        raise ValueError("overlap 必须小于 chunk_size，否则窗口挪不动")
    if len(text) <= size:
        return [text] if text.strip() else []

    for sep in seps:
        if sep == "":
            break
        if sep not in text:
            continue
        parts, buf, out = text.split(sep), "", []
        for i, p in enumerate(parts):
            piece = p + (sep if i < len(parts) - 1 else "")
            if len(buf) + len(piece) <= size:
                buf += piece
            else:
                if buf.strip():
                    out.append(buf)
                buf = (buf[-overlap:] if overlap and buf else "") + piece
        if buf.strip():
            out.append(buf)
        if all(len(c) <= size for c in out):
            return out

    # 所有分隔符都切不动，退到硬切
    step = size - overlap
    return [text[i:i + size] for i in range(0, len(text), step) if text[i:i + size].strip()]


# ---------------------------------------------------------------- ④ 向量化
def tokenize(s):
    """中文逐字拆，英文数字整词保留。"""
    return re.findall(r"[a-zA-Z0-9]+|[\u4e00-\u9fff]", s.lower())


def embed(text):
    """TODO: 换成真实 Embedding 模型，返回定长浮点列表。

    换模型后 similarity() 必须成对修改，否则分数毫无意义却照样能跑。
    """
    vec = {}
    for t in tokenize(text):
        vec[t] = vec.get(t, 0) + 1
    return vec


def similarity(a, b):
    """TODO: 换成标准余弦（点积 / 模长积），与 embed() 成对修改。"""
    common = set(a) & set(b)
    dot = sum(a[t] * b[t] for t 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


# ---------------------------------------------------------------- ⑤ 入库
def build_index(doc_dir=DOC_DIR, index_path=INDEX_PATH):
    """TODO: 换成真实向量库（FAISS / Chroma / pgvector）时只换这个函数。

    注意：向量用来找、原文用来读、来源用来核对，三样必须一起存。
    """
    records, skipped = [], []
    for text, source in load_docs(doc_dir):
        ok, problems = health_check(text, source)
        if not ok:
            skipped.append((source, problems))
            continue
        for i, chunk in enumerate(split(text)):
            records.append({
                "id": "%s#%d" % (source, i),
                "text": chunk,
                "source": source,
                "vector": embed(chunk),     # 真实向量库里这里是浮点列表
            })

    with open(index_path, "w", encoding="utf-8") as f:
        json.dump({"model": EMBED_MODEL, "records": records},
                  f, ensure_ascii=False)

    print("入库 %d 块，落盘 %s（Embedding：%s）"
          % (len(records), index_path, EMBED_MODEL))
    for source, problems in skipped:
        print("  跳过 %s：%s" % (source, "；".join(problems)))
    return records


def self_check(records, probe="出发地是哪里"):
    """建完立刻问一句已知答案，确认库建对了再往下做问答。"""
    if not records:
        print("库是空的，别往下做问答了。")
        return
    qv = embed(probe)
    ranked = sorted(records, key=lambda r: similarity(qv, r["vector"]), reverse=True)
    print("-" * 56)
    print("建库自检，问「%s」取回前 2 块：" % probe)
    for i, r in enumerate(ranked[:2], 1):
        print("  [%d] %.4f  %s" % (i, similarity(qv, r["vector"]),
                                   r["text"].replace("\n", " ")[:46]))
    print("逐条读一遍：取回的内容对不上题，就说明切分或模型有问题，先回去查这两处。")


if __name__ == "__main__":
    self_check(build_index())