SupplyAgent/supply_infra/services/video_discovery_service.py

"""视频发现运行的 Service 层：session 与 ORM 仅在此模块内出现。"""
from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

from sqlalchemy.exc import OperationalError, ProgrammingError

from supply_infra.db.repositories.video_discovery_repo import VideoDiscoveryRepository
from supply_infra.db.session import get_session

_SKIP_STATUSES = frozenset({"finished"})


class RunNotFoundError(LookupError):
    """run_id 在 video_discovery_run 中不存在。"""


def format_db_error(error: Exception) -> str:
    if isinstance(error, (OperationalError, ProgrammingError)):
        return (
            f"数据库表尚未初始化或不可用: {error}。"
            "请先运行 `.venv/bin/python -m supply_infra.db`。"
        )
    return str(error)


def _load_json(value: str | None, default: Any) -> Any:
    if not value:
        return default
    try:
        return json.loads(value)
    except (TypeError, ValueError):
        return default


def _serialize_run(run: Any) -> dict[str, Any]:
    return {
        "run_id": run.run_id,
        "biz_dt": run.biz_dt,
        "demand_grade_id": run.demand_grade_id,
        "demand_word": run.demand_word,
        "seed_video_id": run.seed_video_id,
        "seed_video_title": run.seed_video_title,
        "intent_summary": run.intent_summary,
        "status": run.status,
        "search_count": int(run.search_count or 0),
        "primary_count": int(run.primary_count or 0),
        "stop_reason": run.stop_reason,
    }


def _serialize_search(item: Any) -> dict[str, Any]:
    return {
        "search_id": int(item.id),
        "keyword": item.keyword,
        "query_reason": item.query_reason,
        "source_type": item.source_type,
        "source_value": item.source_value,
        "parent_search_id": item.parent_search_id,
        "provider": item.provider,
        "provider_state": _load_json(item.provider_state_json, {}),
        "content_type": item.content_type,
        "sort_type": item.sort_type,
        "publish_time": item.publish_time,
        "cursor": item.cursor,
        "page_no": item.page_no,
        "results_count": item.results_count,
        "new_candidate_count": item.new_candidate_count,
        "has_more": bool(item.has_more),
        "next_cursor": item.next_cursor,
        "status": item.status,
        "error_message": item.error_message,
    }


def _serialize_candidate(item: Any) -> dict[str, Any]:
    decision_bucket = (
        "pending_evaluation"
        if item.decision_bucket == "unreviewed"
        else item.decision_bucket
    )
    return {
        "candidate_id": int(item.id),
        "search_id": int(item.search_id) if item.search_id is not None else None,
        "aweme_id": item.aweme_id,
        "title": item.title,
        "content_link": item.content_link,
        "author_name": item.author_name,
        "author_sec_uid": item.author_sec_uid,
        "source_keywords": _load_json(item.source_keywords_json, []),
        "source_search_ids": _load_json(item.source_search_ids_json, []),
        "tags": _load_json(item.tags_json, []),
        "play_count": item.play_count,
        "like_count": item.like_count,
        "comment_count": item.comment_count,
        "collect_count": item.collect_count,
        "share_count": item.share_count,
        "relevance_score": (
            float(item.relevance_score) if item.relevance_score is not None else None
        ),
        "elder_score": float(item.elder_score) if item.elder_score is not None else None,
        "share_score": float(item.share_score) if item.share_score is not None else None,
        "value_score": float(item.value_score) if item.value_score is not None else None,
        "decision_bucket": decision_bucket,
        "content_age_evidence": _load_json(item.content_age_evidence_json, {}),
        "account_age_evidence": _load_json(item.account_age_evidence_json, {}),
        "age_normalization": _load_json(item.age_normalization_json, {}),
        "decision_reason": item.decision_reason,
    }


@dataclass(frozen=True)
class PublishableCandidate:
    """脱离 ORM 的发布候选快照。"""

    id: int
    aweme_id: str


class VideoDiscoveryService:
    """视频发现运行、搜索轨迹与候选的读写入口。"""

    def lookup_run(self, run_id: str) -> dict[str, Any] | None:
        with get_session() as session:
            run = VideoDiscoveryRepository(session).get_run(run_id)
            if run is None:
                return None
            return _serialize_run(run)

    def has_candidates(self, run_id: str) -> bool:
        with get_session() as session:
            return VideoDiscoveryRepository(session).has_candidates(run_id)

    def create_run(self, values: dict[str, Any]) -> dict[str, Any]:
        with get_session() as session:
            VideoDiscoveryRepository(session).create_run(values)
        return {
            "run_id": str(values["run_id"]),
            "status": str(values.get("status") or "running"),
        }

    def require_run(self, run_id: str) -> dict[str, Any]:
        run = self.lookup_run(run_id)
        if run is None:
            raise RunNotFoundError(f"run_id 不存在: {run_id}")
        return run

    def save_search_page(
        self,
        run_id: str,
        search_values: dict[str, Any],
        candidate_rows: list[dict[str, Any]],
    ) -> dict[str, Any]:
        with get_session() as session:
            repo = VideoDiscoveryRepository(session)
            if repo.get_run(run_id) is None:
                raise RunNotFoundError(f"run_id 不存在: {run_id}")
            search, candidates = repo.save_search_page(search_values, candidate_rows)
            return {
                **_serialize_search(search),
                "run_id": run_id,
                "new_candidate_count": len(candidates),
                "candidates": [
                    _serialize_candidate(candidate) for candidate in candidates
                ],
            }

    def update_candidates(
        self,
        run_id: str,
        rows: list[dict[str, Any]],
    ) -> dict[str, Any]:
        with get_session() as session:
            repo = VideoDiscoveryRepository(session)
            if repo.get_run(run_id) is None:
                raise RunNotFoundError(f"run_id 不存在: {run_id}")
            candidates = repo.update_candidates(run_id, rows)
            return {
                "updated_count": len(candidates),
                "candidates": [
                    _serialize_candidate(candidate) for candidate in candidates
                ],
            }

    def update_run_status(
        self,
        run_id: str,
        *,
        status: str,
        intent_summary: str | None = None,
        stop_reason: str | None = None,
    ) -> dict[str, Any]:
        with get_session() as session:
            repo = VideoDiscoveryRepository(session)
            if repo.get_run(run_id) is None:
                raise RunNotFoundError(f"run_id 不存在: {run_id}")
            run = repo.finish_run(
                run_id,
                status=status,
                intent_summary=intent_summary,
                stop_reason=stop_reason,
            )
            return _serialize_run(run)

    def get_full_state(
        self,
        run_id: str,
        *,
        include_rejected: bool = True,
        limit: int = 100,
    ) -> dict[str, Any]:
        with get_session() as session:
            repo = VideoDiscoveryRepository(session)
            run = repo.get_run(run_id)
            if run is None:
                raise RunNotFoundError(f"run_id 不存在: {run_id}")
            searches = repo.list_searches(run_id)
            buckets = (
                ("primary", "rejected", "pending_evaluation", "unreviewed")
                if include_rejected
                else ("primary", "pending_evaluation", "unreviewed")
            )
            candidates = repo.list_candidates(
                run_id, buckets=buckets, limit=max(1, min(int(limit), 500))
            )
            return {
                "run": _serialize_run(run),
                "searches": [_serialize_search(item) for item in searches],
                "candidates": [_serialize_candidate(item) for item in candidates],
            }

    def prepare_scheduled_run(
        self,
        *,
        biz_dt: str,
        demand_grade_id: int,
        values: dict[str, Any],
        force: bool = False,
    ) -> tuple[str | None, str | None]:
        """预创建或复用调度运行；返回 (run_id, skip_reason)。"""
        with get_session() as session:
            repo = VideoDiscoveryRepository(session)
            existing = repo.get_by_biz_dt_and_grade(biz_dt, demand_grade_id)
            if existing is not None and not force:
                already_done = (
                    existing.status in _SKIP_STATUSES
                    or repo.has_candidates(str(existing.run_id))
                )
                if already_done:
                    return None, (
                        f"biz_dt={biz_dt} demand_grade_id={demand_grade_id} "
                        f"已执行过 run_id={existing.run_id} status={existing.status}"
                    )
            run_id = existing.run_id if existing is not None else str(values["run_id"])
            payload = {**values, "run_id": run_id}
            repo.upsert_scheduled_run(payload)
            return run_id, None

    def list_skip_grade_ids(self, biz_dt: str) -> set[int]:
        with get_session() as session:
            return VideoDiscoveryRepository(session).list_skip_grade_ids(biz_dt)

    def mark_run_failed(self, run_id: str, *, stop_reason: str | None = None) -> None:
        with get_session() as session:
            VideoDiscoveryRepository(session).mark_run_failed(
                run_id,
                stop_reason=stop_reason,
            )

    def count_passed_videos(self, biz_dt: str) -> int:
        with get_session() as session:
            return VideoDiscoveryRepository(session).count_passed_videos(biz_dt)

    def list_publishable_candidates(
        self,
        *,
        run_id: str | None = None,
        biz_dt: str | None = None,
        skip_published: bool = True,
        limit: int | None = None,
    ) -> list[PublishableCandidate]:
        with get_session() as session:
            rows = VideoDiscoveryRepository(session).list_publishable_candidates(
                run_id=run_id,
                biz_dt=biz_dt,
                skip_published=skip_published,
                limit=limit,
            )
            return [
                PublishableCandidate(id=int(item.id), aweme_id=str(item.aweme_id))
                for item in rows
            ]

    def mark_candidates_aigc_plans(
        self,
        candidate_ids: list[int],
        *,
        crawler_plan_id: str,
        produce_plan_id: str,
        publish_plan_id: str,
        plan_label: str,
    ) -> int:
        with get_session() as session:
            return VideoDiscoveryRepository(session).mark_candidates_aigc_plans(
                candidate_ids,
                crawler_plan_id=crawler_plan_id,
                produce_plan_id=produce_plan_id,
                publish_plan_id=publish_plan_id,
                plan_label=plan_label,
            )


_default_service: VideoDiscoveryService | None = None


def get_video_discovery_service() -> VideoDiscoveryService:
    global _default_service
    if _default_service is None:
        _default_service = VideoDiscoveryService()
    return _default_service