宜博 / lemonboss

/**
 * RetrainPipeline — 带 Checkpoint + Retry 的 HDE 重训流水线
 *
 * 每步完成后写 checkpoint.json 到 conversations 目录。
 * callAgent 失败时自动重试（最多 MAX_RETRIES 次）。
 * 支持从 checkpoint 断点恢复。
 */

const fs = require('fs');
const path = require('path');

const STEPS = ['re_hunt', 're_digest', 'evolve'];
const MAX_RETRIES = 1; // 每步最多重试 1 次（共 2 次尝试）
const AGENT_IDLE_TIMEOUT = 300_000; // 5 分钟 idle timeout

class RetrainPipeline {
  /**
   * @param {object} opts
   * @param {object} opts.job - JobManager 中的 job 对象
   * @param {object} opts.skill - Skill registry 条目
   * @param {object} opts.innovator - InnovatorService 实例
   * @param {object} opts.jobManager - retrainJobManager 实例
   */
  constructor(opts) {
    this.job = opts.job;
    this.skill = opts.skill;
    this.innovator = opts.innovator;
    this.jobManager = opts.jobManager;
    this.instructions = opts.instructions || '';
    this.userId = opts.userId || null;

    this.skillName = opts.skill.skillName;
    // 系统 Agent 目录（l3-hunt 等）
    this.skillsDir = opts.innovator.skillsDir;

    // 目标 Skill 目录：优先 trained_skills/user_{id}/ → data/skills/ → skills/
    const { resolveTrainedSkillDir } = require('./skill-dir-manager');
    const trainedDir = this.userId
      ? resolveTrainedSkillDir(this.skillName, this.userId)
      : null;
    const dataSkillDir = opts.innovator.skillsDataDir
      ? path.join(opts.innovator.skillsDataDir, this.skillName)
      : null;
    const codeSkillDir = path.join(this.skillsDir, this.skillName);
    this.skillDir = this.skill.directory
      || trainedDir
      || (dataSkillDir && fs.existsSync(dataSkillDir) ? dataSkillDir : null)
      || codeSkillDir;

    // Conversations 目录：按 userId 隔离（Innovator 页面读取 user-{id}/ 下的数据）
    this.conversationsDir = path.resolve(__dirname, '../../conversations');
    const userConvDir = this.userId
      ? path.join(this.conversationsDir, `user-${this.userId}`)
      : this.conversationsDir;
    this.conversationId = `retrain-${opts.job.jobId}`;
    this.checkpointPath = path.join(userConvDir, this.conversationId, 'checkpoint.json');
    // 记录实际存储路径
    this._convStorageDir = path.join(userConvDir, this.conversationId);
  }

  // ── Checkpoint I/O ────────────────────────────

  readCheckpoint() {
    try {
      if (fs.existsSync(this.checkpointPath)) {
        return JSON.parse(fs.readFileSync(this.checkpointPath, 'utf-8'));
      }
    } catch (err) {
      console.warn(`[RetrainPipeline] Failed to read checkpoint: ${err.message}`);
    }
    return null;
  }

  writeCheckpoint(step, status, extra = {}) {
    const checkpoint = {
      jobId: this.job.jobId,
      skillName: this.skillName,
      step,
      status,
      retries: extra.retries || 0,
      updatedAt: new Date().toISOString(),
      ...extra,
    };
    try {
      const dir = path.dirname(this.checkpointPath);
      if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
      fs.writeFileSync(this.checkpointPath, JSON.stringify(checkpoint, null, 2), 'utf-8');
    } catch (err) {
      console.error(`[RetrainPipeline] Failed to write checkpoint: ${err.message}`);
    }
    return checkpoint;
  }

  // ── Pipeline 执行 ─────────────────────────────

  /**
   * 执行完整 pipeline，支持从 checkpoint 恢复
   */
  async run() {
    const { ConversationLogger } = require('../services/conversation-logger');
    const taskMgr = require('../services/skill-task-manager');
    const { randomUUID } = require('crypto');
    const { loadConfigFromSkills } = require('../src/core/skill-config-loader');
    const AgentManager = require('../src/core/agent-manager');

    const messageId = `msg-${Date.now()}-${randomUUID().substring(0, 8)}`;
    // 使用用户隔离的目录
    const logBaseDir = path.dirname(this._convStorageDir);
    const logger = new ConversationLogger(logBaseDir, this.conversationId, messageId);

    // 检查是否有未完成的 checkpoint（断点恢复）
    const checkpoint = this.readCheckpoint();
    let startFromIndex = 0;
    if (checkpoint && checkpoint.status === 'done') {
      const doneIndex = STEPS.indexOf(checkpoint.step);
      if (doneIndex >= 0 && doneIndex < STEPS.length - 1) {
        startFromIndex = doneIndex + 1;
        console.log(`[RetrainPipeline] ♻️ 从 checkpoint 恢复: 跳过 ${STEPS.slice(0, startFromIndex).join(', ')}，从 ${STEPS[startFromIndex]} 开始`);
      }
    }

    // 初始化（仅首次运行）
    if (startFromIndex === 0) {
      taskMgr.initTask(this.skillDir, {
        skillName: this.skillName,
        domain: this.skill.domain,
        scope: this.skill.scope,
        description: `重新训练 ${this.skillName}`,
        source: 'retrain',
        jobId: this.job.jobId,
        pipelineType: 'feedback',
      });
      logger.ensureMeta(`[Retrain] ${this.skillName}`);
      logger.logUserMessage(`[Retrain] 重新训练 ${this.skillName}`);
    }

    // 加载 Agent 配置 — 所有 skill（包括 agent）都在 data/skills/ 下
    const backends = JSON.parse(fs.readFileSync(
      path.join(path.resolve(__dirname, '../config'), 'backends.json'), 'utf-8'
    ));
    const agentSkillsDir = this.innovator.skillsDataDir || this.skillsDir;
    const config = loadConfigFromSkills(agentSkillsDir, backends);
    const manager = new AgentManager(config, agentSkillsDir);

    // 读取现有内容
    let existingKnowhow = '';
    let existingBenchmark = '';
    try { existingKnowhow = fs.readFileSync(path.join(this.skillDir, 'hunter', 'knowhow.md'), 'utf-8'); } catch {}
    try { existingBenchmark = fs.readFileSync(path.join(this.skillDir, 'digest', 'benchmark.md'), 'utf-8'); } catch {}

    const instructionBlock = this.instructions
      ? `\n\n## 用户训练要求\n${this.instructions}\n`
      : '';

    const stepConfigs = [
      {
        name: 're_hunt',
        label: '🔍 Step 1/3: ReHunt',
        jobStatus: 'hunting',
        jobProgress: '正在猎取领域知识 (Hunt)...',
        agent: 'innovator-hunt',
        prompt: [
          `请对 Skill "${this.skillName}" 重新猎取领域知识：`,
          `- Skill 描述: ${this.skill.description || '无'}`,
          `- 领域: ${this.skill.domain || '通用'}`,
          `- 目标: 在 ${this.skillName}/hunter/ 下更新 knowhow.md`,
          instructionBlock,
          `现有知识摘要 (前 500 字):`,
          existingKnowhow.substring(0, 500),
          ``,
          `请补充、更新、完善领域知识。`,
        ].join('\n'),
      },
      {
        name: 're_digest',
        label: '📝 Step 2/3: ReDigest',
        jobStatus: 'digesting',
        jobProgress: '正在消化提炼知识 (Digest)...',
        agent: 'innovator-digest',
        prompt: [
          `请对 Skill "${this.skillName}" 重新消化并更新评估基准：`,
          `- 更新后的知识: ${this.skillName}/hunter/knowhow.md`,
          `- 现有 benchmark: ${this.skillName}/digest/benchmark.md`,
          `- 目标: 在 ${this.skillName}/digest/ 下更新 benchmark.md`,
          `- 阈值: 95 分（满分 100）`,
          instructionBlock,
          `现有 benchmark 摘要 (前 300 字):`,
          existingBenchmark.substring(0, 300),
        ].join('\n'),
      },
      {
        name: 'evolve',
        label: '🔄 Step 3/3: Evolve',
        jobStatus: 'evolving',
        jobProgress: '进化优化中 (目标 ≥95分)...',
        agent: 'innovator-evolver',
        prompt: [
          `请对 Skill "${this.skillName}" 进行进化优化：`,
          `- skillName: ${this.skillName}`,
          `- 当前 SKILL.md: ${this.skillName}/SKILL.md`,
          `- 知识库: ${this.skillName}/hunter/knowhow.md`,
          `- 评估基准: ${this.skillName}/digest/benchmark.md`,
          `- 通过阈值: 95 分`,
          `- 最大迭代: 10 次`,
          instructionBlock,
          `请执行 evolve 循环：执行 → 评分 → 分析弱项 → 改写 SKILL.md → 重复，直到 ≥ 95 分。`,
          `结果保存到 ${this.skillName}/evolve/ 目录。`,
          `最终更新 ${this.skillName}/SKILL.md。`,
        ].join('\n'),
      },
    ];

    let evolveResult = null;

    // 逐步执行
    for (let i = startFromIndex; i < stepConfigs.length; i++) {
      const step = stepConfigs[i];
      evolveResult = await this._runStep(step, manager, logger, taskMgr);
    }

    // ── Post-Evolve: 记录进化经验到 Memory ──
    try {
      const { EvolutionMemory } = require('../src/core/evolution-memory');
      const memory = new EvolutionMemory(this.skillDir);
      const output = evolveResult?.output || '';
      const scoreMatch = output.match(/最终分数[：:]\s*(\d+)/);
      const score = scoreMatch ? parseInt(scoreMatch[1]) : null;
      if (score !== null) {
        memory.record({
          iteration: 0,
          strategy: { approach: 'retrain', steps: ['re_hunt', 're_digest', 'evolve'] },
          score,
          weakDimensions: [],
          lessonsLearned: [],
        });
      }
    } catch (err) {
      console.warn(`[RetrainPipeline] EvolutionMemory 记录失败 (非致命): ${err.message}`);
    }

    // ━━ Publish ━━
    await this._publish(evolveResult, logger, taskMgr);
  }

  /**
   * 单步执行 (带重试)
   */
  async _runStep(step, manager, logger, taskMgr) {
    console.log(`[Retrain] ${step.label} — ${this.skillName}`);
    this.jobManager.updateJob(this.job.jobId, { status: step.jobStatus, progress: step.jobProgress });
    taskMgr.updateStep(this.skillDir, step.name, 'running');
    this.writeCheckpoint(step.name, 'running');

    let lastError = null;

    for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
      try {
        if (attempt > 0) {
          console.log(`[Retrain] ⟳ 重试 ${step.name} (第 ${attempt + 1} 次尝试) — ${this.skillName}`);
          this.jobManager.updateJob(this.job.jobId, {
            progress: `${step.jobProgress} (重试 ${attempt}/${MAX_RETRIES})`,
          });
        }

        const result = await manager.callAgent(step.agent, step.prompt, {
          onData: (chunk) => { if (chunk?.trim()) logger.logStreamChunk(step.agent, chunk); },
          idleTimeout: AGENT_IDLE_TIMEOUT,
        });

        // 成功
        taskMgr.updateStep(this.skillDir, step.name, 'done');
        this.writeCheckpoint(step.name, 'done', { retries: attempt });
        console.log(`[Retrain] ✅ ${step.name} 完成: ${this.skillName}${attempt > 0 ? ` (重试 ${attempt} 次后成功)` : ''}`);
        return result;

      } catch (err) {
        lastError = err;
        console.error(`[Retrain] ❌ ${step.name} 失败 (尝试 ${attempt + 1}/${MAX_RETRIES + 1}): ${err.message}`);
        this.writeCheckpoint(step.name, 'failed', { retries: attempt, error: err.message });

        if (attempt >= MAX_RETRIES) break;
        // 等 5 秒后重试
        await new Promise(r => setTimeout(r, 5000));
      }
    }

    // 所有重试用尽
    taskMgr.updateStep(this.skillDir, step.name, 'failed');
    throw new Error(`${step.name} 失败 (已重试 ${MAX_RETRIES} 次): ${lastError?.message || 'unknown'}`);
  }

  /**
   * Publish 阶段：提取分数 + bump 版本 + 写 changelog
   */
  async _publish(evolveResult, logger, taskMgr) {
    let finalScore = null;
    const scoreMatch = (evolveResult?.output || '').match(/最终分数[：:]\s*(\d+)/);
    if (scoreMatch) finalScore = parseInt(scoreMatch[1]);

    // 先同步注册（确保 agent 对 SKILL.md 的修改被识别）
    this.innovator.autoRegisterNewSkills();

    const updatedSkill = this.innovator.registry.bumpVersion(this.skillName, 'minor', finalScore);
    this.innovator.registry.updateSkill(this.skillName, { status: 'enable' });

    // 写入 changelog
    const changelogPath = path.join(this.skillDir, 'changelog.md');
    const changelogEntry = [
      `## v${updatedSkill.version} (${new Date().toISOString().substring(0, 10)})`,
      ``,
      `**触发**: 手动重训 (Hunt → Digest → Evolve)`,
      finalScore ? `**评分**: ${finalScore}/100` : '',
      ``,
      `---`,
      ``,
    ].filter(Boolean).join('\n');

    let existingChangelog = '';
    try { existingChangelog = fs.readFileSync(changelogPath, 'utf-8'); } catch {}
    fs.writeFileSync(changelogPath, changelogEntry + existingChangelog, 'utf-8');

    logger.logResult({
      type: 'retrain',
      jobId: this.job.jobId,
      skillName: this.skillName,
      newVersion: updatedSkill.version,
      finalScore,
    });

    this.jobManager.updateJob(this.job.jobId, {
      status: 'completed',
      progress: `完成，新版本 v${updatedSkill.version}${finalScore ? ` (${finalScore}分)` : ''}`,
      conversationId: this.conversationId,
    });

    // 更新 task manager 的 publish 步骤
    taskMgr.updateStep(this.skillDir, 'publish', 'done');
    this.writeCheckpoint('publish', 'done', { version: updatedSkill.version, finalScore });
    console.log(`[Retrain] ✅ 全部完成: ${this.skillName} → v${updatedSkill.version}`);
  }

  /**
   * Master 编排模式 — 让 Master Agent 驱动 HDE 流程
   *
   * Master 收到每步的结果后自主决策下一步。
   * 产生 master 级别的日志。
   *
   * @param {MasterSession} masterSession
   * @param {object} [opts] - 额外选项
   * @param {function} [opts.onMasterData] - Master 流式输出回调
   * @param {function} [opts.onSubAgentData] - Sub-agent 流式输出回调
   */
  async runWithMaster(masterSession, opts = {}) {
    const { ConversationLogger } = require('../services/conversation-logger');
    const { randomUUID } = require('crypto');

    const messageId = `msg-${Date.now()}-${randomUUID().substring(0, 8)}`;
    const logBaseDir = path.dirname(this._convStorageDir);
    const logger = new ConversationLogger(logBaseDir, this.conversationId, messageId);

    logger.ensureMeta(`[Retrain-Master] ${this.skillName}`);
    logger.logUserMessage(`[Retrain-Master] 由 Master Agent 编排重训 ${this.skillName}`);

    // 读取 checkpoint 构建上下文
    const checkpoint = this.readCheckpoint();
    let checkpointInfo = '无（首次执行）';
    if (checkpoint) {
      checkpointInfo = `step=${checkpoint.step}, status=${checkpoint.status}, retries=${checkpoint.retries || 0}`;
    }

    // 读取现有内容摘要
    let knowhowSummary = '(无)';
    let benchmarkSummary = '(无)';
    try { knowhowSummary = fs.readFileSync(path.join(this.skillDir, 'hunter', 'knowhow.md'), 'utf-8').substring(0, 300); } catch {}
    try { benchmarkSummary = fs.readFileSync(path.join(this.skillDir, 'digest', 'benchmark.md'), 'utf-8').substring(0, 200); } catch {}

    const masterPrompt = [
      `# Retrain 任务编排`,
      ``,
      `你现在作为 **Retrain 任务监工**，负责编排 Skill "${this.skillName}" 的重训流程。`,
      ``,
      `## 任务信息`,
      `- **Skill**: ${this.skillName}`,
      `- **描述**: ${this.skill.description || '无'}`,
      `- **领域**: ${this.skill.domain || '通用'}`,
      `- **Job ID**: ${this.job.jobId}`,
      `- **Checkpoint**: ${checkpointInfo}`,
      this.instructions ? `\n## 用户训练要求\n${this.instructions}\n` : '',
      `## 流程`,
      `按顺序执行以下 3 步，每步完成后审核结果再决策下一步：`,
      ``,
      `1. **Hunt** — 调用 \`hunter\`，猎取领域知识`,
      `   - 现有知识摘要: ${knowhowSummary.substring(0, 150)}...`,
      `2. **Digest** — 调用 \`digester\`，消化提炼为评估基准`,
      `   - 现有 benchmark: ${benchmarkSummary.substring(0, 100)}...`,
      `3. **Evolve** — 调用 \`evolver\`，进化优化 SKILL.md 至 ≥95 分`,
      ``,
      `## 指令`,
      `- 如果 checkpoint 显示某步已完成 (done)，跳过该步`,
      `- 每步完成后，用 \`read_checkpoint\` 检查状态`,
      `- 如果某步失败，用 \`read_log\` 查看日志判断原因，决定重试或跳过`,
      `- 全部完成后输出总结`,
      ``,
      `## 可用工具`,
      `- \`call_sub_agent\`: 调用 sub-agent 执行任务`,
      `- \`read_checkpoint\`: 读取当前 checkpoint 状态`,
      `- \`read_log\`: 读取 sub-agent 日志尾部`,
      `- \`list_retrain_jobs\`: 列出活跃任务`,
      ``,
      `现在请开始执行。如果 checkpoint 已有完成的步骤，从下一步继续。`,
    ].join('\n');

    // 更新 job 状态
    this.jobManager.updateJob(this.job.jobId, {
      status: 'master-orchestrating',
      progress: 'Master Agent 编排中...',
    });
    this.writeCheckpoint('master-init', 'running');

    console.log(`[Retrain] 🧠 Master 编排模式启动: ${this.skillName}`);

    try {
      const result = await masterSession.handleMessage(masterPrompt, {
        onData: (chunk) => {
          if (chunk?.trim()) {
            logger.logStreamChunk('master', chunk);
            if (opts.onMasterData) opts.onMasterData(chunk);
          }
        },
        onSubAgentData: (agentName, chunk) => {
          if (chunk?.trim()) {
            logger.logStreamChunk(agentName, chunk);
            if (opts.onSubAgentData) opts.onSubAgentData(agentName, chunk);
          }
        },
        onToolCall: (tc) => {
          logger.logToolCall(tc.tool, JSON.stringify(tc.params || {}));
          console.log(`[Retrain] 🧠 Master 调用: ${tc.tool}(${JSON.stringify(tc.params).substring(0, 100)})`);

          // 同步 checkpoint — 当 Master 调用 call_sub_agent 时更新
          if (tc.tool === 'call_sub_agent') {
            const agentToStep = {
              'innovator-hunt': 're_hunt',
              'innovator-digest': 're_digest',
              'innovator-evolver': 'evolve',
            };
            const step = agentToStep[tc.params?.agent];
            if (step) {
              this.writeCheckpoint(step, 'running');
              this.jobManager.updateJob(this.job.jobId, {
                status: step === 're_hunt' ? 'hunting' : step === 're_digest' ? 'digesting' : 'evolving',
                progress: `Master 编排: ${step} 执行中...`,
              });
            }
          }
        },
        maxToolRounds: 15, // HDE 需要至少 3 轮 call_sub_agent + 审核
      });

      // Master 完成 — 执行 publish
      console.log(`[Retrain] 🧠 Master 编排完成: ${this.skillName}`);

      // 从 Master 输出或最后的 evolve 结果提取分数
      let finalScore = null;
      const allOutput = result.output || '';
      const scoreMatch = allOutput.match(/最终分数[：:]\s*(\d+)/);
      if (scoreMatch) finalScore = parseInt(scoreMatch[1]);

      // 也检查 tool call 结果中的分数
      if (!finalScore && result.toolCalls) {
        for (const tc of result.toolCalls) {
          if (tc.tool === 'call_sub_agent' && tc.params?.agent === 'innovator-evolver') {
            const evolveOutput = tc.result?.output || '';
            const sm = evolveOutput.match(/最终分数[：:]\s*(\d+)/);
            if (sm) finalScore = parseInt(sm[1]);
          }
        }
      }

      // 先同步注册（确保 agent 对 SKILL.md 的修改被识别）
      this.innovator.autoRegisterNewSkills();

      // Publish — bumpVersion + enable 作为最终权威状态
      const updatedSkill = this.innovator.registry.bumpVersion(this.skillName, 'minor', finalScore);
      this.innovator.registry.updateSkill(this.skillName, { status: 'enable' });

      // 写入 changelog
      const changelogPath = path.join(this.skillDir, 'changelog.md');
      const changelogEntry = [
        `## v${updatedSkill.version} (${new Date().toISOString().substring(0, 10)})`,
        ``,
        `**触发**: 手动重训 (Master 编排 Hunt → Digest → Evolve)`,
        finalScore ? `**评分**: ${finalScore}/100` : '',
        ``,
        `---`,
        ``,
      ].filter(Boolean).join('\n');

      let existingChangelog = '';
      try { existingChangelog = fs.readFileSync(changelogPath, 'utf-8'); } catch {}
      fs.writeFileSync(changelogPath, changelogEntry + existingChangelog, 'utf-8');

      logger.logResult({
        type: 'retrain-master',
        jobId: this.job.jobId,
        skillName: this.skillName,
        newVersion: updatedSkill.version,
        finalScore,
        rounds: result.rounds,
        toolCalls: (result.toolCalls || []).length,
      });

      this.jobManager.updateJob(this.job.jobId, {
        status: 'completed',
        progress: `完成 (Master 编排)，新版本 v${updatedSkill.version}${finalScore ? ` (${finalScore}分)` : ''}`,
        conversationId: this.conversationId,
      });

      this.writeCheckpoint('publish', 'done', { version: updatedSkill.version, finalScore, orchestrator: 'master' });
      console.log(`[Retrain] 🧠✅ Master 编排全部完成: ${this.skillName} → v${updatedSkill.version}`);

      return result;

    } catch (err) {
      console.error(`[Retrain] 🧠❌ Master 编排失败: ${this.skillName}`, err.message);
      this.writeCheckpoint('master-error', 'failed', { error: err.message });

      // 降级: 尝试用原有 Node.js 模式恢复
      console.log(`[Retrain] ⬇️ 降级为 Node.js 串行模式: ${this.skillName}`);
      return this.run();
    }
  }
}

// 导出常量供测试和 watchdog 使用
RetrainPipeline.STEPS = STEPS;
RetrainPipeline.MAX_RETRIES = MAX_RETRIES;
RetrainPipeline.AGENT_IDLE_TIMEOUT = AGENT_IDLE_TIMEOUT;

module.exports = { RetrainPipeline };