xiaoxia-saas/scripts/seed_viral_video_prompts.py

"""爆款视频 5 套 Prompt 模板种子脚本（#2040）。

幂等：以 (prompt_type, version) 为唯一键，存在则更新（UPSERT），重复执行结果一致。
用法：
    python scripts/seed_viral_video_prompts.py            # 自动用应用配置连库
    DATABASE_URL=postgresql+psycopg2://... python scripts/seed_viral_video_prompts.py
"""

from __future__ import annotations

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import sqlalchemy as sa  # noqa: E402

from packages.application.viral_video.prompts import DEFAULT_TEMPLATES  # noqa: E402


def _engine():
    database_url = os.environ.get("DATABASE_URL")
    if database_url:
        return sa.create_engine(database_url)
    # 复用应用自身配置
    from packages.config import get_shared_settings

    url = str(get_shared_settings().database_url)
    return sa.create_engine(url.replace("postgresql+asyncpg://", "postgresql+psycopg2://"))


UPSERT_SQL = sa.text("""
    INSERT INTO viral_video_prompt_templates
        (name, prompt_type, version, system_prompt, user_prompt_template,
         example_output, is_active, updated_at)
    VALUES
        (:name, :prompt_type, :version, :system_prompt, :user_prompt_template,
         :example_output, TRUE, :now_ts)
    ON CONFLICT (prompt_type, version) DO UPDATE SET
        name = EXCLUDED.name,
        system_prompt = EXCLUDED.system_prompt,
        user_prompt_template = EXCLUDED.user_prompt_template,
        example_output = EXCLUDED.example_output,
        is_active = TRUE,
        updated_at = :now_ts
    """)


def seed(engine) -> int:
    count = 0
    from datetime import datetime, timezone

    now_ts = datetime.now(timezone.utc)
    with engine.begin() as conn:
        for item in DEFAULT_TEMPLATES:
            conn.execute(
                UPSERT_SQL,
                {
                    "name": item["name"],
                    "prompt_type": item["prompt_type"],
                    "version": item["version"],
                    "system_prompt": item["system_prompt"],
                    "user_prompt_template": item["user_prompt_template"],
                    "example_output": item["example_output"],
                    "now_ts": now_ts,
                },
            )
            count += 1
    return count


def main() -> int:
    engine = _engine()
    count = seed(engine)
    print(f"seed 完成：{count} 套模板已写入/更新（image_analysis/intent_parsing/copy_fusion/storyboard/review）")
    return 0


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