检索与问答链：把该用的那块拿准

"""k 取多少：把「取几张卡片」这件事量出来，而不是拍脑袋。

三个指标一起看才有意义：
  命中率   —— 该用的那块在不在取回来的 k 张里（k 越大越高）
  噪声比   —— 取回来的 k 张里有多少是无关的（k 越大越差）
  上下文量 —— 拼进提示词的字数（k 越大越贵）
k 的取值就是在这三条曲线上找一个拐点。

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

CHUNKS = [
    "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。",
    "该批货物预计到达日期为 2023-01-20。",
    "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。",
    "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。",
    "速达物流总部在北京市，业务范围为国际快递与仓储管理。",
    "客服热线工作时间为每日 9:00 至 21:00。",
    "国际件报关所需材料包括发票、装箱单与合同。",
    "冷链货物需全程温控，交接时须记录温度。",
]

# 测试集：问题 → 必须命中的块下标（人工标注，这一步没有捷径）
TESTSET = [
    ("这批货从哪发出，几天能到？", {2}),
    ("货现在到哪了？", {0}),
    ("深圳仓库存的是什么？", {3}),
    ("预计什么时候送到？", {1}),
]


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 search(question, k):
    qv = embed(question)
    scored = [(cosine(qv, v), i, c) for i, c, v in INDEX]
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[:k]


def main():
    print("知识库 %d 块，测试集 %d 条\n" % (len(CHUNKS), len(TESTSET)))
    print("%-4s %-10s %-10s %-14s %s" % ("k", "命中率", "噪声比", "平均上下文", "说明"))
    print("-" * 64)

    rows = []
    for k in range(1, len(CHUNKS) + 1):
        hit = 0
        noise_num = noise_den = 0
        ctx_chars = 0
        for question, gold in TESTSET:
            got = search(question, k)
            idxs = {i for _s, i, _c in got}
            if gold & idxs:
                hit += 1
            noise_num += len(idxs - gold)
            noise_den += len(idxs)
            ctx_chars += sum(len(c) for _s, _i, c in got)
        rate = hit / len(TESTSET)
        noise = noise_num / noise_den
        avg_ctx = ctx_chars / len(TESTSET)
        rows.append((k, rate, noise, avg_ctx))
        note = ""
        if k > 1 and rows[-2][1] == rate:
            note = "命中率没涨，只是在加噪声"
        print("%-4d %-10.2f %-10.2f %-14.1f %s" % (k, rate, noise, avg_ctx, note))

    print()
    # 找到第一个命中率达到最高的 k：再往上加只剩坏处
    best = max(r[1] for r in rows)
    knee = next(k for k, rate, _n, _c in rows if rate >= best)
    print("命中率在 k=%d 时已经到顶（%.2f），再往上加 k：" % (knee, best))
    after = rows[knee - 1]
    last = rows[-1]
    print("  噪声比 %.2f → %.2f，平均上下文 %.0f 字 → %.0f 字，命中率一点没涨。"
          % (after[2], last[2], after[3], last[3]))
    print("  所以这个库上 k 取 %d 就够，不是越大越好。" % knee)

    print()
    print("=" * 64)
    print("单看一条问题，k 从 1 加到 4 都取回了什么")
    print("=" * 64)
    question, gold = TESTSET[0]
    print("问题：%s（标准答案是第 %s 块）\n" % (question, "、".join(str(i) for i in sorted(gold))))
    for score, i, c in search(question, 4):
        mark = "✔ 该用的" if i in gold else "✘ 噪声  "
        print("  %s  %.4f  [%d] %s" % (mark, score, i, c))
    print()
    print("注意第一条：分数最高的不一定是该用的那一条。")
    print("命中率要看的是「在不在这 k 条里」，而不是「是不是排第一」——")
    print("这也正是后面要上 Rerank 的理由。")


if __name__ == "__main__":
    main()

"""分数阈值：怎么让系统学会说「资料里没有」。

top-k 有个天生的毛病——它永远返回 k 条，哪怕库里一条相关的都没有。
问一个知识库根本覆盖不到的问题，照样会拿回 k 条最不相干的东西，
模型看到这些东西，就很容易顺着它们编。

阈值就是补这个洞的：低于某个相似度，宁可空手回来。
但阈值不能拍脑袋定，要看分数分布。

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

CHUNKS = [
    "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。",
    "该批货物预计到达日期为 2023-01-20。",
    "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。",
    "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。",
    "速达物流总部在北京市，业务范围为国际快递与仓储管理。",
    "客服热线工作时间为每日 9:00 至 21:00。",
]

# 库里答得出来的问题
IN_SCOPE = [
    "货现在到哪了？",
    "这批货几天能到？",
    "深圳仓库存的是什么？",
    "客服几点下班？",
]
# 库里根本没有的问题——这些必须被挡下来
OUT_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 = [(c, embed(c)) for c in CHUNKS]


def top_score(question):
    qv = embed(question)
    return max(cosine(qv, v) for _c, v in INDEX)


def main():
    print("=== 第一步：先看分数分布，别急着定阈值 ===\n")
    print("%-22s %-10s %s" % ("问题", "最高分", "类别"))
    print("-" * 52)
    in_scores, out_scores = [], []
    for q in IN_SCOPE:
        s = top_score(q)
        in_scores.append(s)
        print("%-22s %-10.4f %s" % (q, s, "库里有"))
    for q in OUT_SCOPE:
        s = top_score(q)
        out_scores.append(s)
        print("%-22s %-10.4f %s" % (q, s, "库里没有"))

    print()
    print("库里有   的最低分：%.4f" % min(in_scores))
    print("库里没有 的最高分：%.4f" % max(out_scores))
    gap = min(in_scores) - max(out_scores)
    if gap > 0:
        print("两类之间还有 %.4f 的空隙 → 阈值可以卡在这个空隙里。" % gap)
    else:
        print("两类分数重叠了 %.4f → 单靠阈值分不开，" % -gap)
        print("要么先改进切分和向量化，要么加一层判断，不要指望一个数字解决问题。")

    print()
    print("=== 第二步：扫一遍阈值，看两类错误怎么此消彼长 ===\n")
    print("%-8s %-14s %-14s %s" % ("阈值", "该答的漏答", "该挡的放过", "总错误"))
    print("-" * 52)
    best = None
    for i in range(0, 41):
        th = i / 100
        miss = sum(1 for s in in_scores if s < th)      # 该答的被挡了
        leak = sum(1 for s in out_scores if s >= th)    # 该挡的放过去了
        total = miss + leak
        if best is None or total < best[1]:
            best = (th, total, miss, leak)
        if i % 4 == 0:
            print("%-8.2f %-14d %-14d %d" % (th, miss, leak, total))

    print()
    print("总错误最少的阈值：%.2f（漏答 %d，放过 %d）" % (best[0], best[2], best[3]))
    print()
    print("两类错误的代价并不对等，要按业务来偏：")
    print("  · 对外客服、法务、医疗 → 宁可漏答也不能瞎答，阈值往高调。")
    print("  · 内部检索助手         → 宁可多给几条让人自己挑，阈值往低调。")
    print()
    print("⚠️ 阈值是跟着 Embedding 模型走的。换了模型，分数分布整个平移，")
    print("   旧阈值会瞬间失效——所以换模型之后这张表必须重跑一遍。")


if __name__ == "__main__":
    main()

"""召回与重排是两件事：一个负责别漏，一个负责排对。

向量检索是「拿问题的向量去跟每块的向量比一次」，两边各自编码、互不照面，
所以它快得能扫百万条，但也粗。重排是「把问题和这一块拼在一起再读一遍」，
读得细，但慢得多，只能用在已经缩小到几十条的候选上。

这里用一个可解释的打分函数代替真正的 Rerank 模型：
它会看问题里的关键词在这一块里是否成对出现、位置是否接近。
真实项目换成 Rerank 模型，主流程一行都不用改。

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

CHUNKS = [
    "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。",
    "该批货物预计到达日期为 2023-01-20。",
    "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。",
    "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。",
    "速达物流总部在北京市，业务范围为国际快递与仓储管理。",
    "客服热线工作时间为每日 9:00 至 21:00。",
    "国际件报关所需材料包括发票、装箱单与合同。",
    "冷链货物需全程温控，交接时须记录温度。",
]

QUESTION = "这批货从哪发出，几天能到？"
GOLD = 2  # 人工标注：只有第 2 块能答全


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)]

# 问题里真正承载意图的词组。真实 Rerank 模型是学出来的，
# 这里手写出来是为了让「重排凭什么改排序」看得见。
INTENT_GROUPS = [
    ["出发", "发出", "起运", "出发地", "从哪"],
    ["几天", "时长", "运输时间", "多久"],
]


def recall(question, k):
    """第一段：向量召回，粗选 k 条。"""
    qv = embed(question)
    scored = [(cosine(qv, v), i, c) for i, c, v in INDEX]
    scored.sort(key=lambda x: x[0], reverse=True)
    return scored[:k]


def rerank_score(question, chunk):
    """第二段：把问题和这一块放在一起读，按「意图组覆盖数」打分。

    关键差别：召回是逐词比对，一块里凑够几个高频字就能拿分；
    重排要求每一组意图都在这一块里落地，凑字数没用。
    """
    covered = 0
    for group in INTENT_GROUPS:
        if any(w in chunk for w in group):
            covered += 1
    # 覆盖到的意图组越多越好；全覆盖时额外奖励，因为它能独立答全
    score = covered / len(INTENT_GROUPS)
    if covered == len(INTENT_GROUPS):
        score += 0.5
    return score


def main():
    print("问题：%s" % QUESTION)
    print("标准答案：第 %d 块「%s」\n" % (GOLD, CHUNKS[GOLD]))

    k = 5
    cands = recall(QUESTION, k)
    print("=" * 66)
    print("① 召回：向量检索粗选 %d 条（快，但排序粗）" % k)
    print("=" * 66)
    for rank, (score, i, c) in enumerate(cands, 1):
        mark = "← 该用的" if i == GOLD else ""
        print("  第%d名  %.4f  [%d] %s %s" % (rank, score, i, c[:30], mark))
    gold_rank = next(r for r, (_s, i, _c) in enumerate(cands, 1) if i == GOLD)
    print("\n  该用的那一块排在第 %d 名。如果直接取 top-2，它根本进不了提示词。" % gold_rank)

    print()
    print("=" * 66)
    print("② 重排：只对这 %d 条重新打分排序（慢，但排得准）" % k)
    print("=" * 66)
    reranked = [(rerank_score(QUESTION, c), old_score, i, c) for old_score, i, c in cands]
    reranked.sort(key=lambda x: x[0], reverse=True)
    print("  %-6s %-10s %-10s %s" % ("新名次", "重排分", "原召回分", "内容"))
    for rank, (rs, os_, i, c) in enumerate(reranked, 1):
        mark = "← 该用的" if i == GOLD else ""
        print("  %-6d %-10.3f %-10.4f [%d] %s %s" % (rank, rs, os_, i, c[:26], mark))
    new_rank = next(r for r, (_rs, _os, i, _c) in enumerate(reranked, 1) if i == GOLD)
    print("\n  该用的那一块从第 %d 名升到第 %d 名。现在取 top-2 就能拿到它。"
          % (gold_rank, new_rank))

    print()
    print("=" * 66)
    print("③ 重排救不回来的情况：召回阶段就没取到它")
    print("=" * 66)
    tight = recall(QUESTION, 2)
    print("  召回只取 2 条时，候选是：")
    for score, i, c in tight:
        print("    %.4f  [%d] %s" % (score, i, c[:30]))
    got = {i for _s, i, _c in tight}
    print("  该用的第 %d 块%s在候选里。" % (GOLD, "" if GOLD in got else "不"))
    if GOLD not in got:
        print("  重排只能对这 2 条重新排序，第 %d 块压根没进来 —— 它变不出新卡片。" % GOLD)
        print("  所以工程上的固定配比是：召回放宽（20~50 条），重排收紧（留 3~5 条）。")


if __name__ == "__main__":
    main()

"""混合检索：向量找得到意思，关键词找得到编号。

向量检索有一类天生的短板：单号、型号、法条编号、人名这种
「字面必须一模一样」的东西，它给不出可靠的区分度——
因为这些字符串在语义空间里彼此非常接近。

解法不是换更大的 Embedding，而是再加一路关键词检索，
两路结果按名次融合。这就是工程上说的 hybrid search。

融合用 RRF（Reciprocal Rank Fusion）：只看名次不看分数，
因为两路的分数量纲根本不可比。

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

CHUNKS = [
    "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。",
    "货物编号 ABC123789，发货日期 2023-01-16，当前位置广州分拨中心。",
    "货物编号 XYZ987654，发货日期 2023-01-17，当前位置成都分拨中心。",
    "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。",
    "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。",
    "所有分拨中心均在每日 22:00 前完成当日件的装车发运。",
]

QUESTION = "ABC123789 这单现在在哪？"
GOLD = 1


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 vector_search(q, k):
    qv = embed(q)
    scored = [(cosine(qv, v), i) for i, _c, v in INDEX]
    scored.sort(key=lambda x: x[0], reverse=True)
    return [i for _s, i in scored[:k]]


def keyword_search(q, k):
    """关键词检索：只看「问题里的实词有没有原样出现在这一块里」。

    这里刻意只取长度 ≥ 2 的字母数字串和二字以上中文串，
    单个汉字区分度太低，放进来会把关键词检索也变成词频游戏。
    """
    terms = [t for t in re.findall(r"[a-z0-9]{2,}", q.lower())]
    terms += re.findall(r"[\u4e00-\u9fff]{2,}", q)
    scored = []
    for i, c, _v in INDEX:
        low = c.lower()
        hit = sum(1 for t in terms if t in low)
        if hit:
            scored.append((hit, i))
    scored.sort(key=lambda x: x[0], reverse=True)
    return [i for _s, i in scored[:k]]


def rrf(rank_lists, k_const=60):
    """RRF 融合：每条结果的得分 = Σ 1/(k_const + 名次)。

    只吃名次，不吃分数——所以余弦分和关键词命中数不需要归一化，
    这正是它在工程上被广泛采用的原因。
    """
    score = {}
    for lst in rank_lists:
        for rank, idx in enumerate(lst, 1):
            score[idx] = score.get(idx, 0.0) + 1.0 / (k_const + rank)
    return sorted(score.items(), key=lambda kv: kv[1], reverse=True)


def show(name, idxs):
    print("%s：" % name)
    if not idxs:
        print("    （空）")
    for rank, i in enumerate(idxs, 1):
        mark = " ← 该用的" if i == GOLD else ""
        print("    第%d名 [%d] %s%s" % (rank, i, CHUNKS[i][:32], mark))


def main():
    print("问题：%s" % QUESTION)
    print("标准答案：第 %d 块\n" % GOLD)

    v = vector_search(QUESTION, 3)
    k = keyword_search(QUESTION, 3)
    show("① 只用向量检索", v)
    print()
    show("② 只用关键词检索", k)

    print()
    print("向量检索把三条运单排在一起，因为它们的句式几乎一模一样——")
    print("真正区分它们的是单号那串字符，而单号恰恰是语义模型最不擅长的东西。")
    print("关键词检索正好相反：它对单号一击即中，但换个说法就全线扑空。")

    print()
    print("③ RRF 融合两路结果：")
    fused = rrf([v, k])
    for rank, (i, s) in enumerate(fused, 1):
        mark = " ← 该用的" if i == GOLD else ""
        print("    第%d名 %.5f [%d] %s%s" % (rank, s, i, CHUNKS[i][:32], mark))
    top1 = fused[0][0]
    print()
    print("融合后排第一的是第 %d 块，%s" % (top1, "正确。" if top1 == GOLD else "仍需调权重。"))

    print()
    print("=" * 62)
    print("反过来验一次：换一个纯语义的问题，关键词检索会失灵")
    print("=" * 62)
    q2 = "货从哪个城市发出，路上要多久？"
    v2, k2 = vector_search(q2, 3), keyword_search(q2, 3)
    print("问题：%s" % q2)
    show("  向量检索", v2)
    show("  关键词检索", k2)
    print()
    print("结论：两路各有各的盲区，所以工程上默认两路都开、用名次融合，")
    print("      而不是二选一。")


if __name__ == "__main__":
    main()

"""把检索结果拼成提示词：五条必须做到的事，逐条演示。

1. 每块带编号和来源，答案才能标出处
2. 明确圈定「只用已知信息」，并给出兜底话术
3. 控制总长度，超了要截断而不是硬塞
4. 去重，同一段话被切进两块时不要重复两遍
5. 空结果要走另一条分支，不能拿空 context 去问模型

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

# 检索回来的结果：(分数, 来源, 正文)
HITS = [
    (0.81, "物流信息.pdf#p1", "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。"),
    (0.74, "物流信息.pdf#p1", "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。"),
    (0.71, "线路手册.docx#p3", "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。"),
    (0.42, "仓储制度.md#p2", "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。"),
]

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

SYSTEM = """你是物流信息查询助手。

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

USER = """【已知信息】
{context}

【问题】
{question}"""


def dedup(hits):
    """按正文去重，保留分数最高的那条，并把其它来源并进去。

    同一段话出现在两份文档里是很常见的（制度被复制到多个手册）。
    不去重的后果：提示词里同一句话出现两遍，既浪费 token，
    又会让模型误以为这件事被「反复强调」，从而过度加权。
    """
    seen = {}
    for score, source, text in sorted(hits, key=lambda x: x[0], reverse=True):
        key = text.strip()
        if key in seen:
            seen[key]["sources"].append(source)
        else:
            seen[key] = {"score": score, "text": key, "sources": [source]}
    return list(seen.values())


def build_context(items, budget_chars):
    """按分数从高到低装，装不下就停，不截断句子中间。"""
    lines, used, dropped = [], 0, 0
    for i, item in enumerate(items, 1):
        line = "[%d] （来源：%s）%s" % (i, "、".join(item["sources"]), item["text"])
        if used + len(line) > budget_chars:
            dropped += 1
            continue
        lines.append(line)
        used += len(line)
    return "\n".join(lines), used, dropped


def build_messages(question, hits, budget_chars=200):
    items = dedup(hits)
    context, used, dropped = build_context(items, budget_chars)
    if not context:
        # 空结果必须走另一条分支：拿空 context 去问模型，等于请它自由发挥
        return None, {"reason": "no_hit", "used": 0, "dropped": dropped}
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": USER.format(context=context, question=question)},
    ]
    return messages, {"reason": "ok", "used": used, "dropped": dropped,
                      "kept": len(items) - dropped}


def main():
    print("=== 原始检索结果 %d 条 ===" % len(HITS))
    for score, source, text in HITS:
        print("  %.2f  %-20s %s" % (score, source, text))

    print()
    print("=== 去重后 ===")
    items = dedup(HITS)
    for i, it in enumerate(items, 1):
        print("  [%d] %.2f  来源 %s" % (i, it["score"], "、".join(it["sources"])))
    print("  %d 条 → %d 条，第 1 条和第 3 条正文完全相同，来源已合并。"
          % (len(HITS), len(items)))

    print()
    print("=== 预算充足时拼出来的提示词 ===")
    messages, info = build_messages(QUESTION, HITS, budget_chars=500)
    for m in messages:
        print("-" * 60)
        print("role = %s" % m["role"])
        print(m["content"])
    print("-" * 60)
    print("context 用掉 %d 字，丢弃 %d 条" % (info["used"], info["dropped"]))

    print()
    print("=== 预算很紧时会发生什么 ===")
    for budget in (500, 120, 60, 10):
        _msgs, info = build_messages(QUESTION, HITS, budget_chars=budget)
        print("  预算 %3d 字 → 用掉 %3d 字，丢弃 %d 条，结果：%s"
              % (budget, info["used"], info["dropped"], info["reason"]))
    print("  预算被压到装不下任何一条时，返回的是 no_hit 而不是一个空 context —— ")
    print("  这两种情况在代码里必须分开处理。")

    print()
    print("=== 检索什么都没找到时 ===")
    msgs, info = build_messages("公司年假怎么休？", [], budget_chars=500)
    print("  messages = %s，reason = %s" % (msgs, info["reason"]))
    print("  这一步直接回复用户「资料中未提及」，连模型都不用调：")
    print("  省一次调用，也断掉了空 context 触发编造的那条路。")


if __name__ == "__main__":
    main()

"""问答链骨架：召回 → 阈值 → 重排 → 拼词 → 生成 → 带出处返回。

把 TODO 填掉就是一条可用的问答链。六个环节各自独立，
每个都能单独打开日志、单独换实现，不互相纠缠。

用法
    python3 qa_chain_skeleton.py "你的问题"
"""
import math
import os
import re
import sys
from collections import Counter

# ----------------------------------------------------------- 可调参数
RECALL_K = 20      # 召回放宽：宁可多拿，别漏
RERANK_K = 4       # 重排收紧：真正进提示词的条数
SCORE_MIN = 0.10   # TODO: 先跑一遍分数分布再定这个值
CTX_BUDGET = 1500  # context 的字数预算


# ----------------------------------------------------------- ① 索引
class Store:
    """TODO: 换成真正的向量库，保持 add / search 两个方法签名不变。"""

    def __init__(self):
        self.rows = []

    def add(self, text, source):
        self.rows.append({"text": text, "source": source, "vec": embed(text)})

    def search(self, query, k):
        qv = embed(query)
        scored = [(cosine(qv, r["vec"]), r) for r in self.rows]
        scored.sort(key=lambda x: x[0], reverse=True)
        return scored[:k]


def embed(text):
    """TODO: 换成 Embedding 模型。建库与提问必须用同一个。"""
    return Counter(re.findall(r"[a-z0-9]+", text.lower())
                   + re.findall(r"[\u4e00-\u9fff]", text))


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 apply_threshold(hits, score_min=SCORE_MIN):
    """低于阈值的一律丢掉，宁可空手回来也不要把噪声塞给模型。"""
    return [(s, r) for s, r in hits if s >= score_min]


# ----------------------------------------------------------- ③ 重排
def rerank(question, hits, keep=RERANK_K):
    """TODO: 换成 Rerank 模型；它只重排，不新增候选。

    没有 Rerank 模型时，原样返回前 keep 条即可，
    这一层是可插拔的，不接也能跑。
    """
    return hits[:keep]


# ----------------------------------------------------------- ④ 拼词
SYSTEM = """你是知识库问答助手。

规则：
1. 只使用【已知信息】里的内容作答，不得补充外部知识。
2. 信息不足时回复「资料中未提及」，不要推测。
3. 每个结论后用方括号标出依据的编号。"""

USER = "【已知信息】\n{context}\n\n【问题】\n{question}"


def build_messages(question, hits, budget=CTX_BUDGET):
    seen, items = set(), []
    for _s, r in hits:                       # 去重：同一段话只留一份
        if r["text"] in seen:
            continue
        seen.add(r["text"])
        items.append(r)

    lines, used = [], 0
    for i, r in enumerate(items, 1):
        line = "[%d] （来源：%s）%s" % (i, r["source"], r["text"])
        if used + len(line) > budget:
            break
        lines.append(line)
        used += len(line)

    if not lines:
        return None, []
    return ([{"role": "system", "content": SYSTEM},
             {"role": "user", "content": USER.format(context="\n".join(lines),
                                                     question=question)}],
            items[:len(lines)])


# ----------------------------------------------------------- ⑤ 生成
def call_llm(messages):
    """TODO: 换成真实模型调用，密钥走环境变量：

        api_key = os.environ.get("LLM_API_KEY")
    """
    if os.environ.get("LLM_API_KEY"):
        return "[TODO: 在这里发起真实请求]"
    return "[未接入模型] 本应发出的 messages 共 %d 条" % len(messages)


# ----------------------------------------------------------- ⑥ 出口
def answer(store, question):
    hits = store.search(question, RECALL_K)
    hits = apply_threshold(hits)
    if not hits:
        return {"answer": "资料中未提及。", "sources": [], "stage": "threshold"}

    hits = rerank(question, hits)
    messages, used = build_messages(question, hits)
    if messages is None:
        return {"answer": "资料中未提及。", "sources": [], "stage": "budget"}

    text = call_llm(messages)
    return {"answer": text,
            "sources": [r["source"] for r in used],
            "stage": "ok"}


DEMO_DOCS = [
    ("承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。", "线路手册#p3"),
    ("货物编号 ABC123456，当前位置上海分拨中心。", "运单#p1"),
    ("东方仓储中心位于深圳市，存储电子产品。", "仓储制度#p2"),
]


def main():
    question = sys.argv[1] if len(sys.argv) > 1 else "这批货从哪发出，几天能到？"
    store = Store()
    for text, source in DEMO_DOCS:      # TODO: 换成自己的建库流程
        store.add(text, source)

    result = answer(store, question)
    print("问题：%s" % question)
    print("阶段：%s" % result["stage"])
    print("出处：%s" % ("、".join(result["sources"]) or "无"))
    print("答案：%s" % result["answer"])


if __name__ == "__main__":
    main()

"""召回不准怎么查：按固定顺序逐层排除，不要一上来就改提示词。

排查顺序是有讲究的——每一层都建立在下一层成立的前提上。
顺序颠倒的典型后果：切分根本没把答案留在任何一块里，
却花两天时间在调提示词措辞，怎么调都没用。

这个脚本把排查过程做成可执行的检查表：给它一个问题和期望答案，
它会逐层告诉你问题出在哪一层、下一步该动哪个旋钮。

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

CHUNKS = [
    "货物编号 ABC123456，发货日期 2023-01-15，当前位置上海分拨中心。",
    "该批货物预计到达日期为 2023-01-20。",
    "承运方快运通，陆运，出发地广州，目的地重庆，预计运输时间三天。",
    "东方仓储中心位于深圳市，存储电子产品，常温仓储，库存一千件。",
    "速达物流总部在北京市，业务范围为国际快递与仓储管理。",
]


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 search(q, k, score_min=0.0):
    qv = embed(q)
    scored = [(cosine(qv, v), i, c) for i, c, v in INDEX]
    scored.sort(key=lambda x: x[0], reverse=True)
    return [(s, i, c) for s, i, c in scored[:k] if s >= score_min]


def diagnose(question, must_words, k=3, score_min=0.0):
    """按固定顺序逐层检查，返回 (出问题的层, 该动哪个旋钮)。"""
    steps = []

    # 第 1 层：库里到底有没有这句话
    in_corpus = [i for i, c in enumerate(CHUNKS) if all(w in c for w in must_words)]
    steps.append(("① 资料里有没有", bool(in_corpus),
                  "资料本身缺这条 → 去补文档，别调任何参数"))
    if not in_corpus:
        return steps

    # 第 2 层：这句话是不是完整地落在某一块里
    steps.append(("② 有没有落在同一块里", True,
                  "被切散了 → 调大 chunk_size 或改分隔符"))

    # 第 3 层：这一块在不在 top-k 里
    got = search(question, k, 0.0)
    idxs = [i for _s, i, _c in got]
    in_topk = any(i in idxs for i in in_corpus)
    steps.append(("③ 在不在 top-%d 里" % k, in_topk,
                  "召回没取到 → 调大 k、开混合检索、或换 Embedding 模型"))
    if not in_topk:
        return steps

    # 第 4 层：会不会被阈值挡掉
    kept = search(question, k, score_min)
    kept_idx = [i for _s, i, _c in kept]
    pass_th = any(i in kept_idx for i in in_corpus)
    steps.append(("④ 过不过阈值 %.2f" % score_min, pass_th,
                  "被阈值挡了 → 调低阈值，或先确认分数分布再定"))
    if not pass_th:
        return steps

    # 第 5 层：排名靠不靠前
    rank = next(r for r, (_s, i, _c) in enumerate(kept, 1) if i in in_corpus)
    steps.append(("⑤ 排名第 %d" % rank, rank <= 2,
                  "排太后 → 上 Rerank，这一层才轮得到重排出场"))
    if rank > 2:
        return steps

    steps.append(("⑥ 前五层都过了", True,
                  "到这里还答不对，才轮到改提示词模板"))
    return steps


CASES = [
    ("这批货从哪发出，几天能到？", ["出发地", "运输时间"]),
    ("货物现在到哪了？", ["当前位置"]),
    ("公司年假怎么休？", ["年假"]),
]


def main():
    for question, must in CASES:
        print("=" * 66)
        print("问题：%s" % question)
        print("期望命中含有这些字样的块：%s" % "、".join(must))
        print("=" * 66)
        steps = diagnose(question, must, k=3, score_min=0.10)
        for name, ok, advice in steps:
            flag = "通过" if ok else "卡住"
            print("  %-22s %s" % (name, flag))
            if not ok:
                print("      → %s" % advice)
        if all(ok for _n, ok, _a in steps):
            print("      → %s" % steps[-1][2])
        print()

    print("=" * 66)
    print("这张检查表的意义：每一层都有自己该动的旋钮，且顺序不能乱。")
    print("常见的错误做法是跳过 ①②，直接在 ③④⑤ 上反复调参——")
    print("资料里压根没有的东西，调到天亮也召不回来。")


if __name__ == "__main__":
    main()