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
    OPENROUTER_IMAGE_MODEL,
    build_ai_image_object_key,
    build_generation_prompts,
    generate_image_bytes,
    sanitize_video_description,
    upload_image_to_oss,
)
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("--limit", type=int, default=1)
    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 ""
    prompts = [
        (p.prompt_type, p.prompt_text)
        for p in build_generation_prompts(
            video_id=args.video_id,
            title=args.title,
            category=args.category,
            features=features,
            sanitized_description=sanitized_description,
        )
    ][: max(1, args.limit)]

    outputs = []
    for prompt_type, prompt_text in prompts:
        image_bytes, content_type, raw = generate_image_bytes(prompt_text, model=args.model)
        ext = ".png" if content_type == "image/png" else ".jpg"
        object_key = build_ai_image_object_key(
            account_id="debug",
            landing_video_id=args.video_id,
            prompt_type=prompt_type,
            extension=ext,
            debug=True,
        )
        oss_url = upload_image_to_oss(image_bytes, content_type, object_key)
        outputs.append({
            "prompt_type": prompt_type,
            "model": args.model,
            "content_type": content_type,
            "oss_url": oss_url,
            "object_key": object_key,
            "prompt_text": prompt_text,
            "raw_response_id": raw.get("id"),
        })

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


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