SupplyAgent/supply_infra/scheduler/auto_assign_grade_plan.py

"""按需求数量自动划分分级计划组（替代统筹 Agent 的默认定时任务路径）。"""
from __future__ import annotations

import logging
from typing import Any

from agents.demand_grade_orchestrator_agent.common.assignment import (
    resolve_planning_state,
)
from agents.demand_grade_orchestrator_agent.common.plan_persist import persist_groups_one_by_one
from agents.demand_grade_orchestrator_agent.common.plan_record import prepare_grade_groups
from agents.demand_grade_orchestrator_agent.common.tree_state import load_tree_state
from supply_infra.scheduler.plan_group_batch import (
    MAX_DEMANDS_PER_BATCH,
    pack_category_units,
    resolve_demands_for_category_ids,
)

logger = logging.getLogger(__name__)

_DEFAULT_GROUPING_STRATEGY = "同分类与同父分类优先，每组约30个需求"
_DEFAULT_REASON = "代码自动分配"
_DEFAULT_TRAITS = _DEFAULT_GROUPING_STRATEGY


def _parent_id(category_id: int, by_id: dict[int, Any]) -> int | None:
    row = by_id.get(category_id)
    if row is None:
        return None
    parent = getattr(row, "parent_id", None)
    if parent in (None, 0):
        return None
    return int(parent)


def build_auto_assign_groups(
    biz_dt: str,
    *,
    unassigned_category_ids: list[int] | None = None,
    max_demands_per_group: int = MAX_DEMANDS_PER_BATCH,
) -> list[dict[str, Any]]:
    """为未分批分类节点生成自动分配批次（仅 category_ids，供 prepare_grade_groups 使用）。"""
    planning_state = resolve_planning_state(biz_dt)
    category_ids = list(unassigned_category_ids or planning_state["unassigned_category_ids"])
    if not category_ids:
        return []

    by_id, _children, _weights = load_tree_state(biz_dt)
    units: list[tuple[int, int | None, list[dict[str, Any]]]] = []
    for category_id in category_ids:
        demands = resolve_demands_for_category_ids(biz_dt, [category_id])
        if not demands:
            continue
        units.append((category_id, _parent_id(category_id, by_id), demands))

    packed = pack_category_units(units, max_demands_per_group=max_demands_per_group)
    return [
        {
            "category_ids": group_category_ids,
            "planning_reason": _DEFAULT_REASON,
            "shared_traits": _DEFAULT_TRAITS,
        }
        for group_category_ids in packed
    ]


def auto_assign_daily_grade_plan(
    *,
    biz_dt: str,
    max_demands_per_group: int = MAX_DEMANDS_PER_BATCH,
) -> dict[str, Any]:
    """自动划分并落库当天分级计划组（每组约 30 个需求，同分类/同父分类优先）。"""
    planning_state = resolve_planning_state(biz_dt)
    if not planning_state["can_plan_more"]:
        logger.info(
            "跳过自动分配: biz_dt=%s reason=%s",
            biz_dt,
            planning_state["skip_reason"],
        )
        return {
            "skipped": True,
            "reason": planning_state["skip_reason"],
            "planning_state": planning_state,
        }

    groups = build_auto_assign_groups(
        biz_dt,
        unassigned_category_ids=planning_state["unassigned_category_ids"],
        max_demands_per_group=max_demands_per_group,
    )
    if not groups:
        logger.info(
            "跳过自动分配：无待分配需求 biz_dt=%s unassigned_categories=%s",
            biz_dt,
            len(planning_state["unassigned_category_ids"]),
        )
        return {
            "skipped": True,
            "reason": "无待分配需求",
            "planning_state": planning_state,
        }

    logger.info(
        "执行自动分配: biz_dt=%s unassigned_categories=%s planned_groups=%s max_demands_per_group=%s",
        biz_dt,
        len(planning_state["unassigned_category_ids"]),
        len(groups),
        max_demands_per_group,
    )

    prepared = prepare_grade_groups(
        biz_dt,
        _DEFAULT_GROUPING_STRATEGY,
        groups,
        assigned_category_ids=set(planning_state["assigned_category_ids"]),
    )
    persist_result = persist_groups_one_by_one(biz_dt, prepared)
    return {
        "skipped": False,
        "planning_state": planning_state,
        "prepared_group_count": len(prepared.get("groups") or []),
        "persist_result": persist_result,
    }