SupplyAgent/supply_infra/aigc/plan_map.py

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, TypeVar

T = TypeVar("T")

# 品类 → (生成计划ID, 发布计划ID)；按 pair 去重后做均匀分发。
AIGC_PLAN_ID_MAP: dict[str, dict[str, str]] = {
    "健康知识": {"生成ID": "20260408092313211598604", "发布ID": "20260408115944193153417"},
    # "历史名人": {"生成ID": "20260408083251311809309", "发布ID": "20260408115139124511126"},
    # "知识科普": {"生成ID": "20260408083824905654920", "发布ID": "20260408115231567509261"},
    "搞笑段子": {"生成ID": "20260408091536533918237", "发布ID": "20260408115842127748387"},
    # "社会风气": {"生成ID": "20260408084318884115213", "发布ID": "20260408115315950129776"},
    "人生忠告": {"生成ID": "20260408085205791658566", "发布ID": "20260408115405410408001"},
    "国际时政": {"生成ID": "20260408090208237400605", "发布ID": "20260408115616925523989"},
    # "生活技巧科普": {"生成ID": "20260408083824905654920", "发布ID": "20260408115231567509261"},
    # "贪污腐败": {"生成ID": "20260408090309503416878", "发布ID": "20260408115653908856043"},
    "民生政策": {"生成ID": "20260408090721867506475", "发布ID": "20260408115727030928177"},
    "对口型表演": {"生成ID": "20260408092122328523262", "发布ID": "20260408115914659162376"},
    "中国战争史": {"生成ID": "20260408090950446586451", "发布ID": "20260408115804931772327"},
    "人财诈骗": {"生成ID": "20260408093140652233649", "发布ID": "20260408120019784463902"},
    "当代正能量人物": {"生成ID": "20260408083148399635274", "发布ID": "20260408115046382803287"},
    "国家科技力量": {"生成ID": "20260408085807674913378", "发布ID": "20260408115542550181196"},
    "国家力量": {"生成ID": "20260408085807674913378", "发布ID": "20260408115542550181196"},
    "通用": {"生成ID": "20260408085649635441036", "发布ID": "20260408115439581604474"},
}


@dataclass(frozen=True)
class AigcPlanPair:
    """一组生成计划 + 发布计划（去重后的分发单元）。"""

    label: str
    produce_plan_id: str
    publish_plan_id: str


def list_unique_plan_pairs(
    plan_map: dict[str, dict[str, str]] | None = None,
) -> list[AigcPlanPair]:
    """从映射表提取不重复的 (生成ID, 发布ID) 计划对。"""
    source = plan_map or AIGC_PLAN_ID_MAP
    seen: set[tuple[str, str]] = set()
    pairs: list[AigcPlanPair] = []
    for label, ids in source.items():
        produce_plan_id = str(ids.get("生成ID") or "").strip()
        publish_plan_id = str(ids.get("发布ID") or "").strip()
        if not produce_plan_id or not publish_plan_id:
            continue
        key = (produce_plan_id, publish_plan_id)
        if key in seen:
            continue
        seen.add(key)
        pairs.append(
            AigcPlanPair(
                label=label,
                produce_plan_id=produce_plan_id,
                publish_plan_id=publish_plan_id,
            )
        )
    return pairs


def distribute_evenly(items: list[T], bucket_count: int) -> list[list[T]]:
    """轮询均匀分配到 bucket_count 个桶。"""
    if bucket_count <= 0:
        raise ValueError("bucket_count 必须大于 0")
    if not items:
        return []
    buckets: list[list[T]] = [[] for _ in range(bucket_count)]
    for index, item in enumerate(items):
        buckets[index % bucket_count].append(item)
    return [bucket for bucket in buckets if bucket]


def chunk_list(items: list[T], size: int) -> list[list[T]]:
    if size <= 0:
        raise ValueError("size 必须大于 0")
    return [items[i : i + size] for i in range(0, len(items), size)]


def assignment_summary(
    plan_pairs: list[AigcPlanPair],
    buckets: list[list[Any]],
) -> list[dict[str, Any]]:
    """返回每个计划分到的数量摘要。"""
    summary: list[dict[str, Any]] = []
    for plan, bucket in zip(plan_pairs, buckets, strict=False):
        if not bucket:
            continue
        summary.append(
            {
                "plan_label": plan.label,
                "produce_plan_id": plan.produce_plan_id,
                "publish_plan_id": plan.publish_plan_id,
                "video_count": len(bucket),
            }
        )
    return summary