PlantAI-rag-new/sdk/app/services/test_privateqa_pick.py

"""test_privateqa_pick.py —— _pick_items 先按 doc_id 聚合再取块（缺陷 #52c，验收 §6 第 3 条）。

运行：cd sdk && /home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest app/services/test_privateqa_pick.py -v
（venv = plantai-rag-new-venv；sdk 的 app/ 布局与 pyproject testpaths 一致）

构造合成数据（不真跑 FAISS）：
- 4 个 doc_id 各若干「作者页块」（高分、分接近）+ 1 条低分「正文块」。
- 旧实现会把高分的作者页块占满 top-8、正文块进不了上下文；
  新实现按 doc 级 max score 取 top 文档 → 各 doc 只占 ≤3 块，
  正文块（本组内最高低分块）能进入 picked。
"""
from __future__ import annotations

from app.services.privateqa import (
    PRIVATE_QA_PER_DOC_MAX,
    _build_context,
    _pick_items,
)


def make_items(doc_specs):
    """doc_specs: [(doc_id, [(score, snippet), ...]), ...] → item 列表（输入序打散，
    模拟 vstore 全局 score 序）。"""
    items = []
    for doc_id, chunks in doc_specs:
        for score, snippet in chunks:
            items.append({"doc_id": doc_id, "score": score, "snippet": snippet, "doc_name": f"资料{doc_id}"})
    return sorted(items, key=lambda it: -it["score"])


def test_dup_docs_no_longer_monopolize_topk():
    """#52c 主场景：4 个 doc_id 各 200 条近同分作者页块（同 doc 内内容相同，
    模拟重复上传残留——同一页面被反复切块/复传）+ 正文块 →
    选出的 items 覆盖多个 doc 且含正文块。"""
    specs = []
    for i in range(4):
        author = [(0.94 - i * 0.001, f"作者页块 {i}") for _ in range(200)]
        body = [(0.30, f"正文块 {i} 方法与结论")]
        specs.append((f"doc-{i}", author + body))
    items = make_items(specs)
    picked = _pick_items(items, top_k=8)

    doc_ids = [it["doc_id"] for it in picked]
    assert len(set(doc_ids)) >= 2, f"重复文档仍占满名额：{doc_ids}"
    for doc_id in set(doc_ids):
        assert doc_ids.count(doc_id) <= PRIVATE_QA_PER_DOC_MAX
    snippets = [it["snippet"] for it in picked]
    assert any("正文块" in s for s in snippets), "正文块未进入上下文"
    # 总分 ≤ top_k；块按分数降序
    assert len(picked) <= 8
    scores = [it["score"] for it in picked]
    assert scores == sorted(scores, reverse=True)


def test_normal_case_semantics_preserved():
    """回归：正常场景 top_k/每 doc 上限语义不变。"""
    items = [
        {"doc_id": "A", "score": 0.95, "snippet": "a1"},
        {"doc_id": "A", "score": 0.90, "snippet": "a2"},
        {"doc_id": "B", "score": 0.85, "snippet": "b1"},
        {"doc_id": "C", "score": 0.80, "snippet": "c1"},
        {"doc_id": "D", "score": 0.75, "snippet": "d1"},
    ]
    picked = _pick_items(items, top_k=5)
    assert len(picked) == 5
    for doc_id in {it["doc_id"] for it in picked}:
        assert sum(1 for it in picked if it["doc_id"] == doc_id) <= PRIVATE_QA_PER_DOC_MAX
    # 每 doc 内按分数降序取 ≤3 块
    a = [it for it in picked if it["doc_id"] == "A"]
    assert [it["snippet"] for it in a] == ["a1", "a2"]


def test_single_doc_caps_at_per_doc_max():
    """回归：单一 doc 大量块时仍 ≤ per-doc 上限。"""
    items = [{"doc_id": "X", "score": 0.9 - i * 0.001, "snippet": f"x{i}"} for i in range(20)]
    picked = _pick_items(items, top_k=8)
    assert 1 <= len(picked) <= PRIVATE_QA_PER_DOC_MAX


def test_ties_are_stable():
    """回归：同分并列时文档选择稳定（按输入首现顺序），结果确定。"""
    items = make_items(
        [
            ("A", [(0.9, "a1"), (0.9, "a2")]),
            ("B", [(0.9, "b1")]),
            ("C", [(0.9, "c1")]),
        ]
    )
    p1 = _pick_items(items, top_k=4)
    p2 = _pick_items(items, top_k=4)
    assert [it["doc_id"] for it in p1] == [it["doc_id"] for it in p2]
    assert len(p1) == 4


def test_malformed_items_ignored():
    """回归：非 dict / 缺 doc_id 条目不炸。"""
    picked = _pick_items(
        [None, "bad", {"score": 0.9, "snippet": "s"}, {"doc_id": "Z", "score": 0.5, "snippet": "z"}],
        top_k=4,
    )
    assert picked


def test_build_context_integration():
    """回归：_build_context 12k 预算截断 + [资料n] 引用编号与 items 顺序一致。"""
    picked = _pick_items(
        [
            {"doc_id": "A", "score": 0.95, "snippet": "正文一"},
            {"doc_id": "A", "score": 0.90, "snippet": "正文二"},
            {"doc_id": "B", "score": 0.80, "snippet": "正文三"},
        ],
        top_k=3,
    )
    context, out_items = _build_context(picked)
    assert "[资料1]" in context
    assert out_items[0]["ref"] == 1
    assert [it["snippet"] for it in out_items] == [it["snippet"] for it in picked]


async def test_ask_private_fetches_wider_pool(monkeypatch):
    """#52c：ask_private 以 top_k×4 向 vstore 取候选池，再收敛到 top_k。"""
    import app.services.privateqa as pq

    calls = {}

    async def fake_search(query, *, top_k, min_score, timeout):
        calls["requested_top_k"] = top_k
        items = []
        for doc_i in range(4):
            for j in range(200):
                items.append(
                    {
                        "doc_id": f"doc-{doc_i}",
                        "score": 0.94 - doc_i * 0.001,
                        # 同 doc 内作者页块内容相同（重复上传残留的真实形态）
                        "snippet": f"作者页块 {doc_i}",
                        "doc_name": f"资料{doc_i}",
                    }
                )
            items.append(
                {
                    "doc_id": f"doc-{doc_i}",
                    "score": 0.30,
                    "snippet": f"正文块 {doc_i} 方法与结论",
                    "doc_name": f"资料{doc_i}",
                }
            )
        return {"ok": True, "mode": "vector", "items": items}

    async def fake_llm(system, user, **kwargs):
        return "方法与结论见[资料1]。"

    monkeypatch.setattr(pq, "search_chunks", fake_search)
    monkeypatch.setattr(pq, "complete_text", fake_llm)

    result = await pq.ask_private("总结这篇论文的主要方法和结论", top_k=8)
    assert calls["requested_top_k"] == 20, "候选池应取 min(top_k×4, vstore 上限 20)"
    assert result["hit"] is True
    assert result["mode"] == "vector"
    doc_ids = [it["doc_id"] for it in result["items"]]
    assert len(set(doc_ids)) >= 2
    assert any("正文块" in (it["snippet"] or "") for it in result["items"])
    assert len(result["items"]) <= 8
    assert all(it["ref"] >= 1 for it in result["items"])


async def test_ask_private_pool_capped_at_20(monkeypatch):
    """#52c：候选池 = min(top_k×4, 20)（vstore CHUNK_TOP_K_MAX clamp）。"""
    import app.services.privateqa as pq

    calls = {}

    async def fake_search(query, *, top_k, min_score, timeout):
        calls["requested_top_k"] = top_k
        return {"ok": True, "mode": "vector", "items": []}

    monkeypatch.setattr(pq, "search_chunks", fake_search)
    await pq.ask_private("问题", top_k=8)
    assert calls["requested_top_k"] == 20
    await pq.ask_private("问题", top_k=5)
    assert calls["requested_top_k"] == 20  # 5×4=20
    await pq.ask_private("问题", top_k=3)
    assert calls["requested_top_k"] == 12  # 3×4=12（低于 20）