SupplyAgent/agents/demand_grade_orchestrator_agent/common/tree_state.py

"""全局树热度状态的加载与格式化。"""
from __future__ import annotations

from collections import defaultdict
from types import SimpleNamespace
from typing import Any

from supply_infra.scoring.ranking import rank_with_scores
from supply_infra.db.repositories.category_tree_weight_repo import CategoryTreeWeightRepository
from supply_infra.db.repositories.global_tree_category_repo import GlobalTreeCategoryRepository
from supply_infra.db.session import get_session


def load_tree_state(biz_dt: str) -> tuple[dict[int, Any], dict[int | None, list[int]], dict[int, Any]]:
    with get_session() as session:
        categories = [
            SimpleNamespace(id=int(row.id), name=row.name, parent_id=row.parent_id, level=row.level)
            for row in GlobalTreeCategoryRepository(session).list_active_categories()
        ]
        weights = [
            SimpleNamespace(
                category_id=int(row.category_id),
                total_score=row.total_score,
                hung_word_count=row.hung_word_count,
                real_rov_7d_avg=row.real_rov_7d_avg,
                real_rov_7d_count=row.real_rov_7d_count,
                real_vov_7d_avg=row.real_vov_7d_avg,
                real_vov_7d_count=row.real_vov_7d_count,
            )
            for row in CategoryTreeWeightRepository(session).list_by_biz_dt(biz_dt)
        ]
    by_id = {row.id: row for row in categories}
    children: dict[int | None, list[int]] = defaultdict(list)
    for row in categories:
        parent_id = int(row.parent_id) if row.parent_id not in (None, 0) else None
        children[parent_id].append(row.id)
    for ids in children.values():
        ids.sort()
    return by_id, children, {row.category_id: row for row in weights}


def format_heat_score(weight: Any | None) -> str:
    """格式化原始 total_score：只保留两位小数，无数据时明确为 null。"""
    if weight is None or weight.total_score is None:
        return "null"
    return f"{float(weight.total_score):.2f}"


def has_hung_demand(weight: Any | None) -> bool:
    """判断分类自身是否挂有需求。"""
    return weight is not None and int(weight.hung_word_count or 0) > 0


def format_demand_count(weight: Any | None) -> str:
    """格式化节点挂载需求数量；无需求时返回空字符串。"""
    if not has_hung_demand(weight):
        return ""
    return str(int(weight.hung_word_count or 0))


HEAT_LEVEL_DEFINITION: dict[str, dict[str, str | float | None]] = {
    "S": {"label": "高热", "min_global_rank_score": 0.90},
    "A": {"label": "较高热", "min_global_rank_score": 0.75},
    "B": {"label": "中热", "min_global_rank_score": 0.50},
    "C": {"label": "较低热", "min_global_rank_score": 0.25},
    "D": {"label": "低热", "min_global_rank_score": 0.00},
    "U": {"label": "数据不足", "min_global_rank_score": None},
}


def global_heat_positions(weights: dict[int, Any]) -> dict[int, dict[str, float | int]]:
    """返回节点 ``total_score`` 在整棵有分节点中的排名位置。"""
    return rank_with_scores([
        (category_id, float(weight.total_score))
        for category_id, weight in weights.items()
        if weight is not None and weight.total_score is not None
    ])


def heat_level(position: dict[str, float | int] | None) -> str:
    """按整树排名分映射批次热度等级；无分节点单列 U。"""
    if position is None:
        return "U"
    score = float(position["normalized_score"])
    for level in ("S", "A", "B", "C"):
        threshold = HEAT_LEVEL_DEFINITION[level]["min_global_rank_score"]
        if threshold is not None and score >= float(threshold):
            return level
    return "D"


def format_rank(position: dict[str, float | int] | None) -> str:
    if position is None:
        return "无排名"
    rank = float(position["rank"])
    rank_text = str(int(rank)) if rank.is_integer() else f"{rank:.1f}"
    return f"{rank_text}/{int(position['total'])}"


def path(category_id: int | None, by_id: dict[int, Any]) -> str:
    names: list[str] = []
    current = by_id.get(category_id) if category_id is not None else None
    while current is not None:
        names.append(current.name or str(current.id))
        parent_id = int(current.parent_id) if current.parent_id not in (None, 0) else None
        current = by_id.get(parent_id) if parent_id is not None else None
    return " > ".join(reversed(names))