mindfulness/server/app/features/personalized_reco/rerank_freqcap/rerank.py

from __future__ import annotations

from typing import Any, Iterable, Optional

from app.features.personalized_reco.rerank_freqcap.defaults import get_default_config
from app.features.personalized_reco.rerank_freqcap.types import RerankConfig, RerankMeta, RerankResult, ScoredCandidate, Scene
from app.features.personalized_reco.rerank_freqcap.utils import as_finite_float, build_tags, clamp, jaccard, normalize_int_id_set


def _sort_by_score_desc(cands: list[ScoredCandidate]) -> list[ScoredCandidate]:
    return sorted(cands, key=lambda x: as_finite_float(x.final_score, default=float("-inf")), reverse=True)


def _dedup_by_seen_ids(
    cands: list[ScoredCandidate],
    *,
    seen_ids: set[int],
) -> tuple[list[ScoredCandidate], int]:
    kept: list[ScoredCandidate] = []
    removed = 0
    for c in cands:
        if int(c.content_id) in seen_ids:
            removed += 1
            continue
        kept.append(c)
    return kept, removed


def _apply_author_template_freqcap(
    cands: list[ScoredCandidate],
    *,
    recent_author_ids: Optional[Iterable[str]],
    recent_template_ids: Optional[Iterable[str]],
) -> tuple[list[ScoredCandidate], dict[str, int], list[str]]:
    """
    V1 策略：
    - 若 recent_*_ids 未提供（None），不执行该维度过滤，但在 meta 记录缺失
    - 若提供，则执行硬过滤
    """

    filtered_counts: dict[str, int] = {"author": 0, "template": 0}
    missing: list[str] = []

    author_set: set[str] | None
    if recent_author_ids is None:
        author_set = None
        missing.append("author")
    else:
        author_set = set([a for a in recent_author_ids if a is not None and str(a).strip() != ""])

    template_set: set[str] | None
    if recent_template_ids is None:
        template_set = None
        missing.append("template")
    else:
        template_set = set([t for t in recent_template_ids if t is not None and str(t).strip() != ""])

    out: list[ScoredCandidate] = []
    for c in cands:
        if author_set is not None and c.author_id and c.author_id in author_set:
            filtered_counts["author"] += 1
            continue
        if template_set is not None and c.template_id and c.template_id in template_set:
            filtered_counts["template"] += 1
            continue
        out.append(c)

    # 只返回真正生效的维度计数（避免 meta 噪音）
    effective_counts: dict[str, int] = {}
    if author_set is not None:
        effective_counts["author"] = int(filtered_counts["author"])
    if template_set is not None:
        effective_counts["template"] = int(filtered_counts["template"])

    missing_sorted = sorted(set(missing))
    return out, effective_counts, missing_sorted


def _sim(a: ScoredCandidate, b: ScoredCandidate, *, tags_a: set[str], tags_b: set[str]) -> float:
    # 离散特征版（V1 推荐），对齐 plan.md
    if int(a.content_id) == int(b.content_id):
        return 1.0

    sim = 0.0
    if a.template_id and b.template_id and a.template_id == b.template_id:
        sim += 0.6
    if a.author_id and b.author_id and a.author_id == b.author_id:
        sim += 0.3

    sim += 0.1 * jaccard(tags_a, tags_b)
    return clamp(sim, 0.0, 1.0)


def _mmr_rerank(
    *,
    candidates: list[ScoredCandidate],
    k: int,
    lam: float,
) -> list[ScoredCandidate]:
    if k <= 0:
        return []

    if not candidates:
        return []

    lam_f = clamp(as_finite_float(lam, default=0.7), 0.0, 1.0)

    # 预计算 tags，避免重复构造
    tags_map: dict[int, set[str]] = {}
    for c in candidates:
        tags_map[int(c.content_id)] = build_tags(c)

    remaining = _sort_by_score_desc(list(candidates))
    selected: list[ScoredCandidate] = []

    # Top1：最高分
    selected.append(remaining.pop(0))

    while remaining and len(selected) < k:
        best_idx = 0
        best_val = float("-inf")

        for idx, c in enumerate(remaining):
            rel = as_finite_float(c.final_score, default=float("-inf"))

            tags_c = tags_map.get(int(c.content_id), set())
            max_sim = 0.0
            for s in selected:
                tags_s = tags_map.get(int(s.content_id), set())
                max_sim = max(max_sim, _sim(c, s, tags_a=tags_c, tags_b=tags_s))

            val = lam_f * float(rel) - (1.0 - lam_f) * float(max_sim)
            if val > best_val:
                best_val = val
                best_idx = idx

        selected.append(remaining.pop(best_idx))

    return selected


def rerank_and_freqcap(
    *,
    scene: Scene,
    scored_candidates: list[ScoredCandidate],
    already_recommended_ids: list[Any],
    touched_or_viewed_ids: list[Any],
    k: int,
    config: Optional[RerankConfig] = None,
    recent_author_ids: Optional[list[str]] = None,
    recent_template_ids: Optional[list[str]] = None,
) -> RerankResult:
    """
    主入口：对 scored_candidates 做去重/频控/重排，输出最终可下发序列。

    V1 约定：
    - 冷却窗口“按天”由调用方保证输入集合已经裁剪到窗口内，本模块以“集合代表窗口内历史”为准
    - Feed 默认只做 dedup + MMR；Push/Widget 做 dedup + freqcap + TopK
    """

    cfg = config or get_default_config(scene)

    # seen_ids = already_recommended_ids ∪ touched_or_viewed_ids
    seen_ids = normalize_int_id_set(list(already_recommended_ids) + list(touched_or_viewed_ids))

    # 先按分数降序，保证 Top1 与 TopK 一致
    base_sorted = _sort_by_score_desc(list(scored_candidates))

    after_dedup, removed_sentence = _dedup_by_seen_ids(base_sorted, seen_ids=seen_ids)
    candidate_pool_size_after_dedup = len(after_dedup)

    missing_history_fields: list[str] = []
    freqcap_counts: dict[str, int] = {"sentence": int(removed_sentence)}

    after_freqcap = after_dedup

    # Push/Widget：作者/模板冷却（增强项）
    if scene in {"push", "widget"}:
        after_freqcap, dim_counts, missing = _apply_author_template_freqcap(
            after_freqcap,
            recent_author_ids=recent_author_ids,
            recent_template_ids=recent_template_ids,
        )
        missing_history_fields = missing
        freqcap_counts.update(dim_counts)
    else:
        # Feed：不强制作者/模板冷却（V1 可选，这里默认跳过）
        missing_history_fields = []

    candidate_pool_size_after_freqcap = len(after_freqcap)

    ranked: list[ScoredCandidate]
    if scene == "feed":
        # MMR 前截断，避免性能问题
        top_n = int(cfg.top_n_for_mmr) if int(cfg.top_n_for_mmr) > 0 else len(after_freqcap)
        mmr_pool = after_freqcap[:top_n]
        ranked = _mmr_rerank(candidates=mmr_pool, k=int(k), lam=cfg.mmr_lambda)
    else:
        ranked = after_freqcap[: max(0, int(k))]

    meta = RerankMeta(
        candidate_pool_size_after_dedup=int(candidate_pool_size_after_dedup),
        candidate_pool_size_after_freqcap=int(candidate_pool_size_after_freqcap),
        missing_history_fields=missing_history_fields,
        freqcap_filtered_counts=freqcap_counts,
    )
    return RerankResult(ranked_items=ranked, meta=meta)