tests/test_rerank_client.py

from math import isclose
from config.schema import CoarseRankFusionConfig, RerankFusionConfig
from search.rerank_client import coarse_resort_hits, fuse_scores_and_resort, run_lightweight_rerank
def test_fuse_scores_and_resort_aggregates_text_components_and_keeps_rerank_primary():
    hits = [
        {
            "_id": "1",
            "_score": 3.2,
            "matched_queries": {
                "base_query": 2.4,
                "base_query_trans_zh": 1.8,
                "knn_query": 0.8,
            },
        },
        {
            "_id": "2",
            "_score": 2.8,
            "matched_queries": {
                "base_query": 9.0,
                "knn_query": 0.2,
            },
        },
    ]
    debug = fuse_scores_and_resort(hits, [0.9, 0.7], debug=True)
    expected_text_1 = 2.4 + 0.25 * (0.8 * 1.8)
    expected_fused_1 = (0.9 + 0.00001) * ((expected_text_1 + 0.1) ** 0.35) * ((0.8 + 0.6) ** 0.2)
    expected_fused_2 = (0.7 + 0.00001) * ((9.0 + 0.1) ** 0.35) * ((0.2 + 0.6) ** 0.2)
    by_id = {hit["_id"]: hit for hit in hits}
    assert isclose(by_id["1"]["_text_score"], expected_text_1, rel_tol=1e-9)
    assert isclose(by_id["1"]["_fused_score"], expected_fused_1, rel_tol=1e-9)
    assert isclose(by_id["2"]["_fused_score"], expected_fused_2, rel_tol=1e-9)
    assert debug[0]["text_source_score"] == 2.4
    assert debug[0]["text_translation_score"] == 1.8
    assert isclose(debug[0]["text_weighted_translation_score"], 1.44, rel_tol=1e-9)
    assert debug[0]["knn_score"] == 0.8
    assert isclose(debug[0]["rerank_factor"], 0.90001, rel_tol=1e-9)
    assert [hit["_id"] for hit in hits] == ["2", "1"]
def test_fuse_scores_and_resort_falls_back_when_matched_queries_missing():
    hits = [
        {"_id": "1", "_score": 0.5},
        {"_id": "2", "_score": 2.0},
    ]
    debug = fuse_scores_and_resort(hits, [0.4, 0.3], debug=True)
    expected_1 = (0.4 + 0.00001) * ((0.5 + 0.1) ** 0.35) * ((0.0 + 0.6) ** 0.2)
    expected_2 = (0.3 + 0.00001) * ((2.0 + 0.1) ** 0.35) * ((0.0 + 0.6) ** 0.2)
    by_id = {hit["_id"]: hit for hit in hits}
    assert isclose(by_id["1"]["_text_score"], 0.5, rel_tol=1e-9)
    assert isclose(by_id["1"]["_fused_score"], expected_1, rel_tol=1e-9)
    assert isclose(by_id["2"]["_text_score"], 2.0, rel_tol=1e-9)
    assert isclose(by_id["2"]["_fused_score"], expected_2, rel_tol=1e-9)
    assert debug[0]["text_score_fallback_to_es"] is True
    assert debug[1]["text_score_fallback_to_es"] is True
    assert [hit["_id"] for hit in hits] == ["2", "1"]
def test_fuse_scores_and_resort_downweights_text_only_advantage():
    hits = [
        {
            "_id": "lexical-heavy",
            "_score": 10.0,
            "matched_queries": {
                "base_query": 10.0,
                "knn_query": 0.0,
            },
        },
        {
            "_id": "rerank-better",
            "_score": 6.0,
            "matched_queries": {
                "base_query": 6.0,
                "knn_query": 0.0,
            },
        },
    ]
    fuse_scores_and_resort(hits, [0.72, 0.98])
    assert [hit["_id"] for hit in hits] == ["rerank-better", "lexical-heavy"]
def test_fuse_scores_and_resort_uses_configurable_fusion_params():
    hits = [
        {
            "_id": "a",
            "_score": 1.0,
            "matched_queries": {"base_query": 2.0, "knn_query": 0.5},
        },
        {
            "_id": "b",
            "_score": 1.0,
            "matched_queries": {"base_query": 3.0, "knn_query": 0.0},
        },
    ]
    fusion = RerankFusionConfig(
        rerank_bias=0.0,
        rerank_exponent=1.0,
        text_bias=0.0,
        text_exponent=1.0,
        knn_bias=0.0,
        knn_exponent=1.0,
    )
    fuse_scores_and_resort(hits, [1.0, 1.0], fusion=fusion)
    # b 的 knn 为 0 -> 融合为 0；a 为 1 * 2 * 0.5
    assert [h["_id"] for h in hits] == ["a", "b"]
    by_id = {h["_id"]: h for h in hits}
    assert isclose(by_id["a"]["_fused_score"], 1.0, rel_tol=1e-9)
    assert isclose(by_id["b"]["_fused_score"], 0.0, rel_tol=1e-9)
def test_fuse_scores_and_resort_boosts_hits_with_selected_sku():
    hits = [
        {
            "_id": "style-selected",
            "_score": 1.0,
            "_style_rerank_suffix": "Blue XL",
            "matched_queries": {"base_query": 1.0, "knn_query": 0.0},
        },
        {
            "_id": "plain",
            "_score": 1.0,
            "matched_queries": {"base_query": 1.0, "knn_query": 0.0},
        },
    ]
    debug = fuse_scores_and_resort(
        hits,
        [1.0, 1.0],
        style_intent_selected_sku_boost=1.2,
        debug=True,
    )
    by_id = {h["_id"]: h for h in hits}
    assert isclose(by_id["style-selected"]["_fused_score"], by_id["plain"]["_fused_score"] * 1.2, rel_tol=1e-9)
    assert by_id["style-selected"]["_style_intent_selected_sku_boost"] == 1.2
    assert by_id["plain"]["_style_intent_selected_sku_boost"] == 1.0
    assert [h["_id"] for h in hits] == ["style-selected", "plain"]
    assert debug[0]["style_intent_selected_sku"] is True
    assert debug[0]["style_intent_selected_sku_boost"] == 1.2
def test_fuse_scores_and_resort_uses_max_of_text_and_image_knn_scores():
    hits = [
        {
            "_id": "mm-hit",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 1.5,
                "knn_query": 0.2,
                "image_knn_query": 0.7,
            },
        }
    ]
    debug = fuse_scores_and_resort(hits, [0.8], debug=True)
    assert isclose(hits[0]["_knn_score"], 0.7, rel_tol=1e-9)
    assert isclose(debug[0]["knn_score"], 0.7, rel_tol=1e-9)
    assert isclose(debug[0]["text_knn_score"], 0.2, rel_tol=1e-9)
    assert isclose(debug[0]["image_knn_score"], 0.7, rel_tol=1e-9)
def test_fuse_scores_and_resort_prefers_exact_knn_scores_over_ann_scores():
    hits = [
        {
            "_id": "exact-mm-hit",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 1.5,
                "knn_query": 0.2,
                "image_knn_query": 0.7,
                "exact_text_knn_query": 0.9,
                "exact_image_knn_query": 0.1,
            },
        }
    ]
    debug = fuse_scores_and_resort(hits, [0.8], debug=True)
    assert isclose(hits[0]["_knn_score"], 0.9, rel_tol=1e-9)
    assert isclose(debug[0]["text_knn_score"], 0.9, rel_tol=1e-9)
    assert isclose(debug[0]["image_knn_score"], 0.1, rel_tol=1e-9)
    assert isclose(debug[0]["exact_text_knn_score"], 0.9, rel_tol=1e-9)
    assert isclose(debug[0]["exact_image_knn_score"], 0.1, rel_tol=1e-9)
    assert isclose(debug[0]["approx_text_knn_score"], 0.2, rel_tol=1e-9)
    assert isclose(debug[0]["approx_image_knn_score"], 0.7, rel_tol=1e-9)
    assert debug[0]["text_knn_source"] == "exact_text_knn_query"
    assert debug[0]["image_knn_source"] == "exact_image_knn_query"
def test_fuse_scores_and_resort_falls_back_to_ann_when_exact_knn_missing():
    hits = [
        {
            "_id": "ann-only-hit",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 1.5,
                "knn_query": 0.4,
                "image_knn_query": 0.5,
            },
        }
    ]
    debug = fuse_scores_and_resort(hits, [0.8], debug=True)
    assert isclose(debug[0]["text_knn_score"], 0.4, rel_tol=1e-9)
    assert isclose(debug[0]["image_knn_score"], 0.5, rel_tol=1e-9)
    assert isclose(debug[0]["approx_text_knn_score"], 0.4, rel_tol=1e-9)
    assert isclose(debug[0]["approx_image_knn_score"], 0.5, rel_tol=1e-9)
    assert debug[0]["text_knn_source"] == "knn_query"
    assert debug[0]["image_knn_source"] == "image_knn_query"
def test_fuse_scores_and_resort_applies_knn_dismax_weights_and_tie_breaker():
    hits = [
        {
            "_id": "mm-hit",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 1.5,
                "knn_query": 0.4,
                "image_knn_query": 0.5,
            },
        }
    ]
    fusion = RerankFusionConfig(
        rerank_bias=0.00001,
        rerank_exponent=1.0,
        text_bias=0.1,
        text_exponent=0.35,
        knn_text_weight=2.0,
        knn_image_weight=1.0,
        knn_tie_breaker=0.25,
        knn_bias=0.0,
        knn_exponent=1.0,
    )
    debug = fuse_scores_and_resort(hits, [0.8], fusion=fusion, debug=True)
    expected_knn = 0.8 + 0.25 * 0.5
    assert isclose(hits[0]["_knn_score"], expected_knn, rel_tol=1e-9)
    assert isclose(debug[0]["weighted_text_knn_score"], 0.8, rel_tol=1e-9)
    assert isclose(debug[0]["weighted_image_knn_score"], 0.5, rel_tol=1e-9)
    assert isclose(debug[0]["knn_primary_score"], 0.8, rel_tol=1e-9)
    assert isclose(debug[0]["knn_support_score"], 0.5, rel_tol=1e-9)
def test_fuse_scores_and_resort_can_add_weighted_text_and_image_knn_factors():
    hits = [
        {
            "_id": "a",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 2.0,
                "knn_query": 0.4,
                "image_knn_query": 0.5,
            },
        }
    ]
    fusion = RerankFusionConfig(
        rerank_bias=0.0,
        rerank_exponent=1.0,
        text_bias=0.0,
        text_exponent=1.0,
        knn_text_weight=2.0,
        knn_image_weight=1.0,
        knn_tie_breaker=0.25,
        knn_bias=0.1,
        knn_exponent=1.0,
        knn_text_exponent=2.0,
        knn_image_exponent=3.0,
    )
    debug = fuse_scores_and_resort(hits, [0.8], fusion=fusion, debug=True)
    weighted_text_knn = 0.8
    weighted_image_knn = 0.5
    expected_knn = weighted_text_knn + 0.25 * weighted_image_knn
    expected_fused = (
        0.8
        * 2.0
        * (expected_knn + 0.1)
        * ((weighted_text_knn + 0.1) ** 2.0)
        * ((weighted_image_knn + 0.1) ** 3.0)
    )
    assert isclose(hits[0]["_fused_score"], expected_fused, rel_tol=1e-9)
    assert isclose(debug[0]["text_knn_factor"], (weighted_text_knn + 0.1) ** 2.0, rel_tol=1e-9)
    assert isclose(debug[0]["image_knn_factor"], (weighted_image_knn + 0.1) ** 3.0, rel_tol=1e-9)
    assert "weighted_text_knn_score=" in debug[0]["fusion_summary"]
    assert "weighted_image_knn_score=" in debug[0]["fusion_summary"]
def test_coarse_resort_hits_can_add_weighted_text_and_image_knn_factors():
    hits = [
        {
            "_id": "coarse-a",
            "_score": 1.0,
            "matched_queries": {
                "base_query": 2.0,
                "knn_query": 0.4,
                "image_knn_query": 0.5,
            },
        }
    ]
    fusion = CoarseRankFusionConfig(
        es_bias=0.0,
        es_exponent=1.0,
        text_bias=0.0,
        text_exponent=1.0,
        knn_text_weight=2.0,
        knn_image_weight=1.0,
        knn_tie_breaker=0.25,
        knn_bias=0.1,
        knn_exponent=1.0,
        knn_text_exponent=2.0,
        knn_image_exponent=3.0,
    )
    debug = coarse_resort_hits(hits, fusion=fusion, debug=True)
    weighted_text_knn = 0.8
    weighted_image_knn = 0.5
    expected_knn = weighted_text_knn + 0.25 * weighted_image_knn
    expected_coarse = (
        1.0
        * 2.0
        * (expected_knn + 0.1)
        * ((weighted_text_knn + 0.1) ** 2.0)
        * ((weighted_image_knn + 0.1) ** 3.0)
    )
    assert isclose(hits[0]["_coarse_score"], expected_coarse, rel_tol=1e-9)
    assert isclose(debug[0]["coarse_text_knn_factor"], (weighted_text_knn + 0.1) ** 2.0, rel_tol=1e-9)
    assert isclose(debug[0]["coarse_image_knn_factor"], (weighted_image_knn + 0.1) ** 3.0, rel_tol=1e-9)
def test_run_lightweight_rerank_sorts_by_fused_stage_score(monkeypatch):
    hits = [
        {
            "_id": "fine-raw-better",
            "_score": 1.0,
            "_source": {"title": {"en": "Alpha"}},
            "matched_queries": {"base_query": 0.5, "knn_query": 0.0},
        },
        {
            "_id": "fusion-better",
            "_score": 1.0,
            "_source": {"title": {"en": "Beta"}},
            "matched_queries": {"base_query": 40.0, "knn_query": 0.0},
        },
    ]
    monkeypatch.setattr(
        "search.rerank_client.call_rerank_service",
        lambda *args, **kwargs: ([0.9, 0.8], {"model": "fine-bge"}),
    )
    scores, meta, debug_rows = run_lightweight_rerank(
        query="toy",
        es_hits=hits,
        language="en",
        debug=True,
    )
    assert scores == [0.9, 0.8]
    assert meta == {"model": "fine-bge"}
    assert [hit["_id"] for hit in hits] == ["fusion-better", "fine-raw-better"]
    assert hits[0]["_fine_fused_score"] > hits[1]["_fine_fused_score"]
    assert debug_rows[0]["fusion_summary"]
    assert "fine_score=" in debug_rows[0]["fusion_summary"]
    assert "text_score=" in debug_rows[0]["fusion_summary"]
def test_fuse_scores_and_resort_uses_hit_level_fine_score_when_not_passed_separately():
    hits = [
        {
            "_id": "with-fine",
            "_score": 1.0,
            "_fine_score": 0.7,
            "matched_queries": {"base_query": 2.0, "knn_query": 0.5},
        }
    ]
    debug = fuse_scores_and_resort(hits, [0.8], debug=True)
    assert isclose(debug[0]["fine_factor"], (0.7 + 0.00001), rel_tol=1e-9)
    assert debug[0]["fusion_inputs"]["fine_score"] == 0.7
    assert "fine_score=" in debug[0]["fusion_summary"]
def test_fuse_scores_and_resort_can_include_raw_es_score_as_factor():
    hits = [
        {
            "_id": "es-strong",
            "_score": 100.0,
            "matched_queries": {"base_query": 1.0, "knn_query": 0.0},
        },
        {
            "_id": "es-weak",
            "_score": 1.0,
            "matched_queries": {"base_query": 1.0, "knn_query": 0.0},
        },
    ]
    fusion = RerankFusionConfig(
        es_bias=0.0,
        es_exponent=1.0,
        rerank_bias=0.0,
        rerank_exponent=1.0,
        text_bias=0.0,
        text_exponent=0.0,
        knn_bias=1.0,
        knn_exponent=0.0,
    )
    debug = fuse_scores_and_resort(hits, [1.0, 1.0], fusion=fusion, debug=True)
    assert [hit["_id"] for hit in hits] == ["es-strong", "es-weak"]
    assert isclose(hits[0]["_raw_es_score"], 100.0, rel_tol=1e-9)
    assert isclose(debug[0]["es_factor"], 100.0, rel_tol=1e-9)
    assert debug[0]["fusion_inputs"]["es_score"] == 100.0