SupplyAgent/agents/generate_demand_agent/tools/dim_constants.py

"""generate_demand_agent 分类树热度维度共享常量与工具函数。"""
from __future__ import annotations

from decimal import Decimal

from supply_infra.db.models.category_tree_weight import CategoryTreeWeight
from supply_infra.db.models.global_tree_category import GlobalTreeCategory

DIM_KEYS: tuple[str, ...] = (
    "ext_pop",
    "plat_sust_pop",
    "plat_ly_pop",
    "recent_pop",
)

DIM_LABEL: dict[str, str] = {
    "ext_pop": "外部热度",
    "plat_sust_pop": "平台持续热度",
    "plat_ly_pop": "平台去年同期热度",
    "recent_pop": "近期热度",
}

# 节点下存在有数据的叶子节点时的标记
HAS_DATA_LEAF_MARK = "+"


def normalize_parent_id(parent_id: int | None) -> int | None:
    if parent_id is None or parent_id == 0:
        return None
    return parent_id


def format_score(avg: Decimal | float | int | None) -> str:
    if avg is None:
        return "—"
    value = float(avg)
    if value >= 100:
        return f"{value:.0f}"
    if value >= 1:
        return f"{value:.1f}"
    return f"{value:.2f}"


def dim_score(
    weight: CategoryTreeWeight | None,
    dim: str,
) -> tuple[float | None, int]:
    """返回 (avg, count)；count>0 即有维度数据（avg 为 0 也算有）。"""
    if weight is None:
        return None, 0
    count = int(getattr(weight, f"{dim}_count", 0) or 0)
    if count <= 0:
        return None, 0
    avg = getattr(weight, f"{dim}_avg", None)
    return (float(avg) if avg is not None else 0.0), count


def build_children_map(
    categories: list[GlobalTreeCategory],
) -> dict[int | None, list[GlobalTreeCategory]]:
    children_map: dict[int | None, list[GlobalTreeCategory]] = {}
    for category in categories:
        parent_key = normalize_parent_id(category.parent_id)
        children_map.setdefault(parent_key, []).append(category)
    for children in children_map.values():
        children.sort(key=lambda c: (c.level or 0, c.id))
    return children_map


def build_score_by_id(
    weights: list[CategoryTreeWeight],
    dim: str,
) -> dict[int, tuple[float | None, int]]:
    return {int(row.category_id): dim_score(row, dim) for row in weights}


def collect_descendant_leaves(
    root_id: int,
    children_map: dict[int | None, list[GlobalTreeCategory]],
) -> list[int]:
    """收集 root 子树内的叶子节点 id（若 root 本身无子节点则含 root）。"""
    leaves: list[int] = []
    stack = [root_id]
    seen: set[int] = set()
    while stack:
        cid = stack.pop()
        if cid in seen:
            continue
        seen.add(cid)
        kids = children_map.get(cid, [])
        if not kids:
            leaves.append(cid)
        else:
            stack.extend(int(c.id) for c in kids)
    leaves.sort()
    return leaves