deerflow-code/offline-backend-20260512/backend/tests/test_llmwiki_title_ranking.py

from __future__ import annotations

from app.gateway.llmwiki_rag import _rrf_merge_source_channels
from deerflow.integrations.weknora.local_index.ranking import title_match_strength


def test_title_phrase_is_stronger_than_unrelated_person_title() -> None:
    assert title_match_strength("海湖庄园", "海湖庄园") > title_match_strength(
        "海湖庄园", "特朗普"
    )


def test_final_raw_and_wiki_fusion_prefers_matching_title() -> None:
    raw_broad = {
        "id": "raw-trump",
        "title": "特朗普人物资料",
        "type": "llmwiki",
        "content": "正文提到海湖庄园。",
        "raw": {"knowledge_base_id": "kb", "rank_score_semantics": "provider_relevance"},
    }
    wiki_title = {
        "id": "wiki-maralago",
        "title": "海湖庄园",
        "type": "llmwiki_wiki",
        "content": "海湖庄园的资料。",
        "raw": {"knowledge_base_id": "kb"},
    }

    results = _rrf_merge_source_channels(
        [[raw_broad], [wiki_title]],
        limit=2,
        query="海湖庄园",
    )

    assert [item["id"] for item in results] == ["wiki-maralago", "raw-trump"]
    assert results[0]["title_match_score"] > results[1]["title_match_score"]


def test_final_fusion_respects_title_boost_switch_and_weight() -> None:
    broad = {
        "id": "broad",
        "title": "特朗普人物资料",
        "type": "llmwiki",
        "content": "正文提到海湖庄园。",
        "raw": {"knowledge_base_id": "kb"},
    }
    title_hit = {
        "id": "title-hit",
        "title": "海湖庄园",
        "type": "llmwiki_wiki",
        "content": "海湖庄园资料。",
        "raw": {"knowledge_base_id": "kb"},
    }

    disabled = _rrf_merge_source_channels(
        [[broad], [title_hit]],
        limit=2,
        query="海湖庄园",
        title_boost_enabled=False,
        title_boost_weight=1.0,
    )
    zero_weight = _rrf_merge_source_channels(
        [[broad], [title_hit]],
        limit=2,
        query="海湖庄园",
        title_boost_enabled=True,
        title_boost_weight=0.0,
    )
    boosted = _rrf_merge_source_channels(
        [[broad], [title_hit]],
        limit=2,
        query="海湖庄园",
        title_boost_enabled=True,
        title_boost_weight=0.25,
    )

    assert [item["id"] for item in disabled] == ["broad", "title-hit"]
    assert [item["id"] for item in zero_weight] == ["broad", "title-hit"]
    assert [item["id"] for item in boosted] == ["title-hit", "broad"]


def test_final_fusion_keeps_raw_and_formatted_wiki_channels_visible() -> None:
    raw_chunks = [
        {
            "id": f"raw-{index}",
            "title": "海湖庄园资料",
            "type": "llmwiki",
            "content": f"原文切片 {index}",
            "raw": {
                "knowledge_base_id": "kb",
                "retrieval_mode": "raw_chunk",
            },
        }
        for index in range(8)
    ]
    formatted_wiki = [
        {
            "id": f"wiki-{index}",
            "title": f"格式化主题 {index}",
            "type": "llmwiki_wiki",
            "content": f"格式化 Wiki 页面 {index}",
            "raw": {
                "knowledge_base_id": "kb",
                "wiki_slug": f"formatted/{index}",
            },
        }
        for index in range(8)
    ]

    results = _rrf_merge_source_channels(
        [raw_chunks, formatted_wiki],
        limit=8,
        query="海湖庄园",
        title_boost_enabled=True,
        title_boost_weight=0.25,
        min_items_per_channel=4,
    )

    assert [item["type"] for item in results] == [
        "llmwiki",
        "llmwiki_wiki",
        "llmwiki",
        "llmwiki_wiki",
        "llmwiki",
        "llmwiki_wiki",
        "llmwiki",
        "llmwiki_wiki",
    ]


def test_final_fusion_backfills_unused_channel_quota() -> None:
    raw_chunks = [
        {
            "id": f"raw-{index}",
            "title": "海湖庄园资料",
            "type": "llmwiki",
            "raw": {"knowledge_base_id": "kb", "retrieval_mode": "raw_chunk"},
        }
        for index in range(8)
    ]
    formatted_wiki = [
        {
            "id": "wiki-only",
            "title": "格式化主题",
            "type": "llmwiki_wiki",
            "raw": {"knowledge_base_id": "kb", "wiki_slug": "formatted/only"},
        }
    ]

    results = _rrf_merge_source_channels(
        [raw_chunks, formatted_wiki],
        limit=8,
        query="海湖庄园",
        min_items_per_channel=4,
    )

    assert len(results) == 8
    assert sum(item["type"] == "llmwiki_wiki" for item in results) == 1
    assert results[1]["type"] == "llmwiki_wiki"


def test_final_fusion_counts_a_duplicate_only_once_per_channel() -> None:
    duplicate = {
        "id": "raw-duplicate",
        "title": "重复切片",
        "type": "llmwiki",
        "raw": {"knowledge_base_id": "kb", "retrieval_mode": "raw_chunk"},
    }

    results = _rrf_merge_source_channels(
        [[duplicate, dict(duplicate)]],
        limit=8,
    )

    assert len(results) == 1
    assert results[0]["rrf_score"] == 1 / 61