Agent/examples/auto_put_ad_mini/debug_generate_ai_material.py

"""手动调试 AI 图片生成和 OSS 上传。

默认只生成图片并上传 OSS,打印 JSON 输出;不会创建广告/创意。

示例:
  python debug_generate_ai_material.py --video-id 71187017
需要环境变量:
  OPENROUTER_API_KEY
  ALIYUN_OSS_ENDPOINT
  ALIYUN_OSS_BUCKET
  ALIYUN_OSS_ACCESS_KEY_ID
  ALIYUN_OSS_ACCESS_KEY_SECRET
  AI_IMAGE_PUBLIC_BASE_URL
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from dotenv import load_dotenv

_HERE = Path(__file__).parent
load_dotenv(_HERE / ".env")
sys.path.insert(0, str(_HERE.parent.parent))
sys.path.insert(0, str(_HERE))

from tools.ai_generated_material import (  # noqa: E402
    AI_IMAGE_PATTERN_PLACEMENT,
    AI_IMAGE_PATTERN_TOP_K,
    OPENROUTER_TEXT_MODEL,
    OPENROUTER_IMAGE_MODEL,
    build_ai_image_object_key,
    build_generation_prompts,
    build_pattern_generation_prompts,
    generate_image_bytes,
    insert_and_review_generated_material,
    sanitize_video_description,
    upload_image_to_oss,
)
from tools.video_recall import LandingVideo  # noqa: E402
from tools.video_feature_query import read_cached_video_element_features  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--video-id", type=int, default=0)
    parser.add_argument("--title", default="")
    parser.add_argument("--category", default="")
    parser.add_argument("--model", default=OPENROUTER_IMAGE_MODEL)
    parser.add_argument("--text-model", default=OPENROUTER_TEXT_MODEL)
    parser.add_argument("--limit", type=int, default=1)
    parser.add_argument("--account-id", type=int, default=0)
    parser.add_argument("--adgroup-id", type=int, default=0)
    parser.add_argument("--crowd-package", default="")
    parser.add_argument("--placement", default=AI_IMAGE_PATTERN_PLACEMENT)
    parser.add_argument("--pattern-top-k", type=int, default=AI_IMAGE_PATTERN_TOP_K)
    parser.add_argument("--use-pattern-selector", action="store_true")
    parser.add_argument("--write-db", action="store_true")
    args = parser.parse_args()

    if not args.video_id:
        raise SystemExit("必须提供 --video-id")
    features_by_vid = read_cached_video_element_features([args.video_id])
    features = features_by_vid.get(args.video_id) or []
    topic = next((f.standard_element for f in features if f.element_dimension == "解构选题"), "")
    sanitized_description = sanitize_video_description(topic) if topic else ""
    if args.use_pattern_selector:
        built_prompts = build_pattern_generation_prompts(
            video_id=args.video_id,
            title=args.title,
            category=args.category,
            features=features,
            sanitized_description=sanitized_description,
            crowd_package=args.crowd_package,
            placement=args.placement,
            top_k=args.pattern_top_k,
            text_model=args.text_model,
        )
    else:
        built_prompts = build_generation_prompts(
            video_id=args.video_id,
            title=args.title,
            category=args.category,
            features=features,
            sanitized_description=sanitized_description,
        )
    prompts = built_prompts[: max(1, args.limit)]

    outputs = []
    landing = LandingVideo(
        video_id=args.video_id,
        title=args.title,
        cover_url="",
        video_url="",
        score=0,
        rov=0,
        sim=0,
        visit_uv=0,
        category=args.category,
        standard_element="",
        category_name="",
        demand_content_title="",
        demand_content_topic="",
        demand_content_id="",
        demand_type="",
        point_type="",
        dimension="",
        experiment_id="",
        raw={},
    )
    for prompt in prompts:
        image_bytes, content_type, raw = generate_image_bytes(prompt.prompt_text, model=args.model)
        ext = ".png" if content_type == "image/png" else ".jpg"
        object_key = build_ai_image_object_key(
            account_id=args.account_id or "debug",
            landing_video_id=args.video_id,
            prompt_type=prompt.prompt_type,
            extension=ext,
            debug=not args.write_db,
        )
        oss_url = upload_image_to_oss(image_bytes, content_type, object_key)
        asset_id = None
        review_output = None
        if args.write_db:
            asset, review = insert_and_review_generated_material(
                account_id=args.account_id,
                adgroup_id=args.adgroup_id,
                crowd_package=args.crowd_package or "debug",
                landing=landing,
                prompt=prompt,
                model=args.model,
                object_key=object_key,
                oss_url=oss_url,
                raw_response=raw,
            )
            asset_id = asset.id
            review_output = {
                "status": review.status,
                "score": review.score,
                "reason": review.reason,
                "risk_tags": review.risk_tags,
                "ocr_text": review.ocr_text,
            }
        outputs.append({
            "asset_id": asset_id,
            "ai_review": review_output,
            "prompt_type": prompt.prompt_type,
            "model": args.model,
            "text_model": args.text_model,
            "content_type": content_type,
            "oss_url": oss_url,
            "object_key": object_key,
            "prompt_text": prompt.prompt_text,
            "feature_hits": prompt.feature_hits,
            "raw_response_id": raw.get("id"),
        })

    print(json.dumps(outputs, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())