memind/scripts/benchmark-token-optimization.mjs

#!/usr/bin/env node
/**
 * Local A/B benchmark: DeepSeek token usage vs scenario quality.
 *
 * Usage:
 *   node scripts/benchmark-token-optimization.mjs
 *   node scripts/benchmark-token-optimization.mjs --scenario john2-suzhou-page
 *   node scripts/benchmark-token-optimization.mjs --phase baseline
 *   node scripts/benchmark-token-optimization.mjs --phase optimized
 *   node scripts/benchmark-token-optimization.mjs --restore
 *
 * Does NOT touch production. Restores model + memory admin config on exit.
 */
import fs from 'node:fs';
import path from 'node:path';
import { fileURLToPath } from 'node:url';
import { createDbPool } from '../db.mjs';
import { createLlmProviderService } from '../llm-providers.mjs';
import { createMemoryV2AdminConfigService } from '../memory-v2-admin-config.mjs';
import { loadH5Environment } from './load-env.mjs';
import {
  createReporter,
  createAgentRun,
  extractAssistantTexts,
  extractPublicLinks,
  getSession,
  loadScenario,
  loginViaApi,
  resolvePortalBase,
  snapshotPublicHtml,
  verifyPageAccess,
  waitForAssistantGrowth,
  waitForRunTerminal,
} from './scenario-test-lib.mjs';

const scriptDir = path.dirname(fileURLToPath(import.meta.url));
const root = path.join(scriptDir, '..');
const statePath = path.join(root, '.token-benchmark-state.json');

loadH5Environment(scriptDir);

const PROFILES = {
  baseline: {
    label: 'baseline (deepseek-v4-pro + 默认记忆预算)',
    model: 'deepseek-v4-pro',
    memoryPatch: null,
  },
  optimized: {
    label: 'optimized (deepseek-v4-flash + 收紧记忆/Router)',
    model: 'deepseek-v4-flash',
    memoryPatch: {
      runtimeControl: {
        agentResolveLimit: '3',
      },
      retriever: {
        limit: '4',
        tokenBudget: '900',
      },
      persona: {
        maxTokens: '300',
      },
      chatIntentRouter: {
        enabled: true,
        shadowMode: false,
        model: 'deepseek-chat',
        memoryResolveLimit: '2',
      },
    },
  },
};

function parseArgs(argv) {
  let scenarioId = 'token-benchmark-page';
  let phase = 'both';
  let port = Number(process.env.H5_PORT ?? 8081);
  for (let i = 2; i < argv.length; i += 1) {
    const arg = argv[i];
    if (arg === '--scenario' && argv[i + 1]) scenarioId = argv[++i];
    else if (arg === '--phase' && argv[i + 1]) phase = argv[++i];
    else if (arg === '--port' && argv[i + 1]) port = Number(argv[++i]);
    else if (arg === '--restore') phase = 'restore';
    else if (arg === '-h' || arg === '--help') {
      console.log(`Usage:
  node scripts/benchmark-token-optimization.mjs [--scenario <id>] [--phase baseline|optimized|both|restore]`);
      process.exit(0);
    } else {
      throw new Error(`Unknown argument: ${arg}`);
    }
  }
  if (!['baseline', 'optimized', 'both', 'restore'].includes(phase)) {
    throw new Error(`Invalid --phase: ${phase}`);
  }
  return { scenarioId, phase, port };
}

async function ensurePortalReady(baseUrl) {
  const response = await fetch(`${baseUrl}/auth/status`);
  if (!response.ok) {
    throw new Error(`Portal 未就绪: ${baseUrl}/auth/status -> ${response.status}`);
  }
}

async function readCurrentModel(pool) {
  const [rows] = await pool.query(
    `SELECT default_model FROM h5_llm_provider_keys WHERE is_selected = 1 AND status = 'active' LIMIT 1`,
  );
  return rows[0]?.default_model ?? null;
}

async function saveState(pool, memoryConfigService) {
  const model = await readCurrentModel(pool);
  const memoryAdmin = await memoryConfigService.getAdminConfig().catch(() => null);
  const payload = {
    savedAt: Date.now(),
    model,
    memoryConfig: memoryAdmin?.config ?? null,
  };
  fs.writeFileSync(statePath, `${JSON.stringify(payload, null, 2)}\n`, 'utf8');
  return payload;
}

function loadSavedState() {
  if (!fs.existsSync(statePath)) return null;
  return JSON.parse(fs.readFileSync(statePath, 'utf8'));
}

async function applyProfile(pool, llmProviderService, memoryConfigService, profileKey) {
  const profile = PROFILES[profileKey];
  if (!profile) throw new Error(`Unknown profile: ${profileKey}`);

  const modelResult = await llmProviderService.setGlobalModel(profile.model);
  if (!modelResult?.ok) {
    throw new Error(`切换模型失败 (${profile.model}): ${modelResult?.message ?? 'unknown'}`);
  }

  if (profile.memoryPatch) {
    await memoryConfigService.updateAdminConfig(profile.memoryPatch, {
      updatedBy: 'token-benchmark',
    });
  }

  const currentModel = await readCurrentModel(pool);
  console.log(`\n==> 已应用 ${profile.label}`);
  console.log(`    全局模型: ${currentModel}`);
  if (profile.memoryPatch) {
    console.log('    记忆/Router 预算: 已写入 h5_memory_v2_admin_config');
  }
}

async function restoreSavedState(pool, llmProviderService, memoryConfigService) {
  const saved = loadSavedState();
  if (!saved) {
    console.log('无 .token-benchmark-state.json，跳过恢复');
    return false;
  }

  if (saved.model) {
    const result = await llmProviderService.setGlobalModel(saved.model);
    if (!result?.ok) {
      console.warn(`恢复模型失败: ${result?.message ?? 'unknown'}`);
    } else {
      console.log(`已恢复全局模型: ${saved.model}`);
    }
  }

  if (saved.memoryConfig) {
    await memoryConfigService.updateAdminConfig(saved.memoryConfig, {
      updatedBy: 'token-benchmark-restore',
    });
    console.log('已恢复 memory v2 admin config');
  }

  fs.unlinkSync(statePath);
  return true;
}

async function collectUsageMetrics(pool, requestIds) {
  if (!requestIds.length) {
    return { inputTokens: 0, outputTokens: 0, costCents: 0, records: [] };
  }
  const placeholders = requestIds.map(() => '?').join(', ');
  const [rows] = await pool.query(
    `SELECT request_id, input_tokens, output_tokens, cost_cents, billing_source, created_at
     FROM h5_usage_records
     WHERE request_id IN (${placeholders})
     ORDER BY created_at ASC`,
    requestIds,
  );
  let inputTokens = 0;
  let outputTokens = 0;
  let costCents = 0;
  for (const row of rows) {
    inputTokens += Number(row.input_tokens ?? 0);
    outputTokens += Number(row.output_tokens ?? 0);
    costCents += Number(row.cost_cents ?? 0);
  }
  return { inputTokens, outputTokens, costCents, records: rows };
}

async function collectRunEventTokens(pool, runIds) {
  if (!runIds.length) return { inputTokens: 0, outputTokens: 0, events: [] };
  const placeholders = runIds.map(() => '?').join(', ');
  const [rows] = await pool.query(
    `SELECT run_id, data_json, created_at
     FROM h5_agent_run_events
     WHERE run_id IN (${placeholders}) AND event_type = 'session_finished'
     ORDER BY created_at ASC`,
    runIds,
  );
  let inputTokens = 0;
  let outputTokens = 0;
  const events = [];
  for (const row of rows) {
    const data = typeof row.data_json === 'string'
      ? JSON.parse(row.data_json)
      : row.data_json;
    const tokenState = data?.tokenState ?? data?.token_state ?? null;
    if (!tokenState) continue;
    const input = Number(
      tokenState.accumulatedInputTokens
      ?? tokenState.accumulated_input_tokens
      ?? tokenState.inputTokens
      ?? tokenState.input_tokens
      ?? 0,
    );
    const output = Number(
      tokenState.accumulatedOutputTokens
      ?? tokenState.accumulated_output_tokens
      ?? tokenState.outputTokens
      ?? tokenState.output_tokens
      ?? 0,
    );
    inputTokens += input;
    outputTokens += output;
    events.push({ runId: row.run_id, input, output, tokenState });
  }
  return { inputTokens, outputTokens, events };
}

async function runScenarioBenchmark(scenario, port, { profileKey = 'run', runId = Date.now() } = {}) {
  const reporter = createReporter();
  const baseUrl = resolvePortalBase(port);
  const account = {
    username: scenario.account?.username ?? 'john2',
    password: process.env.JOHN_PASSWORD
      ?? process.env.H5_ACCESS_PASSWORD
      ?? scenario.account?.password
      ?? '888888',
  };

  await ensurePortalReady(baseUrl);

  const auth = await loginViaApi(baseUrl, account, reporter);
  const publishKey = auth.user?.id ?? auth.user?.publishSlug ?? account.username;

  let sessionId = null;
  let assistantCount = 0;
  let assistantCombinedLength = 0;
  const tracked = {
    requestIds: [],
    runIds: [],
    sessionIds: [],
    replies: [],
    pageLinks: [],
    elapsedMs: 0,
    profileKey,
    runId,
  };
  const started = Date.now();

  for (const step of scenario.steps ?? []) {
    if (step.action !== 'chat') continue;

    const message = String(step.message ?? '')
      .replaceAll('{{PROFILE}}', profileKey)
      .replaceAll('{{RUN_ID}}', String(runId));

    let htmlBefore = [];
    if (step.expect?.page) {
      htmlBefore = await snapshotPublicHtml(publishKey);
    }

    const run = await createAgentRun(baseUrl, auth.cookie, {
      message,
      sessionId: null,
      selectedChatSkill: step.selectedChatSkill ?? null,
    });
    tracked.requestIds.push(run.requestId);
    tracked.runIds.push(run.runId);

    const terminal = await waitForRunTerminal(
      baseUrl,
      auth.cookie,
      run.runId,
      step.expect?.timeoutMs ?? 600_000,
    );
    sessionId = terminal.sessionId ?? terminal.agent_session_id ?? run.sessionId ?? sessionId;
    if (sessionId) tracked.sessionIds.push(sessionId);

    if (terminal.status === 'failed') {
      reporter.fail('run 终态', terminal.error ?? 'failed');
      continue;
    }
    reporter.pass('run 终态', terminal.status);

    const reply = await waitForAssistantGrowth(baseUrl, auth.cookie, sessionId, {
      previousCount: assistantCount,
      previousCombinedLength: assistantCombinedLength,
      minChars: step.expect?.assistantMinChars ?? 1,
      timeoutMs: step.expect?.timeoutMs ?? 120_000,
    });

    if (!reply) {
      reporter.fail('assistant 回复', '超时未收到新回复');
      continue;
    }

    assistantCount = reply.count;
    assistantCombinedLength = reply.combined.length;
    tracked.replies.push({
      label: step.label ?? step.message,
      chars: reply.combined.length,
      elapsedMs: reply.elapsedMs,
      text: reply.combined,
    });
    reporter.pass('assistant 回复', `${reply.combined.length} 字 / ${reply.elapsedMs}ms`);

    const keywords = step.expect?.replyKeywords ?? [];
    if (keywords.length) {
      const hit = keywords.filter((word) => reply.combined.includes(word));
      if (hit.length === 0) reporter.fail('回复关键词', `未命中: ${keywords.join(', ')}`);
      else reporter.pass('回复关键词', hit.join(', '));
    }

    tracked.pageLinks.push(...extractPublicLinks(reply.combined, baseUrl));

    if (step.expect?.page) {
      await verifyPageAccess({
        baseUrl,
        cookie: auth.cookie,
        publishKey,
        replyText: reply.combined,
        htmlBefore,
        expect: step.expect.page,
        uploadedAssetIds: [],
        reporter,
      });
    }
  }

  tracked.elapsedMs = Date.now() - started;
  return {
    reporter,
    tracked,
    passed: reporter.issues.length === 0,
    issueCount: reporter.issues.length,
  };
}

function pctDelta(before, after) {
  if (!before) return after ? 100 : 0;
  return Math.round(((after - before) / before) * 1000) / 10;
}

function printComparison(baseline, optimized) {
  console.log('\n========================================');
  console.log('Token 优化 A/B 对比（本地）');
  console.log('========================================\n');

  const rows = [
    ['指标', 'Baseline', 'Optimized', '变化'],
    [
      '输入 Token (usage)',
      String(baseline.usage.inputTokens),
      String(optimized.usage.inputTokens),
      `${pctDelta(baseline.usage.inputTokens, optimized.usage.inputTokens)}%`,
    ],
    [
      '输出 Token (usage)',
      String(baseline.usage.outputTokens),
      String(optimized.usage.outputTokens),
      `${pctDelta(baseline.usage.outputTokens, optimized.usage.outputTokens)}%`,
    ],
    [
      '总 Token (usage)',
      String(baseline.usage.inputTokens + baseline.usage.outputTokens),
      String(optimized.usage.inputTokens + optimized.usage.outputTokens),
      `${pctDelta(
        baseline.usage.inputTokens + baseline.usage.outputTokens,
        optimized.usage.inputTokens + optimized.usage.outputTokens,
      )}%`,
    ],
    [
      '扣费 (分)',
      String(baseline.usage.costCents),
      String(optimized.usage.costCents),
      `${pctDelta(baseline.usage.costCents, optimized.usage.costCents)}%`,
    ],
    [
      '场景耗时 (s)',
      String(Math.round(baseline.tracked.elapsedMs / 1000)),
      String(Math.round(optimized.tracked.elapsedMs / 1000)),
      `${pctDelta(baseline.tracked.elapsedMs, optimized.tracked.elapsedMs)}%`,
    ],
    [
      '回复字数',
      String(baseline.tracked.replies.at(-1)?.chars ?? 0),
      String(optimized.tracked.replies.at(-1)?.chars ?? 0),
      `${pctDelta(
        baseline.tracked.replies.at(-1)?.chars ?? 0,
        optimized.tracked.replies.at(-1)?.chars ?? 0,
      )}%`,
    ],
    [
      '质量检查',
      baseline.passed ? '通过' : `失败 ${baseline.issueCount} 项`,
      optimized.passed ? '通过' : `失败 ${optimized.issueCount} 项`,
      optimized.passed === baseline.passed ? '持平' : (optimized.passed ? '改善' : '下降'),
    ],
  ];

  const widths = [0, 1, 2, 3].map((col) => Math.max(...rows.map((row) => row[col].length)));
  for (const row of rows) {
    console.log(row.map((cell, i) => cell.padEnd(widths[i] + 2)).join(''));
  }

  console.log('\n--- Baseline 页面链接 ---');
  for (const link of baseline.tracked.pageLinks) console.log(link);
  console.log('\n--- Optimized 页面链接 ---');
  for (const link of optimized.tracked.pageLinks) console.log(link);

  if (baseline.passed && optimized.passed) {
    const totalBefore = baseline.usage.inputTokens + baseline.usage.outputTokens;
    const totalAfter = optimized.usage.inputTokens + optimized.usage.outputTokens;
    const saved = totalBefore - totalAfter;
    console.log(`\n结论: 质量检查均通过。Token 总量 ${totalBefore} → ${totalAfter}（${saved >= 0 ? '节省' : '增加'} ${Math.abs(saved)}）`);
    if (saved > 0 && optimized.passed) {
      console.log('效果未打折，可继续观察更多场景。');
    } else if (saved <= 0) {
      console.log('Token 未下降，需检查模型/Router/记忆配置是否生效。');
    }
  } else {
    console.log('\n结论: 存在质量回归，不建议直接上生产。');
    if (!optimized.passed) {
      console.log('Optimized 失败项:');
      for (const issue of optimized.reporter.issues) {
        console.log(`  - ${issue.label}: ${issue.detail}`);
      }
    }
  }
}

async function runPhase(pool, llmProviderService, memoryConfigService, profileKey, scenario, port) {
  await applyProfile(pool, llmProviderService, memoryConfigService, profileKey);
  const runId = `${profileKey}-${Date.now()}`;
  console.log(`\n>>> 开始 ${profileKey} 场景跑分 (runId=${runId})...`);
  const result = await runScenarioBenchmark(scenario, port, { profileKey, runId });
  const usage = await collectUsageMetrics(pool, result.tracked.requestIds);
  const eventTokens = await collectRunEventTokens(pool, result.tracked.runIds);
  console.log(`\n--- ${profileKey} token 汇总 ---`);
  console.log(`usage records: in=${usage.inputTokens} out=${usage.outputTokens} cost=${usage.costCents}分`);
  console.log(`run events:    in=${eventTokens.inputTokens} out=${eventTokens.outputTokens}`);
  return { ...result, usage, eventTokens, profileKey };
}

async function main() {
  const { scenarioId, phase, port } = parseArgs(process.argv);
  const pool = createDbPool(process.env);
  const llmProviderService = createLlmProviderService(pool, {
    apiTarget: process.env.TKMIND_API_TARGET ?? 'https://127.0.0.1:18006',
    apiSecret: process.env.TKMIND_SERVER__SECRET_KEY ?? 'local-dev-secret',
  });
  const memoryConfigService = createMemoryV2AdminConfigService(pool);

  try {
    if (phase === 'restore') {
      await restoreSavedState(pool, llmProviderService, memoryConfigService);
      return;
    }

    const scenario = await loadScenario(scenarioId);
    console.log(`场景: ${scenario.name ?? scenario.id}`);
    console.log(`Portal: ${resolvePortalBase(port)}`);

    if (phase === 'both') {
      await saveState(pool, memoryConfigService);
    }

    let baselineResult = null;
    let optimizedResult = null;

    if (phase === 'both' || phase === 'baseline') {
      baselineResult = await runPhase(
        pool,
        llmProviderService,
        memoryConfigService,
        'baseline',
        scenario,
        port,
      );
    }

    if (phase === 'both' || phase === 'optimized') {
      optimizedResult = await runPhase(
        pool,
        llmProviderService,
        memoryConfigService,
        'optimized',
        scenario,
        port,
      );
    }

    if (baselineResult && optimizedResult) {
      printComparison(baselineResult, optimizedResult);
    }

    if (phase === 'both') {
      await restoreSavedState(pool, llmProviderService, memoryConfigService);
      console.log('\n已恢复 benchmark 前的模型与 memory 配置。');
    }
  } finally {
    await pool.end?.();
  }
}

main().catch(async (error) => {
  console.error(error instanceof Error ? error.stack ?? error.message : error);
  try {
    const pool = createDbPool(process.env);
    const llmProviderService = createLlmProviderService(pool, {
      apiTarget: process.env.TKMIND_API_TARGET ?? 'https://127.0.0.1:18006',
      apiSecret: process.env.TKMIND_SERVER__SECRET_KEY ?? 'local-dev-secret',
    });
    const memoryConfigService = createMemoryV2AdminConfigService(pool);
    await restoreSavedState(pool, llmProviderService, memoryConfigService);
    await pool.end?.();
  } catch {
    // best effort restore
  }
  process.exit(1);
});