SupplyAgent/agents/demand_grade_agent/tools/tree_local.py

"""demand_grade_agent 专用的全局树局部热度查询辅助（不对外注册为工具）。"""
from __future__ import annotations

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

from agents.demand_grade_agent.tools.shared import (
    build_category_heat_summary,
    build_category_path,
    format_category_weight_lines,
)
from supply_agent.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 _normalize_parent_id(parent_id: int | None) -> int | None:
    if parent_id is None or parent_id == 0:
        return None
    return int(parent_id)


def _materialize_category(row: Any) -> SimpleNamespace:
    """在 session 内抽出标量，避免 DetachedInstanceError。"""
    return SimpleNamespace(
        id=int(row.id),
        name=row.name,
        parent_id=row.parent_id,
        level=row.level,
    )


def _materialize_weight(row: Any) -> SimpleNamespace:
    """在 session 内抽出标量，避免 DetachedInstanceError。"""
    return SimpleNamespace(
        category_id=int(row.category_id),
        biz_dt=row.biz_dt,
        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,
    )


def _score_value(weight: Any | None) -> float | None:
    if weight is None or weight.total_score is None:
        return None
    return float(weight.total_score)


def load_local_tree_state(biz_dt: str) -> tuple[dict[int, Any], dict[int | None, list[int]], dict[int, Any]]:
    with get_session() as session:
        categories = [
            _materialize_category(row)
            for row in GlobalTreeCategoryRepository(session).list_active_categories()
        ]
        weights = [
            _materialize_weight(row)
            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:
        children[_normalize_parent_id(row.parent_id)].append(row.id)
    for ids in children.values():
        ids.sort()
    weight_by_id = {row.category_id: row for row in weights}
    return by_id, children, weight_by_id


def _describe_node(
    category_id: int,
    by_id: dict[int, Any],
    weight_by_id: dict[int, Any],
    *,
    biz_dt: str,
    global_positions: dict[int, dict[str, float | int]],
) -> dict[str, Any]:
    category = by_id.get(category_id)
    weight = weight_by_id.get(category_id)
    position = global_positions.get(category_id)
    return {
        "category_id": category_id,
        "name": category.name if category is not None else None,
        "path": build_category_path(category_id, by_id),
        "hung_word_count": int(weight.hung_word_count or 0) if weight is not None else 0,
        "heat": build_category_heat_summary(weight, biz_dt=biz_dt),
        "global_tree_position": {
            "rank": float(position["rank"]) if position is not None else None,
            "scored_node_count": int(position["total"]) if position is not None else 0,
            "normalized_rank_score": (
                float(position["normalized_score"]) if position is not None else None
            ),
        },
    }


def _sibling_ids(category_id: int, parent_id: int | None, children: dict[int | None, list[int]]) -> list[int]:
    if parent_id is not None:
        return list(children.get(parent_id, []))
    return list(children.get(None, []))


def build_local_heat_snapshot(
    biz_dt: str,
    category_ids: list[int],
    *,
    tree_state: tuple[dict[int, Any], dict[int | None, list[int]], dict[int, Any]] | None = None,
) -> list[dict[str, Any]]:
    """构建局部环境快照：自身、父节点、全部兄弟节点的全局热度与后验。"""
    if tree_state is None:
        by_id, children, weight_by_id = load_local_tree_state(biz_dt)
    else:
        by_id, children, weight_by_id = tree_state
    snapshots: list[dict[str, Any]] = []
    global_positions = rank_with_scores([
        (category_id, float(weight.total_score))
        for category_id, weight in weight_by_id.items()
        if weight is not None and weight.total_score is not None
    ])
    for category_id in dict.fromkeys(int(value) for value in category_ids):
        category = by_id.get(category_id)
        if category is None:
            snapshots.append({
                "category_id": category_id,
                "error": "分类不存在",
            })
            continue

        parent_id = _normalize_parent_id(category.parent_id)
        sibling_ids = _sibling_ids(category_id, parent_id, children)
        ranked_siblings = sorted(
            sibling_ids,
            key=lambda cid: (
                _score_value(weight_by_id.get(cid)) is None,
                -(_score_value(weight_by_id.get(cid)) or 0),
                cid,
            ),
        )
        sibling_total = len(ranked_siblings)
        siblings = []
        for index, sibling_id in enumerate(ranked_siblings, start=1):
            item = _describe_node(
                sibling_id,
                by_id,
                weight_by_id,
                biz_dt=biz_dt,
                global_positions=global_positions,
            )
            item["rank_among_siblings"] = index
            item["sibling_count"] = sibling_total
            item["is_self"] = sibling_id == category_id
            siblings.append(item)

        snapshots.append({
            "category_id": category_id,
            "biz_dt": biz_dt,
            "global_scored_node_count": len(global_positions),
            "self": _describe_node(
                category_id,
                by_id,
                weight_by_id,
                biz_dt=biz_dt,
                global_positions=global_positions,
            ),
            "parent": (
                _describe_node(
                    parent_id,
                    by_id,
                    weight_by_id,
                    biz_dt=biz_dt,
                    global_positions=global_positions,
                )
                if parent_id is not None
                else None
            ),
            "siblings": siblings,
            "sibling_count": sibling_total,
        })
    return snapshots


def _append_node_block(
    lines: list[str],
    *,
    title: str,
    node: dict[str, Any] | None,
    biz_dt: str,
    weight_by_id: dict[int, Any] | None = None,
) -> None:
    if node is None:
        lines.append(f"{title}: 无")
        return
    lines.append(f"{title}:")
    weight = weight_by_id.get(int(node["category_id"])) if weight_by_id is not None else None
    lines.extend(format_category_weight_lines(
        weight,
        category_id=int(node["category_id"]),
        name=node.get("name"),
        path=node.get("path"),
        biz_dt=biz_dt,
        hung_word_count=int(node.get("hung_word_count") or 0),
    ))
    position = node.get("global_tree_position") or {}
    rank = position.get("rank")
    if rank is None:
        lines.append("  整树位置=无排名（不是低热，表示缺少 total_score）")
    else:
        lines.append(
            f"  整树位置={float(rank):g}/{int(position.get('scored_node_count') or 0)} "
            f"整树排名归一分={float(position['normalized_rank_score']):.4f}"
        )


def format_local_heat_snapshot(snapshot: dict[str, Any], *, weight_by_id: dict[int, Any] | None = None) -> list[str]:
    if snapshot.get("error"):
        return [f"--- category_id={snapshot['category_id']} ---", f"错误: {snapshot['error']}"]

    biz_dt = snapshot["biz_dt"]
    lines = [f"--- category_id={snapshot['category_id']} ---"]
    _append_node_block(lines, title="自身节点", node=snapshot["self"], biz_dt=biz_dt, weight_by_id=weight_by_id)
    _append_node_block(lines, title="父节点", node=snapshot.get("parent"), biz_dt=biz_dt, weight_by_id=weight_by_id)

    siblings = snapshot.get("siblings") or []
    lines.append(f"全部兄弟节点（共 {len(siblings)} 个，按 total_score 降序，* 为当前节点）:")
    for sibling in siblings:
        marker = "*" if sibling.get("is_self") else " "
        rank_note = f" 兄弟排名 {sibling.get('rank_among_siblings')}/{sibling.get('sibling_count')}"
        lines.append(f"{marker} {'-' * 8}{rank_note}")
        _append_node_block(lines, title="  节点", node=sibling, biz_dt=biz_dt, weight_by_id=weight_by_id)
    return lines


def render_local_heat_report(biz_dt: str, category_ids: list[int]) -> str:
    tree_state = load_local_tree_state(biz_dt)
    _, _, weight_by_id = tree_state
    snapshots = build_local_heat_snapshot(biz_dt, category_ids, tree_state=tree_state)
    if not snapshots:
        return "category_ids 不能为空"
    lines = [
        f"biz_dt={biz_dt} | 局部环境包含节点自身、父节点、全部兄弟节点的全局热度 total_score 与后验数据；",
        "每个节点同时给出整棵树名次；不得只参考附近节点，也不得只参考部分兄弟节点。",
    ]
    for snapshot in snapshots:
        lines.extend(format_local_heat_snapshot(snapshot, weight_by_id=weight_by_id))
        lines.append("")
    return "\n".join(lines).rstrip()