how/step3_generate_inspirations.py

"""
Step3: 基于匹配节点生成灵感点

基于 Step1 的 Top1 匹配结果，以匹配到的人设要素作为锚点，
让 Agent 分析可以产生哪些灵感点
"""
import os
import sys
import json
import asyncio
from pathlib import Path

from agents import Agent, Runner, trace
from agents.tracing.create import custom_span
from lib.my_trace import set_trace_smith as set_trace
from lib.client import get_model
from lib.data_loader import load_persona_data, load_inspiration_list, select_inspiration

# 模型配置
MODEL_NAME = "google/gemini-2.5-pro"


# ========== System Prompt ==========
GENERATE_INSPIRATIONS_PROMPT = """
# 任务
基于给定的人设体系和锚点要素，分析和生成可能的灵感点。

## 输入说明

- **<人设体系></人设体系>**: 完整的人设系统，用于理解整体风格和定位
- **<锚点要素></锚点要素>**: 作为锚点的人设要素（一级或二级分类）
- **<要素定义></要素定义>**: 该要素在人设中的完整定义（如果有）
- **<要素上下文></要素上下文>**: 该要素的上下文信息（所属视角、一级分类等）

## 分析方法

### 核心原则：基于人设特征和要素定位发散灵感

1. 先理解整个人设体系的风格、调性和定位
2. 深入理解锚点要素在人设中的含义和作用
3. 基于人设特征，推理可能触发该要素的灵感来源

### 分析步骤

1. **理解人设体系**
   - 分析人设的整体风格和内容定位
   - 理解灵感点、目的点、关键点的组织方式
   - 把握人设的表达特色和价值取向

2. **理解锚点要素**
   - 结合要素定义，理解该要素的核心含义
   - 分析该要素在人设体系中的层级位置
   - 理解该要素代表的内容类型或表达方式

3. **发散灵感思考**
   - 基于人设整体风格，思考什么样的灵感会触发该要素
   - 考虑不同的场景、话题、情感、事件等
   - 确保生成的灵感符合人设的调性和定位
   - 灵感应该多样化，但都要能映射到该要素

4. **生成灵感点列表**
   - 每个灵感点应该简洁明确（一句话）
   - 灵感点之间应有一定的多样性
   - 灵感点应该能够触发该人设要素
   - 纯粹基于人设和要素特征推理

---

## 输出格式（严格JSON）

```json
{
  "人设理解": {
    "整体风格": "简要描述人设的整体风格和定位（1-2句话）",
    "内容调性": "该人设的内容调性和表达特点"
  },
  "要素分析": {
    "核心特征": "结合人设和要素定义，描述该要素的核心特征（1-2句话）",
    "适用场景": "该要素在这个人设中适用的内容场景"
  },
  "灵感点列表": [
    {
      "灵感点": "具体的灵感点描述",
      "说明": "为什么这个灵感可能触发该要素（结合人设特征说明）"
    },
    {
      "灵感点": "具体的灵感点描述",
      "说明": "为什么这个灵感可能触发该要素（结合人设特征说明）"
    }
  ]
}
```

**输出要求**：
1. 必须严格按照上述JSON格式输出
2. 所有字段都必须填写
3. **人设理解**：体现对整个人设体系的把握
4. **要素分析**：结合人设和要素定义进行分析
5. **灵感点列表**：生成 5-10 个灵感点
6. 每个灵感点包含：
   - **灵感点**：简洁的灵感描述（一句话）
   - **说明**：解释为什么这个灵感可能触发该要素，要结合人设特征（1-2句话）
7. 灵感点应该多样化，覆盖不同角度和场景
8. 纯粹基于人设和要素特征进行推理，不依赖外部参考
""".strip()


def create_generate_agent(model_name: str) -> Agent:
    """创建灵感生成的 Agent

    Args:
        model_name: 模型名称

    Returns:
        Agent 实例
    """
    agent = Agent(
        name="Inspiration Generator Expert",
        instructions=GENERATE_INSPIRATIONS_PROMPT,
        model=get_model(model_name),
        tools=[],
    )

    return agent


def parse_generate_response(response_content: str) -> dict:
    """解析生成响应

    Args:
        response_content: Agent 返回的响应内容

    Returns:
        解析后的字典
    """
    try:
        # 如果响应包含在 markdown 代码块中，提取 JSON 部分
        if "```json" in response_content:
            json_start = response_content.index("```json") + 7
            json_end = response_content.index("```", json_start)
            json_text = response_content[json_start:json_end].strip()
        elif "```" in response_content:
            json_start = response_content.index("```") + 3
            json_end = response_content.index("```", json_start)
            json_text = response_content[json_start:json_end].strip()
        else:
            json_text = response_content.strip()

        return json.loads(json_text)
    except Exception as e:
        print(f"解析响应失败: {e}")
        return {
            "人设理解": {
                "整体风格": "解析失败",
                "内容调性": "解析失败"
            },
            "要素分析": {
                "核心特征": "解析失败",
                "适用场景": "解析失败"
            },
            "灵感点列表": []
        }


def format_persona_system(persona_data: dict) -> str:
    """格式化完整人设系统为文本

    Args:
        persona_data: 人设数据

    Returns:
        格式化的人设系统文本
    """
    lines = ["# 人设系统"]

    # 处理三个部分：灵感点列表、目的点、关键点列表
    for section_key, section_title in [
        ("灵感点列表", "【灵感点】灵感的来源和性质"),
        ("目的点", "【目的点】创作的目的和价值导向"),
        ("关键点列表", "【关键点】内容的核心主体和表达方式")
    ]:
        section_data = persona_data.get(section_key, [])
        if not section_data:
            continue

        lines.append(f"\n## {section_title}\n")

        for perspective in section_data:
            perspective_name = perspective.get("视角名称", "")
            lines.append(f"\n### 视角：{perspective_name}")

            for pattern in perspective.get("模式列表", []):
                pattern_name = pattern.get("分类名称", "")
                pattern_def = pattern.get("核心定义", "")
                lines.append(f"\n  【一级】{pattern_name}")
                if pattern_def:
                    lines.append(f"    定义：{pattern_def}")

                # 二级细分
                for sub in pattern.get("二级细分", []):
                    sub_name = sub.get("分类名称", "")
                    sub_def = sub.get("分类定义", "")
                    lines.append(f"    【二级】{sub_name}：{sub_def}")

    return "\n".join(lines)


def find_element_definition(persona_data: dict, element_name: str) -> str:
    """从人设数据中查找要素的定义

    Args:
        persona_data: 人设数据
        element_name: 要素名称

    Returns:
        要素定义文本，如果未找到则返回空字符串
    """
    # 在灵感点列表中查找
    for section_key in ["灵感点列表", "目的点", "关键点列表"]:
        section_data = persona_data.get(section_key, [])
        for perspective in section_data:
            for pattern in perspective.get("模式列表", []):
                # 检查一级分类
                if pattern.get("分类名称", "") == element_name:
                    definition = pattern.get("核心定义", "")
                    if definition:
                        return definition

                # 检查二级分类
                for sub in pattern.get("二级细分", []):
                    if sub.get("分类名称", "") == element_name:
                        return sub.get("分类定义", "")

    return ""


def find_step1_file(persona_dir: str, inspiration: str, model_name: str) -> str:
    """查找 step1 输出文件

    Args:
        persona_dir: 人设目录
        inspiration: 灵感点名称
        model_name: 模型名称

    Returns:
        step1 文件路径

    Raises:
        SystemExit: 找不到文件时退出
    """
    step1_dir = os.path.join(persona_dir, "how", "灵感点", inspiration)
    model_name_short = model_name.replace("google/", "").replace("/", "_")
    step1_file_pattern = f"*_step1_*_{model_name_short}.json"

    step1_files = list(Path(step1_dir).glob(step1_file_pattern))
    if not step1_files:
        print(f"❌ 找不到 step1 输出文件")
        print(f"查找路径: {step1_dir}/{step1_file_pattern}")
        sys.exit(1)

    return str(step1_files[0])


async def process_step3_generate_inspirations(
    step1_top1: dict,
    persona_data: dict,
    current_time: str = None,
    log_url: str = None
) -> dict:
    """执行灵感生成分析（核心业务逻辑）

    Args:
        step1_top1: step1 的 top1 匹配结果
        persona_data: 完整的人设数据
        current_time: 当前时间戳
        log_url: trace URL

    Returns:
        生成结果字典
    """
    # 从 step1 结果中提取信息
    business_info = step1_top1.get("业务信息", {})
    input_info = step1_top1.get("输入信息", {})

    matched_element = business_info.get("匹配要素", "")
    element_context = input_info.get("A_Context", "")

    # 格式化人设系统
    persona_system_text = format_persona_system(persona_data)

    # 查找要素定义
    element_definition = find_element_definition(persona_data, matched_element)

    print(f"\n开始灵感生成分析")
    print(f"锚点要素: {matched_element}")
    print(f"要素定义: {element_definition if element_definition else '(未找到定义)'}")
    print(f"模型: {MODEL_NAME}\n")

    # 构建任务描述（包含完整人设系统、锚点要素、要素定义、要素上下文）
    task_description = f"""## 本次分析任务

<人设体系>
{persona_system_text}
</人设体系>

<锚点要素>
{matched_element}
</锚点要素>

<要素定义>
{element_definition if element_definition else '无'}
</要素定义>

<要素上下文>
{element_context}
</要素上下文>

请基于上述完整的人设体系和锚点要素，深入理解人设的整体风格和该要素的定位，推理并生成可能的灵感点列表，严格按照系统提示中的 JSON 格式输出结果。"""

    # 构造消息
    messages = [{
        "role": "user",
        "content": [
            {
                "type": "input_text",
                "text": task_description
            }
        ]
    }]

    # 使用 custom_span 追踪生成过程
    with custom_span(
        name=f"Step3: 灵感生成 - {matched_element}",
        data={
            "锚点要素": matched_element,
            "模型": MODEL_NAME,
            "步骤": "基于要素生成灵感点"
        }
    ):
        # 创建 Agent
        agent = create_generate_agent(MODEL_NAME)

        # 运行 Agent
        result = await Runner.run(agent, input=messages)

    # 解析响应
    parsed_result = parse_generate_response(result.final_output)

    # 构建输出
    return {
        "元数据": {
            "current_time": current_time,
            "log_url": log_url,
            "model": MODEL_NAME,
            "步骤": "Step3: 基于匹配节点生成灵感点"
        },
        "锚点信息": {
            "人设要素": matched_element,
            "要素定义": element_definition if element_definition else "无",
            "要素上下文": element_context
        },
        "step1_结果": step1_top1,
        "生成结果": parsed_result
    }


async def main(current_time: str, log_url: str, force: bool = False):
    """主函数

    Args:
        current_time: 当前时间戳
        log_url: 日志链接
        force: 是否强制重新执行（跳过已存在文件检查）
    """
    # 解析命令行参数
    persona_dir = sys.argv[1] if len(sys.argv) > 1 else "data/阿里多多酱/out/人设_1110"
    inspiration_arg = sys.argv[2] if len(sys.argv) > 2 else "0"

    # 第三个参数：force（如果从命令行调用且有该参数，则覆盖函数参数）
    if len(sys.argv) > 3 and sys.argv[3] == "force":
        force = True

    print(f"{'=' * 80}")
    print(f"Step3: 基于匹配节点生成灵感点")
    print(f"{'=' * 80}")
    print(f"人设目录: {persona_dir}")
    print(f"灵感参数: {inspiration_arg}")

    # 加载数据
    persona_data = load_persona_data(persona_dir)
    inspiration_list = load_inspiration_list(persona_dir)

    # 选择灵感
    try:
        inspiration_index = int(inspiration_arg)
        if 0 <= inspiration_index < len(inspiration_list):
            test_inspiration = inspiration_list[inspiration_index]
            print(f"使用灵感[{inspiration_index}]: {test_inspiration}")
        else:
            print(f"❌ 灵感索引超出范围: {inspiration_index}")
            sys.exit(1)
    except ValueError:
        if inspiration_arg in inspiration_list:
            test_inspiration = inspiration_arg
            print(f"使用灵感: {test_inspiration}")
        else:
            print(f"❌ 找不到灵感: {inspiration_arg}")
            sys.exit(1)

    # 查找并加载 step1 结果
    step1_file = find_step1_file(persona_dir, test_inspiration, MODEL_NAME)
    step1_filename = os.path.basename(step1_file)
    step1_basename = os.path.splitext(step1_filename)[0]

    print(f"Step1 输入文件: {step1_file}")

    # 构建输出文件路径
    output_dir = os.path.join(persona_dir, "how", "灵感点", test_inspiration)
    model_name_short = MODEL_NAME.replace("google/", "").replace("/", "_")
    scope_prefix = step1_basename.split("_")[0]
    result_index = 0
    output_filename = f"{scope_prefix}_step3_top{result_index + 1}_生成灵感_{model_name_short}.json"
    output_file = os.path.join(output_dir, output_filename)

    # 检查文件是否已存在
    if not force and os.path.exists(output_file):
        print(f"\n✓ 输出文件已存在，跳过执行: {output_file}")
        print(f"提示: 如需重新执行，请添加 'force' 参数\n")
        return

    with open(step1_file, 'r', encoding='utf-8') as f:
        step1_data = json.load(f)

    actual_inspiration = step1_data.get("灵感", "")
    step1_results = step1_data.get("匹配结果列表", [])

    if not step1_results:
        print("❌ step1 结果为空")
        sys.exit(1)

    print(f"灵感: {actual_inspiration}")

    # 默认处理 top1
    selected_result = step1_results[result_index]
    print(f"处理第 {result_index + 1} 个匹配结果（Top{result_index + 1}）\n")

    # 执行核心业务逻辑
    output = await process_step3_generate_inspirations(
        step1_top1=selected_result,
        persona_data=persona_data,
        current_time=current_time,
        log_url=log_url
    )

    # 在元数据中添加 step1 匹配索引
    output["元数据"]["step1_匹配索引"] = result_index + 1

    # 保存结果
    os.makedirs(output_dir, exist_ok=True)

    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(output, f, ensure_ascii=False, indent=2)

    # 输出生成的灵感点预览
    generated = output.get("生成结果", {})
    inspirations = generated.get("灵感点列表", [])

    print(f"\n{'=' * 80}")
    print(f"生成了 {len(inspirations)} 个灵感点:")
    print(f"{'=' * 80}")
    for i, item in enumerate(inspirations[:5], 1):
        print(f"{i}. {item.get('灵感点', '')}")
    if len(inspirations) > 5:
        print(f"... 还有 {len(inspirations) - 5} 个")

    print(f"\n完成！结果已保存到: {output_file}")
    if log_url:
        print(f"Trace: {log_url}\n")


if __name__ == "__main__":
    # 设置 trace
    current_time, log_url = set_trace()

    # 使用 trace 上下文包裹整个执行流程
    with trace("Step3: 生成灵感点"):
        asyncio.run(main(current_time, log_url))