SupplyAgent/supply_infra/scheduler/plan_group_batch.py

"""计划组需求明细：物化与读取。"""
from __future__ import annotations

from typing import Any

from sqlalchemy.orm import Session

from supply_infra.db.repositories.demand_belong_category_repo import DemandBelongCategoryRepository
from supply_infra.db.repositories.demand_belong_pool_rel_repo import DemandBelongPoolRelRepository
from supply_infra.db.repositories.multi_demand_pool_di_repo import MultiDemandPoolDiRepository
from supply_infra.db.session import get_session

MAX_DEMANDS_PER_BATCH = 30


def _sort_category_units(
    units: list[tuple[int, int | None, list[Any]]],
) -> list[tuple[int, int | None, list[Any]]]:
    """同父分类相邻，同分类作为不可分割单元。"""
    return sorted(
        units,
        key=lambda item: (
            item[1] is None,
            item[1] if item[1] is not None else -1,
            item[0],
        ),
    )


def pack_category_units(
    units: list[tuple[int, int | None, list[Any]]],
    *,
    max_demands_per_group: int = MAX_DEMANDS_PER_BATCH,
) -> list[list[int]]:
    """将分类需求单元打包为计划组；每组尽量不超过 max_demands_per_group 条需求。"""
    cap = max(1, int(max_demands_per_group))
    groups: list[list[int]] = []
    current: list[int] = []
    current_size = 0

    for category_id, _parent_id, demands in _sort_category_units(units):
        size = len(demands)
        if size <= 0:
            continue

        if size > cap:
            if current:
                groups.append(current)
                current = []
                current_size = 0
            groups.append([category_id])
            continue

        if current and current_size + size > cap:
            groups.append(current)
            current = []
            current_size = 0

        current.append(category_id)
        current_size += size

    if current:
        groups.append(current)

    return groups


def split_even_batches(items: list[Any], *, max_per_batch: int = MAX_DEMANDS_PER_BATCH) -> list[list[Any]]:
    """按总数均分子批次：总数不超过上限则一批；否则递增组数直到每组不超过上限，余数从前组分配。"""
    n = len(items)
    if n == 0:
        return []
    cap = max(1, int(max_per_batch))
    if n <= cap:
        return [items]

    k = 2
    while True:
        base = n // k
        rem = n % k
        max_size = base + 1 if rem > 0 else base
        if max_size <= cap:
            break
        k += 1

    batches: list[list[Any]] = []
    idx = 0
    for i in range(k):
        size = base + (1 if i < rem else 0)
        batches.append(items[idx : idx + size])
        idx += size
    return batches


def _priority_sort_key(name: str, priority_index: dict) -> tuple:
    return (
        priority_index.get(name, {}).get("source_rank_score") is None,
        -float(priority_index.get(name, {}).get("source_rank_score") or 0),
        float(priority_index.get(name, {}).get("global_demand_rank") or float("inf")),
        name,
    )


def resolve_demands_for_category_ids(
    biz_dt: str,
    category_ids: list[int],
    *,
    session: Session | None = None,
) -> list[dict[str, Any]]:
    """按分类节点解析全部需求池记录（pool_id + demand_name）。"""
    selected_ids = list(dict.fromkeys(int(value) for value in category_ids))
    if session is None:
        with get_session() as owned_session:
            return resolve_demands_for_category_ids(
                biz_dt,
                selected_ids,
                session=owned_session,
            )

    belongs = DemandBelongCategoryRepository(session).list_by_category_ids(selected_ids)
    pool_ids_by_belong = DemandBelongPoolRelRepository(session).get_pool_ids_by_belong_ids(
        [int(row.id) for row in belongs]
    )
    pool_ids = sorted({pool_id for values in pool_ids_by_belong.values() for pool_id in values})
    pool_repo = MultiDemandPoolDiRepository(session)
    pool_rows = pool_repo.get_by_ids(pool_ids)
    from agents.demand_grade_agent.tools.demand_priority import build_demand_priority_index

    priority_index = build_demand_priority_index(pool_repo.list_by_biz_dt(biz_dt))
    candidates: list[dict[str, Any]] = []
    for row in pool_rows:
        if row.biz_dt != biz_dt or not row.demand_name:
            continue
        candidates.append({
            "pool_id": int(row.id),
            "demand_name": str(row.demand_name),
        })

    candidates.sort(
        key=lambda item: (
            *_priority_sort_key(item["demand_name"], priority_index),
            item["pool_id"],
        )
    )
    return candidates