Agent/examples/process_pipeline/script/extract_workflow.py

"""
逐 case 提取 workflow + fragments (v6版本)

从 case.json 读取，按 index 遍历每个 case，
调用 LLM 同时提取 workflow（薄壳 steps）和 fragments（原子操作，含完整 capability 字段），
按 index 原位回填到 case.json

v6 架构特性：
- workflow.steps 是薄壳：step_id / order / phase / relation / body，不含 capability 字段
- fragments 是原子操作列表：每个 fragment 含完整 capability 字段 + workflow_step_ref + is_alternative_to
- 步内多原子操作 + 步内 alternative 都在 fragment 层表达
- standalone fragment（workflow_step_ref=null）用于无 workflow 上下文的能力提及
"""

import asyncio
import json
from pathlib import Path
from typing import Any, Dict, Optional, List

from examples.process_pipeline.script.llm_helper import call_llm_with_retry


# v5 词库文件路径
SCRIPT_DIR = Path(__file__).resolve().parent
METHOD_VOCAB_PATH = SCRIPT_DIR / "resource" / "method_vocab_v5.json"

# 默认词库（如果文件不存在时使用）
DEFAULT_METHOD_VOCAB = {
    "流程角色": [
        "生成指令", "编辑指令", "约束条件", "参考素材", "控制信号",
        "区域控制", "参数配置", "模型资源", "源素材", "中间产物",
        "成品", "模板", "评估结果"
    ],
    "模态": ["文本", "图片", "视频", "音频", "特征点", "参数", "模型", "向量", "表格"],
    "主动作": [
        "生成", "编辑", "提取", "改写", "合成", "修复", "增强",
        "训练", "评估", "剪辑", "模板化", "排版", "转写", "配音",
        "匹配", "扩展", "导出"
    ],
    "动作方式": [
        "直接生成", "一致性保持", "结构约束", "质量收束", "局部重绘",
        "扩图", "换背景", "提示词反推", "模板化", "多图融合", "清晰化",
        "风格迁移", "常规编辑", "变体生成", "动画化", "镜头延展",
        "换主体", "换装", "擦除", "调色", "前后景融合", "图文合成",
        "音画合成", "分层叠加", "特征提取", "蒙版提取", "关键帧提取",
        "字幕提取", "风格提取", "片段拼接", "节奏压缩", "转场编排",
        "字幕对齐", "音画同步", "降噪", "补帧", "超分", "稳定化",
        "质感增强", "结构抽象", "变量抽象", "版式套用", "格式转换", "压缩导出"
    ],
}


def load_method_vocab() -> Dict[str, list]:
    """从 JSON 文件加载结构化词库（v5）"""
    if METHOD_VOCAB_PATH.exists():
        try:
            with open(METHOD_VOCAB_PATH, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"Warning: Failed to load method_vocab.json: {e}, using default")
    return DEFAULT_METHOD_VOCAB


def load_prompt_template(prompt_name: str) -> str:
    base_dir = Path(__file__).parent.parent
    prompt_path = base_dir / "prompts" / f"{prompt_name}.prompt"
    with open(prompt_path, "r", encoding="utf-8") as f:
        content = f.read()
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 3:
            content = parts[2]
    content = content.replace("$system$", "").replace("$user$", "")
    return content.strip()


def render_method_vocab_block(vocab: Dict[str, list]) -> str:
    """渲染结构化接口词库说明（v5）"""
    lines = [
        "\n# 结构化接口词库（v5，必须遵守）",
        "只输出结构化 inputs / outputs / action。",
        "- `role/流程角色` 只写接口职责，不写具体内容 what。",
        "- `modality/模态` 只写媒介或数据形态；统一用 `图片`，不要写 `图像`；统一用 `文本`，不要写 `文字`。",
        "- `artifact_type/工件类型` 写该模态下的具体工件，如 `正向提示词`、`蒙版`。",
        "- `action.main_action` 写主动作；`action.mechanism` 写动作内部机制。",
        "- 只有词库确实不够时才新增术语；新增术语也必须抽象、短、可复用。",
        "",
        "当前词库：",
    ]
    for key, values in vocab.items():
        lines.append(f"- {key}：{'、'.join(values)}")
    return "\n".join(lines)




async def extract_workflow_from_case(
    case_item: Dict[str, Any],
    llm_call: Any,
    model: str = "anthropic/claude-sonnet-4-5"
) -> tuple[Optional[Dict[str, Any]], Optional[List[Dict[str, Any]]], float]:
    """
    从单个 case item 同时提取 workflow（薄壳 steps）和 fragments（原子操作列表）。

    Returns:
        (workflow_dict, fragments_list, cost)
        workflow_dict 为 None 表示 skip 或提取失败
        fragments_list 为 None 表示 skip 或提取失败
    """
    images = case_item.get("images", [])

    case_copy = dict(case_item)
    case_copy.pop("images", None)
    case_copy.pop("_raw", None)
    case_copy.pop("workflow", None)
    case_copy.pop("fragments", None)
    case_copy.pop("capabilities", None)

    if not case_copy and not images:
        return None, None, 0.0

    title = case_item.get("title", "")[:20] or "untitled"
    context = json.dumps(case_copy, ensure_ascii=False, indent=2)

    try:
        prompt_template = load_prompt_template("extract_workflow")
        method_vocab = load_method_vocab()
        vocab_block = render_method_vocab_block(method_vocab)

        if "%context%" in prompt_template:
            prompt = prompt_template.replace("%context%", context)
        else:
            prompt = prompt_template + f"\n\n## 帖子内容\n{context}"

        if "{interface_vocab}" in prompt:
            prompt = prompt.replace("{interface_vocab}", vocab_block)
        elif vocab_block not in prompt:
            prompt = prompt + "\n" + vocab_block

    except Exception as e:
        print(f"Warning: Failed to load prompt template: {e}, using fallback")
        method_vocab = load_method_vocab()
        vocab_block = render_method_vocab_block(method_vocab)
        prompt = f"""将以下帖子内容总结为AI图片生成的工序和原子操作，以JSON格式输出。

# 输出格式（v6）
{{
  "skip": false,
  "skip_reason": "",
  "workflow": {{
    "workflow_id": null,
    "steps": [
      {{
        "step_id": "s1",
        "order": 1,
        "phase": "生成",
        "relation": "[去向.最终成品]",
        "body": "string | null"
      }}
    ]
  }},
  "fragments": [
    {{
      "fragment_id": "f_s1_0",
      "action": {{"main_action": "生成", "mechanism": "直接生成"}},
      "inputs": [{{"modality": "文本", "description": "...", "relation": "[来源.原始输入]"}}],
      "outputs": [{{"modality": "图片", "description": "...", "relation": "[去向.最终成品]"}}],
      "body": "string | null",
      "effects": [
        {{
          "statement": "实现XXX",
          "criteria": "判断标准",
          "judge_method": "vlm",
          "negative_examples": []
        }}
      ],
      "control_target": [],
      "artifact_type": null,
      "tools": [],
      "apply_to_draft": {{"实质": ["..."], "形式": ["..."]}},
      "workflow_step_ref": {{"workflow_id": null, "step_id": "s1"}},
      "is_alternative_to": []
    }}
  ]
}}

{vocab_block}

## 帖子内容
{context}

请严格按照上述格式输出JSON，不要包含其他内容。"""

    if images:
        image_urls = [img for img in images[:9] if isinstance(img, str) and img.startswith("http")]
        if image_urls:
            content_array = [{"type": "text", "text": prompt}]
            for url in image_urls:
                content_array.append({"type": "image_url", "image_url": {"url": url}})
            messages = [{"role": "user", "content": content_array}]
        else:
            messages = [{"role": "user", "content": prompt}]
    else:
        messages = [{"role": "user", "content": prompt}]

    result_data, cost = await call_llm_with_retry(
        llm_call=llm_call,
        messages=messages,
        model=model,
        temperature=0.1,
        max_tokens=10000,
        max_retries=3,
        schema_name="extract_workflow",
        task_name=f"Workflow_{title}",
    )

    if not result_data:
        return None, None, cost

    if result_data.get("skip"):
        return None, None, cost

    workflow_data = result_data.get("workflow")
    fragments_data = result_data.get("fragments", [])

    return workflow_data, fragments_data, cost


async def extract_workflow(
    case_file: Path,
    llm_call: Any,
    model: str = "anthropic/claude-sonnet-4-5",
    max_concurrent: int = 3,
    case_indices: Optional[List[int]] = None
) -> Dict[str, Any]:
    """
    按 index 遍历 case.json，提取 workflow

    Args:
        case_file: case.json 文件路径
        llm_call: LLM 调用函数
        model: 使用的模型
        max_concurrent: 最大并发数
        case_indices: 可选，指定要处理的 case index 列表。如果为 None，处理所有 case
    """
    with open(case_file, "r", encoding="utf-8") as f:
        case_data = json.load(f)

    cases = case_data.get("cases", [])

    # 如果指定了 case_indices，只处理这些 case
    if case_indices is not None:
        cases_to_process = [c for c in cases if c.get("index") in case_indices]
        print(f"Extracting workflow from {len(cases_to_process)} cases (filtered by indices: {case_indices})...")
    else:
        cases_to_process = cases
        print(f"Extracting workflow from {len(cases)} cases...")

    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_with_semaphore(case_item):
        async with semaphore:
            index = case_item.get("index", 0)
            raw = case_item.get("_raw", {})
            case_id = raw.get("case_id", "unknown")
            title = case_item.get("title", "")

            print(f"  -> [{index}] [{case_id}] extracting workflow: {title[:60]}")

            workflow, fragments, cost = await extract_workflow_from_case(case_item, llm_call, model)

            frag_count = len(fragments) if fragments else 0
            status = f"ok ({frag_count} fragments)" if workflow else "null"
            print(f"  <- [{index}] [{case_id}] workflow {status}")

            result = dict(case_item)
            result["workflow"] = workflow
            result["fragments"] = fragments if fragments is not None else []
            return result, cost

    tasks = [process_with_semaphore(case) for case in cases_to_process]
    results_with_costs = await asyncio.gather(*tasks)

    results = [r[0] for r in results_with_costs]
    costs = [r[1] for r in results_with_costs]
    total_cost = sum(costs)

    success_count = sum(1 for r in results if r.get("workflow") and r.get("fragments"))
    failed_count = len(results) - success_count

    # 如果是部分更新，需要合并回原始 cases 列表
    if case_indices is not None:
        # 创建一个 index -> result 的映射
        result_map = {r.get("index"): r for r in results}
        # 更新原始 cases 列表中对应的项
        for i, case in enumerate(cases):
            if case.get("index") in result_map:
                cases[i] = result_map[case.get("index")]
        results = cases

    results.sort(key=lambda x: x.get("index", 0))
    case_data["cases"] = results

    case_file.parent.mkdir(parents=True, exist_ok=True)
    with open(case_file, "w", encoding="utf-8") as f:
        json.dump(case_data, f, ensure_ascii=False, indent=2)

    fragments_count = sum(len(r.get("fragments") or []) for r in results)

    return {
        "total": len(results),
        "success": success_count,
        "failed": failed_count,
        "fragments_total": fragments_count,
        "total_cost": total_cost,
        "output_file": str(case_file),
    }


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python extract_workflow.py <output_dir>")
        sys.exit(1)

    output_dir = Path(sys.argv[1])
    case_file = output_dir / "case.json"

    if not case_file.exists():
        print(f"Error: {case_file} not found")
        sys.exit(1)

    print("Please use this module through run_pipeline.py")