SupplyAgent/supply_infra/pipeline/registry.py

from __future__ import annotations

import hashlib
import json
from collections.abc import Callable
from pathlib import Path
from typing import Any

from supply_infra.config import get_infra_settings
from supply_infra.db.repositories.pipeline_outbox_repo import PipelineOutboxRepository
from supply_infra.db.session import get_session
from supply_infra.pipeline.contracts import StepContext

StepHandler = Callable[[StepContext], dict[str, Any]]
_REPO_ROOT = Path(__file__).resolve().parents[2]
_MAX_INLINE_EFFECT_BYTES = 512_000


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,
    )


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

    settings = get_infra_settings()
    payload = publish_videos_from_discovery(biz_dt=context.biz_dt)
    publish_success = bool(payload.get("success", True))
    canonical = json.dumps(
        payload,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        default=str,
    )
    encoded = canonical.encode("utf-8")
    payload_hash = hashlib.sha256(encoded).hexdigest()
    idempotency_key = f"aigc_write_record:{context.biz_dt}:{payload_hash}"
    payload_uri: str | None = None
    stored_payload: dict[str, Any] | None = payload
    if len(encoded) > _MAX_INLINE_EFFECT_BYTES:
        log_dir = Path(settings.pipeline_log_dir)
        if not log_dir.is_absolute():
            log_dir = _REPO_ROOT / log_dir
        effect_dir = log_dir / context.run_id / "effects"
        effect_dir.mkdir(parents=True, exist_ok=True)
        effect_path = effect_dir / f"{context.step_run_id}-{payload_hash}.json"
        temporary_path = effect_path.with_suffix(".json.tmp")
        temporary_path.write_bytes(encoded)
        temporary_path.replace(effect_path)
        payload_uri = str(effect_path)
        stored_payload = {
            "externalized": True,
            "payload_bytes": len(encoded),
            "payload_hash": payload_hash,
        }
    with get_session() as session:
        record = PipelineOutboxRepository(session).record_effect(
            run_id=context.run_id,
            step_run_id=context.step_run_id,
            effect_type="aigc_write_record",
            idempotency_key=idempotency_key,
            payload_hash=payload_hash,
            payload=stored_payload,
            payload_uri=payload_uri,
        )
        outbox_id = record.outbox_id

    result: dict[str, Any] = {
        "success": publish_success,
        "effect_recorded": True,
        "external_request_made": int(payload.get("candidate_count", 0) or 0) > 0,
        "outbox_id": outbox_id,
        "payload_hash": payload_hash,
        "payload_uri": payload_uri,
        "candidate_count": payload.get("candidate_count", 0),
        "batch_count": payload.get("batch_count", 0),
        "failed_batch_count": payload.get("failed_batch_count", 0),
    }
    if not publish_success:
        result["error"] = str(
            payload.get("error")
            or f"AIGC publish failed ({result['failed_batch_count']} batch(es))"
        )
        result["error_code"] = "aigc_publish_failed"
    return result


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,
    "aigc_write_record": _aigc_write_record,
}


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