SupplyAgent/supply_infra/pipeline/registry.py

from __future__ import annotations

from collections.abc import Callable
from typing import Any

from supply_infra.pipeline.contracts import StepContext

StepHandler = Callable[[StepContext], dict[str, Any]]


def _global_tree(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.sync_global_tree_odps_to_mysql import (
        sync_global_tree_odps_to_mysql,
    )

    return sync_global_tree_odps_to_mysql(
        partition_date=str(context.date_snapshot["global_tree_partition"])
    )


def _demand_source(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.sync import _sync_pool_rows

    return _sync_pool_rows(str(context.date_snapshot["demand_pool_partition"]))


def _demand_classify(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.sync import _classify_words

    return _classify_words(context.biz_dt)


def _demand_rel(_context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.belong_rel import (
        sync_demand_belong_pool_rel,
    )

    return sync_demand_belong_pool_rel()


def _real_metrics(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.sync import enrich_real_rov_vov_7d

    return enrich_real_rov_vov_7d(context.biz_dt)


def _popularity(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.sync import compute_popularity_stats

    return compute_popularity_stats(context.biz_dt)


def _tree_weight(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.tree_weight import (
        compute_category_tree_weight,
    )

    return compute_category_tree_weight(context.biz_dt)


def _source_videos(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.demand_pool.videos import sync_multi_demand_videos

    return sync_multi_demand_videos(
        limit=None,
        decode_dt=str(context.date_snapshot["video_decode_partition"]),
    )


def _grade(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.grade_demand_pool import grade_demand_pool

    return grade_demand_pool(context.biz_dt)


def _expand(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.jobs.expand_demand_from_video_points import (
        expand_demand_from_video_points,
    )

    return expand_demand_from_video_points(
        context.biz_dt,
        workers=5,
    )


def _discover(context: StepContext) -> dict[str, Any]:
    from supply_infra.scheduler.constants import PIPELINE_FIND_AGENT_WORKERS
    from supply_infra.scheduler.jobs.discover_videos_from_demands import (
        discover_videos_from_demands,
    )

    return discover_videos_from_demands(
        context.biz_dt,
        workers=PIPELINE_FIND_AGENT_WORKERS,
    )


STEP_REGISTRY: dict[str, StepHandler] = {
    "global_tree_sync": _global_tree,
    "demand_pool_source_sync": _demand_source,
    "demand_classify": _demand_classify,
    "demand_belong_rel_sync": _demand_rel,
    "real_metrics_sync": _real_metrics,
    "popularity_stats": _popularity,
    "category_tree_weight": _tree_weight,
    "source_video_sync": _source_videos,
    "demand_grade": _grade,
    "demand_expand": _expand,
    "video_discovery": _discover,
}


def execute_registered_step(context: StepContext) -> dict[str, Any]:
    handler = STEP_REGISTRY.get(context.step_key)
    if handler is None:
        raise KeyError(f"Unknown pipeline step: {context.step_key}")
    payload = handler(context)
    if not isinstance(payload, dict):
        return {"success": True, "result": payload}
    return payload