Agent/examples/auto_put_ad_mini/tools/ai_material_review.py

"""AI review for generated image materials.

This module is intentionally separate from image generation. It can be used
both inline after generation and later for rescanning historical assets.
"""

from __future__ import annotations

import json
import logging
import os
import re
from dataclasses import dataclass
from typing import Any

import httpx

logger = logging.getLogger(__name__)

OPENROUTER_CHAT_COMPLETIONS_URL = os.getenv(
    "OPENROUTER_CHAT_COMPLETIONS_URL",
    "https://openrouter.ai/api/v1/chat/completions",
)
AI_MATERIAL_REVIEW_MODEL = os.getenv("AI_MATERIAL_REVIEW_MODEL", "google/gemini-3-flash-preview")


@dataclass(frozen=True)
class MaterialReviewResult:
    status: str
    score: int
    reason: str
    risk_tags: list[str]
    ocr_text: str
    raw: dict[str, Any]

    @property
    def passed(self) -> bool:
        return self.status == "pass"


def _openrouter_api_key() -> str:
    key = os.getenv("OPEN_ROUTER_API_KEY") or os.getenv("OPENROUTER_API_KEY")
    if not key:
        raise RuntimeError("缺少 OPENROUTER_API_KEY/OPEN_ROUTER_API_KEY,无法进行AI素材审核")
    return key


def _extract_chat_completion_text(data: dict) -> str:
    choices = data.get("choices") or []
    if not choices:
        raise RuntimeError("OpenRouter AI审核响应缺 choices")
    message = choices[0].get("message") or {}
    content = message.get("content")
    if isinstance(content, str):
        return content.strip()
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and isinstance(item.get("text"), str):
                parts.append(item["text"])
        return "\n".join(parts).strip()
    return ""


def _extract_json_object(text: str) -> dict[str, Any]:
    raw = str(text or "").strip()
    if raw.startswith("```"):
        raw = re.sub(r"^```(?:json)?", "", raw).strip()
        raw = re.sub(r"```$", "", raw).strip()
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError:
        match = re.search(r"\{.*\}", raw, flags=re.S)
        if not match:
            raise
        parsed = json.loads(match.group(0))
    if not isinstance(parsed, dict):
        raise ValueError("AI审核模型返回不是JSON object")
    return parsed


def _normalize_status(value: Any) -> str:
    status = str(value or "").strip().lower()
    if status in {"pass", "approve", "approved", "通过"}:
        return "pass"
    if status in {"hold", "review", "manual_review", "人工复核"}:
        return "hold"
    return "reject"


def review_generated_material(
    *,
    image_url: str,
    prompt_type: str,
    prompt_text: str,
    feature_hits: list[dict],
    model: str = AI_MATERIAL_REVIEW_MODEL,
) -> MaterialReviewResult:
    """Review one generated material image with multimodal model."""
    payload = {
        "task": "审核一张AI生成的腾讯广告信息流中文封面图是否可进入人工投放审批候选",
        "review_goals": [
            "判断图片标题是否有乱码、错字、悬挂标点、分裂文案、不可读文字",
            "判断是否包含强领取承诺、专家背书、假官方、假按钮、二维码、下载/播放按钮伪装",
            "判断是否涉及医疗疗效、健康恐吓、违法血腥、迷信、低俗、名人肖像、外国人物",
            "判断图片是否和视频主题及pattern相关",
            "判断是否适合60-75岁中老年用户的信息流广告封面",
        ],
        "hard_reject_rules": [
            "出现乱码、明显错字、不可读标题、标题被截断或标点悬挂",
            "出现国家发钱、官方发放、已到账、马上到账、保证到账、立即领取、直接领取、点这里领取等明确承诺或行动诱导",
            "出现假微信/假聊天/假按钮/二维码/下载按钮/播放按钮伪装",
            "出现医疗疗效承诺、健康恐吓、医院药品专家治疗背书",
            "出现真实名人肖像感、外国人物、logo、水印、畸形人物",
            "图片与视频主题或pattern明显不相关",
        ],
        "soft_review_rules": [
            "政策、退休补贴、养老金、清单讲解类素材中,领取、领全、能领、补贴等词不单独构成拒绝理由",
            "如果标题只是提出疑问或提示核对,例如“这几项你知道吗”“你问清了吗”“你领全了吗”,可判为hold或pass,不要仅因词语本身reject",
            "专家详细解读、老师讲解、讲清楚、一图看懂不单独构成拒绝理由;只有和医疗疗效、投资收益、官方承诺、领取承诺强绑定时才reject",
            "如果存在轻度政策诱导但画面质量好、主题相关,优先hold;只有明确承诺结果或伪官方才reject",
        ],
        "prompt_type": prompt_type,
        "feature_hits": feature_hits[:3],
        "output_schema": {
            "decision": "pass/reject/hold",
            "score": "0-100整数",
            "ocr_text": "识别到的主标题文字",
            "risk_tags": ["命中的风险标签"],
            "reason": "一句话说明审核结论",
            "theme_relevance": "high/medium/low",
            "title_quality": "good/medium/bad",
            "visual_quality": "good/medium/bad",
        },
    }
    messages = [
        {
            "role": "system",
            "content": (
                "你是腾讯广告中文信息流素材的AI预审员。"
                "只输出JSON object,不要解释正文。"
                "审核要区分硬性违规和轻度风险:标题不可读、明确承诺、伪官方或主题不相关才reject;"
                "政策信息讲解中的领取相关疑问或专家解读表述,不应仅凭单个词直接reject,可按风险程度给hold或pass。"
            ),
        },
        {
            "role": "user",
            "content": [
                {"type": "text", "text": json.dumps(payload, ensure_ascii=False)},
                {"type": "image_url", "image_url": {"url": image_url}},
            ],
        },
    ]
    resp = httpx.post(
        OPENROUTER_CHAT_COMPLETIONS_URL,
        headers={
            "Authorization": f"Bearer {_openrouter_api_key()}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        json={
            "model": model,
            "messages": messages,
            "temperature": 0.1,
            "max_tokens": 900,
        },
        timeout=90,
    )
    resp.raise_for_status()
    parsed = _extract_json_object(_extract_chat_completion_text(resp.json()))
    try:
        score = int(float(parsed.get("score", 0)))
    except (TypeError, ValueError):
        score = 0
    status = _normalize_status(parsed.get("decision"))
    if score < 60 and status == "pass":
        status = "hold"
    return MaterialReviewResult(
        status=status,
        score=max(0, min(100, score)),
        reason=str(parsed.get("reason") or "").strip(),
        risk_tags=[str(v) for v in parsed.get("risk_tags") or [] if str(v).strip()],
        ocr_text=str(parsed.get("ocr_text") or "").strip(),
        raw=parsed,
    )


def update_material_review_result(
    material_id: int,
    result: MaterialReviewResult,
    *,
    model: str = AI_MATERIAL_REVIEW_MODEL,
) -> None:
    from db.connection import get_connection

    status = "generated" if result.passed else result.status
    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                UPDATE ai_generated_material
                SET status=%s,
                    ai_review_status=%s,
                    ai_review_score=%s,
                    ai_review_model=%s,
                    ai_review_reason=%s,
                    ai_review_json=%s,
                    ai_reviewed_at=CURRENT_TIMESTAMP,
                    error=CASE WHEN %s='pass' THEN error ELSE %s END,
                    updated_at=CURRENT_TIMESTAMP
                WHERE id=%s
                """,
                (
                    status,
                    result.status,
                    result.score,
                    model,
                    result.reason[:2000],
                    json.dumps(result.raw, ensure_ascii=False, default=str)[:16000000],
                    result.status,
                    result.reason[:2000],
                    int(material_id),
                ),
            )
        conn.commit()
    finally:
        conn.close()