cang-ying/video-ai-system/lib/doubao_chat.py

# LIB · 豆包对话模型适配器（剧本拆解/分镜生成）
# D180 · 2026-07-10
"""火山方舟 ARK Chat API → doubao-seed 系列模型"""
import json
import os
import subprocess
import tempfile

from lib.secrets import get_ark_key
ARK_KEY = get_ark_key()
ARK_CHAT_URL = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"

# 可用模型
MODELS = {
    "pro": "doubao-seed-2-1-pro-260628",     # 最强·适合复杂剧本
    "lite": "doubao-seed-2-0-lite-260215",    # 轻量·适合批量
    "deepseek": "deepseek-v4-pro-260425",     # 深度思<E5BAA6><E6809D><EFBFBD>·适合拆解
}

def _curl_api(payload):
    """通过 subprocess curl 调用 API（绕过 Windows urllib 超时问题）"""
    import subprocess, tempfile
    tf = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8')
    json.dump(payload, tf, ensure_ascii=False)
    tf.close()
    try:
        result = subprocess.run([
            "/mingw64/bin/curl", "-s", "--max-time", "180", ARK_CHAT_URL,
            "-H", f"Authorization: Bearer {ARK_KEY}",
            "-H", "Content-Type: application/json",
            "-d", f"@{tf.name}"
        ], capture_output=True, text=True, timeout=180)
        os.unlink(tf.name)
        return json.loads(result.stdout) if result.stdout else {"error": result.stderr}
    except subprocess.TimeoutExpired:
        os.unlink(tf.name)
        return {"error": "timeout"}

def chat(prompt, system="你是一个专业的短剧剧本分析专家", model="pro", temperature=0.3, max_tokens=4096):
    """调用豆包对话模型"""
    payload = {
        "model": MODELS.get(model, model),
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": prompt}
        ],
        "temperature": temperature,
        "max_tokens": max_tokens
    }
    result = _curl_api(payload)
    if "error" in result:
        return result
    choice = result.get("choices", [{}])[0]
    return {
        "content": choice.get("message", {}).get("content", ""),
        "model": result.get("model", ""),
        "tokens": result.get("usage", {}),
        "finish": choice.get("finish_reason", "")
    }

def breakdown_script(script_text, episode_num=1):
    """拆解一集剧本为结构化分镜"""
    prompt = f"""请将以下短剧剧本第{episode_num}集拆解为结构化分镜。 

要求：
1. 按镜头拆分，每个镜头包含：镜头编号、景别（特写/近景/中景/全景/POV）、时长（秒）
2. 提取每个镜头中出现的角色、场景、道具
3. 写出每个镜头的画面描述（50字以内）
4. 标注镜头类型：establishing/action/reaction/closeup/transition

输出JSON格式：
{{
  "episode": {episode_num},
  "shots": [
    {{
      "shot_number": "S01",
      "description": "画面描述",
      "camera": "POV|特写|近景|中景|全景",
      "duration": 6,
      "type": "establishing|action|reaction|closeup|transition",
      "characters": ["角色名"],
      "scenes": ["场景名"],
      "props": ["道具名"],
      "dialogue": null
    }}
  ]
}}

剧本内容：
{script_text}"""
    return chat(prompt, system="你是专业的短剧剧本拆解专家，输出纯JSON，不添加任何解释。")

def generate_keyframes(shot_description, character_refs, scene_refs):
    """为单个镜头生成关键帧 prompt"""
    prompt = f"""基于以下镜头描述，生成即梦4.0图像生成提示词。
    
镜头：{shot_description}
可用角色：{json.dumps(character_refs, ensure_ascii=False)}
可用场景：{json.dumps(scene_refs, ensure_ascii=False)}

要求：
1. 提示词用英文
2. 包含景别、角色位置、场景细节、光照、风格
3. 不超过150词
4. 输出纯提示词，不加任何解释"""
    return chat(prompt, model="pro", temperature=0.5, max_tokens=300)

# ====== CLI ======
if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python doubao_chat.py <command> [args]")
        print("  chat <prompt>    直接对话")
        print("  breakdown <file> 拆解剧本文件")
        sys.exit(1)
    cmd = sys.argv[1]
    if cmd == "chat":
        r = chat(sys.argv[2])
        print(r["content"] if "content" in r else r)
    elif cmd == "breakdown":
        with open(sys.argv[2], "r", encoding="utf-8") as f:
            script = f.read()
        r = breakdown_script(script)
        print(r["content"] if "content" in r else r)