yingqing/core/backend/apps/assets/review.py

"""火山人像素材库审核编排:真人资产静默送审 + 轮询绿/红状态。

策略(用户定):一团队一素材组,后台静默上传,前端只显示绿盾(active)/红标(failed,提示改提示词重生)。
全部 best-effort:审核未启用/出错都不影响主流程(生图/采用照常)。
"""
import logging
from datetime import timedelta

from django.db import IntegrityError, transaction
from django.utils import timezone

from apps.assets import assets_client
from apps.assets.models import Asset, AssetReviewGroup
from apps.assets.storage import TosStorage

logger = logging.getLogger(__name__)

# processing 超过此时长仍无终态 → 判超时失败,停止再查(防永久 processing + 无限轮询)
_PROCESSING_TIMEOUT = timedelta(minutes=15)

# 火山 Status → 本地 review_status
_STATUS_MAP = {"Active": "active", "Failed": "failed", "Processing": "processing", "Pending": "processing"}


def _asset_url(asset: Asset) -> str:
    """资产主图可公开访问 URL(火山要从 URL 抓图;TOS 签名 URL 即可,会立即抓取)。"""
    f = asset.files.filter(is_primary=True).first() or asset.files.first()
    if f is None:
        return ""
    if f.preview_url:
        return f.preview_url
    try:
        return TosStorage().presigned_get_url(object_key=f.object_key)
    except Exception:  # noqa: BLE001
        return ""


def get_or_create_team_group(team) -> AssetReviewGroup:
    """取/建团队的火山素材组(一团队一组)。并发安全:
    ① DB 行靠 OneToOne 唯一约束去重(只有一个赢,输的 catch 后复用);
    ② 远程建组用 select_for_update 串行化,只有第一个把 remote_group_id 从空写非空的那次才真建组,
    避免并发双建远程组(孤儿)+ 丢送审。"""
    grp = AssetReviewGroup.objects.filter(team=team).first()
    if grp and grp.remote_group_id:
        return grp
    name = f"airshelf-team-{team.id}"
    if grp is None:
        try:
            grp = AssetReviewGroup.objects.create(team=team, name=name, remote_group_id="")
        except IntegrityError:
            grp = AssetReviewGroup.objects.get(team=team)
    if not grp.remote_group_id:
        with transaction.atomic():
            locked = AssetReviewGroup.objects.select_for_update().get(pk=grp.pk)
            if not locked.remote_group_id:
                locked.remote_group_id = assets_client.create_asset_group(name=name, description="AirShelf 真人素材审核")
                locked.save(update_fields=["remote_group_id"])
            grp = locked
    return grp


def submit_asset_for_review(asset: Asset) -> None:
    """真人资产静默送审:建组(若无)→ 传素材 → 标 processing。出错只记日志,不抛。"""
    if not assets_client.is_enabled() or asset.category != Asset.Category.PERSON:
        return
    url = _asset_url(asset)
    if not url:
        return
    try:
        grp = get_or_create_team_group(asset.team)
        remote_id = assets_client.create_asset(group_id=grp.remote_group_id, image_url=url, name=(asset.name or "person")[:64])
        if not remote_id:
            # 火山没回 Id:不要标 processing(否则 remote_id 为空、poll 永远早退、卡死黄),留空可重试
            logger.warning("create_asset 返回空 id,asset %s 暂不送审(可重试)", asset.id)
            return
        asset.review_remote_id = remote_id
        asset.review_status = "processing"
        asset.review_error = ""
        asset.save(update_fields=["review_remote_id", "review_status", "review_error", "updated_at"])
    except Exception as exc:  # noqa: BLE001
        logger.warning("submit_asset_for_review failed for asset %s: %s", asset.id, exc)


def poll_asset_review(asset: Asset) -> str:
    """查单个真人资产审核状态并更新 review_status。返回最新状态。
    只在状态变化时落库(保留 updated_at 作为「进入 processing 的时刻」);processing 超时兜底为 failed。"""
    if not assets_client.is_enabled() or not asset.review_remote_id:
        return asset.review_status
    # 超时兜底:processing 太久(火山卡住 / remote_id 失效每次抛错)→ 判失败,退出永久轮询
    if asset.review_status == "processing" and asset.updated_at and (timezone.now() - asset.updated_at) > _PROCESSING_TIMEOUT:
        asset.review_status = "failed"
        asset.review_error = "审核超时,请重新生成"
        asset.save(update_fields=["review_status", "review_error", "updated_at"])
        return "failed"
    try:
        data = assets_client.get_asset(asset.review_remote_id)
        raw = data.get("Status")
        status = _STATUS_MAP.get(raw)
        if status is None:
            logger.warning("未知审核 Status %r(asset %s),暂按 processing 处理", raw, asset.id)
            status = "processing"
        if status != asset.review_status:  # 只在变化时写,避免每次 poll 刷新 updated_at 让超时永不触发
            asset.review_status = status
            asset.review_error = (data.get("ErrorMessage") or "") if status == "failed" else ""
            asset.save(update_fields=["review_status", "review_error", "updated_at"])
    except Exception as exc:  # noqa: BLE001
        logger.warning("poll_asset_review failed for asset %s: %s", asset.id, exc)
    return asset.review_status


def poll_team_reviews(team) -> dict:
    """轮询该团队所有「审核中」真人资产,更新状态。返回 {asset_id: status}。"""
    out: dict[str, str] = {}
    pending = Asset.objects.filter(
        team=team, category=Asset.Category.PERSON, review_status="processing", is_deleted=False
    )
    for asset in pending:
        out[str(asset.id)] = poll_asset_review(asset)
    return out