SupplyAgent/api/services/growth_category_tree.py

"""Build the growth heat map from global_category_v2 and its daily weights."""

from __future__ import annotations

from typing import Any

from sqlalchemy import select

from supply_infra.db.models.global_category_content_weight import (
    GlobalCategoryContentWeight,
)
from supply_infra.db.models.global_v2 import GlobalCategoryV2
from supply_infra.db.repositories.global_category_content_weight_repo import (
    GlobalCategoryContentWeightRepository,
)
from supply_infra.db.session import get_session

GROWTH_DIM_META: list[dict[str, str]] = [
    {"key": "read_rate", "label": "阅读率"},
    {"key": "avg_read_rate", "label": "平均阅读率"},
    {"key": "like_rate", "label": "点赞率"},
]


def _normalize_parent_id(parent_id: int | None, category_ids: set[int]) -> int | None:
    if parent_id in (None, 0) or int(parent_id) not in category_ids:
        return None
    return int(parent_id)


def _build_growth_tree(
    categories: list[GlobalCategoryV2],
    weights: list[GlobalCategoryContentWeight],
) -> list[dict[str, Any]]:
    category_ids = {int(category.stable_id) for category in categories}
    children_by_parent: dict[int | None, list[GlobalCategoryV2]] = {}
    for category in categories:
        parent_id = _normalize_parent_id(category.parent_stable_id, category_ids)
        children_by_parent.setdefault(parent_id, []).append(category)
    for children in children_by_parent.values():
        children.sort(key=lambda row: (row.level or 0, int(row.stable_id)))

    weight_by_stable_id = {int(row.stable_id): row for row in weights}

    def to_node(category: GlobalCategoryV2, ancestors: frozenset[int]) -> dict[str, Any]:
        stable_id = int(category.stable_id)
        weight = weight_by_stable_id.get(stable_id)
        count = int(weight.source_element_count or 0) if weight else 0
        next_ancestors = ancestors | {stable_id}
        children = [
            to_node(child, next_ancestors)
            for child in children_by_parent.get(stable_id, [])
            if int(child.stable_id) not in next_ancestors
        ]
        return {
            "id": stable_id,
            "name": category.name,
            "level": category.level,
            "description": category.description,
            "weights": {
                "read_rate": float(weight.read_rate_score) if weight else None,
                "avg_read_rate": float(weight.avg_read_rate_score) if weight else None,
                "like_rate": float(weight.like_rate_score) if weight else None,
            },
            "counts": {
                "read_rate": count,
                "avg_read_rate": count,
                "like_rate": count,
            },
            "children": children,
        }

    return [to_node(root, frozenset()) for root in children_by_parent.get(None, [])]


def build_growth_category_tree(biz_dt: str | None = None) -> dict[str, Any]:
    """Return the V2 category tree with the requested/latest complete daily scores."""
    with get_session() as session:
        categories = list(
            session.scalars(
                select(GlobalCategoryV2).order_by(
                    GlobalCategoryV2.level, GlobalCategoryV2.stable_id
                )
            ).all()
        )
        weight_repo = GlobalCategoryContentWeightRepository(session)
        resolved_dt = biz_dt or weight_repo.get_latest_completed_biz_dt()
        weights = weight_repo.list_completed_models_by_biz_dt(resolved_dt) if resolved_dt else []
        nodes = _build_growth_tree(categories, weights)

    return {
        "biz_dt": resolved_dt,
        "dims": GROWTH_DIM_META,
        "nodes": nodes,
    }