Spaces:

seonglae
/

AgentStateGraph-Browser

Paused

AgentStateGraph-Browser / scripts /build-data.ts
Interactive research article
9e2b9b8 about 2 months ago
38.5 kB
/**
 * Preprocesses experiment results into a single JSON for the browser.
 * Computes per-trace FSM traversals for the interpretability view.
 * Run: npx tsx scripts/build-data.ts
 */
import { readFileSync, writeFileSync, mkdirSync } from 'fs'
import path from 'path'
import { fileURLToPath } from 'url'

const __dirname = path.dirname(fileURLToPath(import.meta.url))
const ROOT = path.resolve(__dirname, '../..')
// Active 12 datasets (paper set). tau-bench v1 (taubench-airline/retail) was dropped
// and replaced by tau2-bench; excluded here so the dashboard reflects 12, not 14.
const DATASETS = [
  'swesmith', 'whoandwhen', 'mind2web',
  'sweagent',
  'gui-odyssey', 'webarena', 'agentnet',
  'tau2bench-airline', 'tau2bench-retail', 'tau2bench-telecom',
  'atbench', 'osworld',
] as const

// ---------- Types ----------

interface TraceStep {
  idx: number
  role: string
  activity: string
  fromState: string
  toState: string | null  // null = FSM failure
  consumed: boolean
}

interface TraceDetail {
  id: string
  split: 'train' | 'test'
  success: boolean | null
  fullFitness: number | null
  successFitness: number | null
  steps: TraceStep[]
  stateSequence: string[]
  fitness: number
  consumed: number
  total: number
  firstFailIdx: number  // -1 if no failure
  stateVisits: Record<string, number>
}

interface DatasetBundle {
  id: string
  meta: {
    name: string
    shortName: string
    venue: string
    traces: number
    trainSize: number
    testSize: number
    activities: number
    description: string
  }
  baselines: Array<{
    method: string
    label: string
    trainFitness: number
    testFitness: number
    states: number | null
    transitions: number | null
    randomAcceptRate: number | null
    permutedAcceptRate: number | null
    pmPrecision: number | null
    category: string
    note: string
  }>
  convergence: {
    stateConvergenceAt: number
    transitionConvergenceAt: number
    fitnessConvergenceAt: number
    curve: Array<{ n: number; states: number; transitions: number; fitness: number }>
  }
  traces: TraceDetail[]
  downstream: {
    fitnessProfile: {
      mean: number; std: number; min: number; max: number
      anomalyCount: number; anomalies: string[]
      histogram: Array<{ min: number; max: number; count: number }>
    } | null
    behavioralProfiling: {
      stateVisitDistribution: Array<{ state: string; count: number; fraction: number }>
      transitionFrequency: Array<{ transition: string; count: number; fraction: number }>
      deadTransitions: string[]
      transitionCoverage: number
      avgUniqueStatesPerTrace: number
    } | null
    structuralAnalysis: {
      density: number
      hubs: Array<{ id: string; outDegree: number; inDegree: number }>
      sinks: string[]
      sources: string[]
      selfLoops: string[]
      nonTrivialSCCs: string[][]
    } | null
    differentialFSM: {
      successOnlyTransitions: string[]
      failureOnlyTransitions: string[]
      sharedTransitions: string[]
    } | null
    failurePrediction: Record<string, unknown> | null
    mistakeLocalization: Record<string, unknown> | null
    neuralProbe: {
      seq_mlp: { cv: number; std: number; holdout: number }
      fsm_mlp: { cv: number; std: number; holdout: number }
      seq_gru: { cv: number; std: number; holdout: number }
      fsm_gru: { cv: number; std: number; holdout: number }
      seq_transf: { cv: number; std: number; holdout: number }
      fsm_transf: { cv: number; std: number; holdout: number }
    } | null
    workflowMemory: {
      n: number
      noMemory: { accuracy: number; correct: number }
      awm: { accuracy: number; correct: number }
      fsm: { accuracy: number; correct: number }
      fsmGain: number
    } | null
    crossModel: {
      models: Array<{ name: string; shortName: string; traces: number; successRate: number }>
      matrices: {
        fitness: number[][]
        auroc: number[][]
        stateCount: number[][]
      }
      summary: {
        diagonalAUROC: { mean: number; std: number }
        offDiagonalAUROC: { mean: number; std: number }
        aurocTransferGap: number
      }
    } | null
    monitor: {
      trainSize: number; testSize: number; succSize: number; failSize: number
      configs: Array<{
        config: string; gamma: number; window: number
        monitoringCurve: Array<{ point: number; threshold: number; f1: number; precision: number; recall: number }>
      }>
      baselines: Record<string, { auroc: number; curve?: Array<{ point: number; f1: number }> }>
    } | null
    counterfactual: {
      nSuccess: number; nFailure: number
      decisionPoints: Array<{ state: string; branching: number; targets: string[] }>
      pathDiversity: { uniqueSuccPaths: number; uniqueFailPaths: number; overlap: number }
      dpDivergence: Array<{ state: string; jsd: number }>
      editDistance: { meanInternal?: number; meanCross?: number } | null
    } | null
  }
  multiseed: {
    numSeeds: number
    ourFSM: { testFitness: { mean: number; std: number }; states: { mean: number; std: number } }
    rpni: { testFitness: { mean: number; std: number }; states: { mean: number; std: number } }
    awm: { testFitness: { mean: number; std: number } }
  } | null
  graph: {
    states: Array<{ id: string }>
    transitions: Array<{ source: string; target: string; subject: string }>
  }
}

const META: Record<string, DatasetBundle['meta']> = {
  swesmith: {
    name: 'SWE-smith',
    shortName: 'SWE',
    venue: 'NeurIPS 2025 D&B',
    traces: 500,
    trainSize: 400,
    testSize: 100,
    activities: 9,
    description: 'Coding agent trajectories with tool calls (bash, str_replace_editor, submit)',
  },
  whoandwhen: {
    name: 'Who & When to Delegate',
    shortName: 'W&W',
    venue: 'arXiv 2505.00212',
    traces: 184,
    trainSize: 147,
    testSize: 37,
    activities: 8,
    description: 'Multi-agent delegation failures across 8 specialized actors',
  },
  mind2web: {
    name: 'Mind2Web',
    shortName: 'M2W',
    venue: 'NeurIPS 2023',
    traces: 500,
    trainSize: 400,
    testSize: 100,
    activities: 7,
    description: 'Web navigation agent traces with click, type, select, hover, enter actions',
  },
  'taubench-airline': {
    name: 'tau-bench (airline)',
    shortName: 'tau-air',
    venue: 'Sierra 2024',
    traces: 200,
    trainSize: 160,
    testSize: 40,
    activities: 5,
    description: 'Customer service agent traces for airline booking/cancellation with 12 APIs',
  },
  'taubench-retail': {
    name: 'tau-bench (retail)',
    shortName: 'tau-ret',
    venue: 'Sierra 2024',
    traces: 460,
    trainSize: 368,
    testSize: 92,
    activities: 5,
    description: 'Customer service agent traces for retail order management with 14 APIs',
  },
  sweagent: {
    name: 'SWE-agent',
    shortName: 'SWE-a',
    venue: 'NeurIPS 2024',
    traces: 2000,
    trainSize: 1600,
    testSize: 400,
    activities: 68,
    description: 'SWE-agent coding trajectories with 68 unique bash/editor commands',
  },
  atbench: {
    name: 'ATBench',
    shortName: 'ATBench',
    venue: 'arXiv 2604.02022',
    traces: 1000,
    trainSize: 800,
    testSize: 200,
    activities: 14,
    description: 'Agent trajectory safety benchmark with balanced safe/unsafe outcomes (503 safe / 497 unsafe)',
  },
  osworld: {
    name: 'OSWorld',
    shortName: 'OSWorld',
    venue: 'NeurIPS 2024',
    traces: 2166,
    trainSize: 1732,
    testSize: 434,
    activities: 27,
    description: 'Desktop GUI agent traces (Gelato-30B, GTA1-32B) across real OS applications',
  },
  'gui-odyssey': {
    name: 'GUI-Odyssey',
    shortName: 'GUI',
    venue: 'CVPR 2025',
    traces: 7735,
    trainSize: 6188,
    testSize: 1547,
    activities: 6,
    description: 'Mobile GUI agent traces with click, type, scroll, swipe, enter, home actions',
  },
  webarena: {
    name: 'WebArena',
    shortName: 'WA',
    venue: 'ICLR 2024 Oral',
    traces: 8337,
    trainSize: 6670,
    testSize: 1667,
    activities: 24,
    description: 'Web agent traces on realistic websites with navigation, form-filling, and content extraction',
  },
  agentnet: {
    name: 'AgentNet',
    shortName: 'ANet',
    venue: 'NeurIPS 2024',
    traces: 5000,
    trainSize: 4000,
    testSize: 1000,
    activities: 24,
    description: 'Desktop GUI agent traces across diverse applications',
  },
  'tau2bench-airline': {
    name: 'tau2-bench (airline)',
    shortName: 'tau2-air',
    venue: 'arXiv 2506.07982',
    traces: 800,
    trainSize: 640,
    testSize: 160,
    activities: 17,
    description: 'Customer service agent traces (4 LLMs) for airline booking with 12 APIs',
  },
  'tau2bench-retail': {
    name: 'tau2-bench (retail)',
    shortName: 'tau2-ret',
    venue: 'arXiv 2506.07982',
    traces: 1824,
    trainSize: 1459,
    testSize: 365,
    activities: 18,
    description: 'Customer service agent traces (4 LLMs) for retail order management with 14 APIs',
  },
  'tau2bench-telecom': {
    name: 'tau2-bench (telecom)',
    shortName: 'tau2-tel',
    venue: 'arXiv 2506.07982',
    traces: 1824,
    trainSize: 1459,
    testSize: 365,
    activities: 4,
    description: 'Customer service agent traces (4 LLMs) for telecom workflows',
  },
}

// ---------- Activity extraction (must match shared.ts exactly) ----------

function extractActivity(role: string, content: string, agentName?: string): string {
  const actor = agentName || role
  const c = content.trim()
  const bracketMatch = c.match(/^\[(\w+)\]/)
  if (bracketMatch) return `${actor}:${bracketMatch[1].toLowerCase()}`
  if (c.includes('function_call') || c.includes('tool_call')) {
    // Match "tool_call:toolname" anywhere in content (loader format)
    const prefixMatch = c.match(/tool_call:(\w+)/)
    if (prefixMatch) return `${actor}:tool:${prefixMatch[1]}`
    // Match "name='toolname'" or "name: toolname" (OpenAI/generic format)
    const nameMatch = c.match(/name['"=:\s]+(\w+)/)
    return nameMatch ? `${actor}:tool:${nameMatch[1]}` : `${actor}:tool_call`
  }
  return `${actor}:text`
}

// ---------- FSM replay ----------

type TransitionFn = (state: string, symbol: string) => string | null

function graphToTransitionFn(graph: { transitions: Array<{ source: string; target: string }> }): TransitionFn {
  const valid = new Set<string>()
  for (const t of graph.transitions) {
    valid.add(`${t.source}|${t.target}`)
  }
  return (current: string, symbol: string) =>
    valid.has(`${current}|${symbol}`) ? symbol : null
}

function replayTrace(
  transitionFn: TransitionFn,
  messages: Array<{ role: string; content: string; name?: string }>,
): { steps: TraceStep[]; stateSequence: string[]; fitness: number; consumed: number; total: number; firstFailIdx: number; stateVisits: Record<string, number> } {
  let current = 'init'
  const steps: TraceStep[] = []
  const stateSequence: string[] = ['init']
  const stateVisits: Record<string, number> = { init: 1 }
  let consumed = 0
  let firstFailIdx = -1

  for (let i = 0; i < messages.length; i++) {
    const msg = messages[i]
    const activity = extractActivity(msg.role, msg.content, msg.name)
    const next = transitionFn(current, activity)

    if (next !== null) {
      steps.push({ idx: i, role: msg.role, activity, fromState: current, toState: next, consumed: true })
      stateSequence.push(next)
      stateVisits[next] = (stateVisits[next] || 0) + 1
      current = next
      consumed++
    } else {
      steps.push({ idx: i, role: msg.role, activity, fromState: current, toState: null, consumed: false })
      stateSequence.push(`FAIL:${activity}`)
      if (firstFailIdx === -1) firstFailIdx = i
    }
  }

  return {
    steps,
    stateSequence,
    fitness: messages.length > 0 ? consumed / messages.length : 1,
    consumed,
    total: messages.length,
    firstFailIdx,
    stateVisits,
  }
}

// ---------- Main ----------

function loadJSON(p: string) {
  return JSON.parse(readFileSync(path.join(ROOT, p), 'utf-8'))
}

function tryLoadJSON(p: string): Record<string, unknown> | null {
  try { return loadJSON(p) } catch { return null }
}

function loadTauBenchRaw(domain: string): Array<{ id: string; messages: Array<{ role: string; content: string; name?: string }>; annotations?: { success?: boolean } }> {
  const raw: Array<{ task_id: number; reward: number; traj: Array<{ role: string; content: string; tool_calls?: Array<{ function: { name: string } }> }> }> =
    loadJSON(`data/taubench/${domain}.json`)
  return raw.map(entry => ({
    id: `taubench-${domain}-${entry.task_id}`,
    messages: entry.traj.map(step => {
      if (step.role === 'assistant' && step.tool_calls && step.tool_calls.length > 0) {
        return { role: 'assistant', content: `tool_call:${step.tool_calls[0].function.name}` }
      }
      if (step.role === 'tool') {
        const c = step.content || ''
        return { role: 'tool', content: c.length > 500 ? c.slice(0, 500) : c }
      }
      return { role: step.role || 'user', content: step.content || '' }
    }),
    annotations: { success: entry.reward === 1.0 },
  }))
}

const TAU2_FILES: Record<string, string[]> = {
  airline: [
    'data/tau2bench/gpt-4.1-2025-04-14_airline_default_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/claude-3-7-sonnet-20250219_airline_default_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/gpt-4.1-mini-2025-04-14_airline_base_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/o4-mini-2025-04-16_airline_default_gpt-4.1-2025-04-14_4trials.json',
  ],
  retail: [
    'data/tau2bench/gpt-4.1-2025-04-14_retail_default_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/claude-3-7-sonnet-20250219_retail_default_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/gpt-4.1-mini-2025-04-14_retail_base_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/o4-mini-2025-04-16_retail_default_gpt-4.1-2025-04-14_4trials.json',
  ],
  telecom: [
    'data/tau2bench/telecom_gpt41.json',
    'data/tau2bench/claude-3-7-sonnet-20250219_telecom_default_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/gpt-4.1-mini-2025-04-14_telecom_base_gpt-4.1-2025-04-14_4trials.json',
    'data/tau2bench/o4-mini-2025-04-16_telecom-workflow_default_gpt-4.1-2025-04-14_4trials.json',
  ],
}

function loadTau2BenchRaw(domain: string): Array<{ id: string; messages: Array<{ role: string; content: string; name?: string }>; annotations?: { success?: boolean } }> {
  const files = TAU2_FILES[domain] || []
  const allTraces: Array<{ id: string; messages: Array<{ role: string; content: string; name?: string }>; annotations?: { success?: boolean } }> = []
  for (const f of files) {
    try {
      const data: { simulations: Array<{ task_id: string; messages: Array<{ role: string; content: string; tool_calls?: Array<{ name?: string; function?: { name: string } }> }>; reward_info: { reward: number } }> } = loadJSON(f)
      const modelTag = f.replace(/.*\//, '').replace(/\.json$/, '').slice(0, 30)
      for (let idx = 0; idx < data.simulations.length; idx++) {
        const sim = data.simulations[idx]
        const messages = sim.messages
          .filter((m: { role: string }) => m.role !== 'system')
          .map((step: { role: string; content: string; tool_calls?: Array<{ name?: string; function?: { name: string } }> }) => {
            if (step.tool_calls && step.tool_calls.length > 0) {
              const tc = step.tool_calls[0]
              const toolName = tc.name || tc.function?.name || 'unknown'
              return { role: 'assistant', content: `tool_call:${toolName}` }
            }
            if (step.role === 'tool') {
              const c = step.content || ''
              return { role: 'tool', content: c.length > 500 ? c.slice(0, 500) : c }
            }
            return { role: step.role || 'user', content: step.content || '' }
          })
        allTraces.push({
          id: `tau2bench-${domain}-${modelTag}-${sim.task_id}-${idx}`,
          messages,
          annotations: { success: sim.reward_info.reward === 1.0 },
        })
      }
    } catch { console.log(`  Warning: could not load ${f}`) }
  }
  return allTraces
}

const neuralProbe: Record<string, DatasetBundle['downstream']['neuralProbe']> =
  tryLoadJSON('experiments/fsm-neural-probe/results.json') || {}

// Paper Table 4 (LLM-judged top-1, gpt-4.1-mini, fsmWorkflow format, FULL validation split).
// Authoritative for the 8 datasets in the paper; older poc-fsm-memory JSONs used gpt-4o-mini
// with a different prompt format and are not directly comparable.
const PAPER_TABLE_4: Record<string, { n: number; awm: number; fsm: number }> = {
  webarena: { n: 4800, awm: 0.655, fsm: 0.812 },
  swesmith: { n: 300, awm: 0.747, fsm: 1.000 },
  sweagent: { n: 1200, awm: 0.677, fsm: 0.705 },
  'tau2bench-telecom': { n: 1095, awm: 0.285, fsm: 0.456 },
  'tau2bench-retail': { n: 1095, awm: 0.529, fsm: 0.651 },
  'tau2bench-airline': { n: 480, awm: 0.565, fsm: 0.573 },
  atbench: { n: 600, awm: 0.478, fsm: 0.625 },
  osworld: { n: 1286, awm: 0.550, fsm: 0.707 },
}

const monitorRoot: Record<string, DatasetBundle['downstream']['monitor']> =
  tryLoadJSON('experiments/poc-online-monitor/online-monitor-results.json') || {}

const cfRoot: Record<string, DatasetBundle['downstream']['counterfactual']> =
  tryLoadJSON('experiments/counterfactual/counterfactual-results.json') || {}

// Incremental-convergence runs (6 datasets). The per-dataset evolution-vN folders are
// LFS stubs for everything except atbench/osworld, so the dashboard convergence tab only
// saw 2 of 6 — the same 'taubench-*' vs 'tau2bench-*' keying issue. Build the curve from
// this file (fitnessTrajectory + final state/edge counts) when evolution snapshots are absent.
const convRateRoot: Record<string, {
  finalStates: number; finalEdges: number; trainTraces: number
  fitnessTrajectory: Array<{ pctData: number; fitness: number }>
}> = tryLoadJSON('experiments/convergence-rate/convergence-results.json') || {}

function convRateCurve(id: string): { curve: DatasetBundle['convergence']['curve']; fitnessAt: number } | null {
  const c = convRateRoot[id] || convRateRoot[monitorKey(id)]
  if (!c || !c.fitnessTrajectory?.length) return null
  const curve = c.fitnessTrajectory.map(p => ({
    n: Math.round(p.pctData * c.trainTraces),
    states: c.finalStates,        // per-point counts unavailable; state count converges near-instantly
    transitions: c.finalEdges,
    fitness: p.fitness,
  }))
  const hit = c.fitnessTrajectory.find(p => p.fitness >= 0.95)
  return { curve, fitnessAt: hit ? Math.round(hit.pctData * c.trainTraces) : 0 }
}

// monitor/counterfactual JSONs key tau2-bench suites as 'taubench-*' (legacy
// single-tau); the dataset registry uses 'tau2bench-*'. Map before lookup so all
// four evaluated datasets (swesmith, sweagent, retail, airline) resolve, not just two.
function monitorKey(id: string): string {
  return id.replace(/^tau2bench-/, 'taubench-')
}

function loadMonitor(id: string): DatasetBundle['downstream']['monitor'] {
  return monitorRoot[id] || monitorRoot[monitorKey(id)] || null
}

function loadCounterfactual(id: string): DatasetBundle['downstream']['counterfactual'] {
  return cfRoot[id] || cfRoot[monitorKey(id)] || null
}

function loadCrossModel(id: string): DatasetBundle['downstream']['crossModel'] {
  const d = tryLoadJSON(`experiments/${id}-cross-model/transfer-matrix.json`)
  if (!d || typeof d !== 'object') return null
  // Matrices in the experiment JSON are stored as {labels, values}; extract the 2-D values
  // (older files may already be a raw number[][]).
  type Mat = number[][] | { values?: number[][] }
  const matVals = (m: Mat | undefined): number[][] =>
    Array.isArray(m) ? m : (m?.values || [])
  const obj = d as {
    models?: Array<{ name: string; shortName: string; totalTraces: number; successRate: number }>
    matrices?: { fitness?: Mat; auroc?: Mat; stateCount?: Mat }
    summary?: {
      diagonalAUROC?: { mean: number; std: number }
      offDiagonalAUROC?: { mean: number; std: number }
      aurocTransferGap?: number
    }
  }
  if (!obj.models || !obj.matrices || !obj.summary) return null
  return {
    models: obj.models.map(m => ({
      name: m.name, shortName: m.shortName, traces: m.totalTraces, successRate: m.successRate,
    })),
    matrices: {
      fitness: matVals(obj.matrices.fitness),
      auroc: matVals(obj.matrices.auroc),
      stateCount: matVals(obj.matrices.stateCount),
    },
    summary: {
      diagonalAUROC: obj.summary.diagonalAUROC || { mean: 0, std: 0 },
      offDiagonalAUROC: obj.summary.offDiagonalAUROC || { mean: 0, std: 0 },
      aurocTransferGap: obj.summary.aurocTransferGap || 0,
    },
  }
}

function loadWorkflowMemory(id: string): DatasetBundle['downstream']['workflowMemory'] {
  const paper = PAPER_TABLE_4[id]
  if (paper) {
    // Use paper-authoritative numbers; pull noMemory from the older JSON if available.
    const json = tryLoadJSON(`experiments/poc-fsm-memory/${id}-results.json`)
    const noMem = (json && typeof json === 'object'
      ? ((json as { llm?: { noMemory?: { accuracy: number } } }).llm?.noMemory?.accuracy ?? null)
      : null)
    return {
      n: paper.n,
      noMemory: { accuracy: noMem ?? 0, correct: noMem != null ? Math.round(noMem * paper.n) : 0 },
      awm: { accuracy: paper.awm, correct: Math.round(paper.awm * paper.n) },
      fsm: { accuracy: paper.fsm, correct: Math.round(paper.fsm * paper.n) },
      fsmGain: paper.fsm - paper.awm,
    }
  }
  return null
}

function buildDataset(id: string): DatasetBundle {
  // tau2-bench: the combined 4-model run lives in the UNVERSIONED folder; the versioned
  // (-v11) folders are stale single-model runs (telecom 5 states vs combined 43). Prefer
  // the combined run for tau2 so baselines/graph match the paper.
  const combinedFirst = id.startsWith('tau2bench-')

  // Load baselines: (tau2: unversioned combined first) → v12 → v11 → v8 → v7 → v6 → unversioned
  const baseline = (combinedFirst ? tryLoadJSON(`experiments/${id}-baselines/baseline-comparison.json`) : null)
    || tryLoadJSON(`experiments/${id}-baselines-v12/baseline-comparison.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v11/baseline-comparison.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v8/baseline-comparison.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v7/baseline-comparison.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v6/baseline-comparison.json`)
    || loadJSON(`experiments/${id}-baselines/baseline-comparison.json`)

  // Load evolution: v9 → v8 → v7 → v6 → v5
  let evolution: Record<string, unknown>
  const ev = tryLoadJSON(`experiments/${id}-evolution-v9/evolution-results.json`)
    || tryLoadJSON(`experiments/${id}-evolution-v8/evolution-results.json`)
    || tryLoadJSON(`experiments/${id}-evolution-v7/evolution-results.json`)
    || tryLoadJSON(`experiments/${id}-evolution-v6/evolution-results.json`)
    || tryLoadJSON(`experiments/${id}-evolution-v5/evolution-results.json`)
  evolution = ev || {}
  if (ev) console.log(`  Loaded evolution for ${id}`)
  else console.log(`  No evolution data for ${id}`)

  // Load downstream: v15 → v14 → v13 → v12 → v11 → v10-fixed → v9 → v8 → v7 → v2 → v6 → v5
  const downstream: Record<string, unknown> = tryLoadJSON(`experiments/${id}-downstream-v15/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v14/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v13/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v12/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v11/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v10-fixed/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v9/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v8/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v7/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v2/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v6/downstream-eval.json`)
    || tryLoadJSON(`experiments/${id}-downstream-v5/downstream-eval.json`)
    || {}
  if (Object.keys(downstream).length > 0) console.log(`  Loaded downstream for ${id}`)
  else console.log(`  No downstream data for ${id}`)

  // Load graph: (tau2: unversioned combined first) → v12 → v11 → v8 → v7 → v6 → unversioned
  const graph = (combinedFirst ? tryLoadJSON(`experiments/${id}-baselines/train-graph.json`) : null)
    || tryLoadJSON(`experiments/${id}-baselines-v12/train-graph.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v11/train-graph.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v8/train-graph.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v7/train-graph.json`)
    || tryLoadJSON(`experiments/${id}-baselines-v6/train-graph.json`)
    || loadJSON(`experiments/${id}-baselines/train-graph.json`)

  // Load raw traces for FSM replay
  type RawTrace = { id: string; messages: Array<{ role: string; content: string; name?: string }>; annotations?: { success?: boolean } }
  let rawTraces: RawTrace[]
  if (id === 'taubench-airline') {
    rawTraces = loadTauBenchRaw('airline')
  } else if (id === 'taubench-retail') {
    rawTraces = loadTauBenchRaw('retail')
  } else if (id.startsWith('tau2bench-')) {
    const domain = id.replace('tau2bench-', '')
    rawTraces = loadTau2BenchRaw(domain)
  } else {
    rawTraces = loadJSON(`data/${id}/traces.json`)
  }
  const traceMap = new Map(rawTraces.map(t => [t.id, t]))

  // For tau2bench: also register traces under the old ID format (without modelTag)
  // so baselines generated with the old loader still match.
  // Old format: tau2bench-{domain}-{task_id}-{idx}
  // New format: tau2bench-{domain}-{modelTag}-{task_id}-{idx}
  if (id.startsWith('tau2bench-')) {
    const domain = id.replace('tau2bench-', '')
    const files = TAU2_FILES[domain] || []
    for (const f of files) {
      try {
        const data: { simulations: Array<{ task_id: string }> } = loadJSON(f)
        for (let idx = 0; idx < data.simulations.length; idx++) {
          const sim = data.simulations[idx]
          const modelTag = f.replace(/.*\//, '').replace(/\.json$/, '').slice(0, 30)
          const newId = `tau2bench-${domain}-${modelTag}-${sim.task_id}-${idx}`
          const oldId = `tau2bench-${domain}-${sim.task_id}-${idx}`
          const trace = traceMap.get(newId)
          if (trace && !traceMap.has(oldId)) {
            traceMap.set(oldId, trace)
          }
        }
      } catch { /* skip */ }
    }
  }

  console.log(`  Loaded ${rawTraces.length} raw traces for replay (traceMap: ${traceMap.size} keys)`)

  // Build transition function from graph
  const transitionFn = graphToTransitionFn(graph)

  // Extract baselines
  const baselines: DatasetBundle['baselines'] = []

  // Our FSM
  const sg = baseline.stateGraph
  baselines.push({
    method: 'stateGraph',
    label: 'Our FSM',
    trainFitness: sg.trainFitness,
    testFitness: sg.testFitness,
    states: sg.stateCount,
    transitions: sg.transitionCount,
    randomAcceptRate: sg.precision?.randomAcceptanceRate ?? null,
    permutedAcceptRate: sg.precision?.permutedAcceptanceRate ?? null,
    pmPrecision: null,
    category: 'ours',
    note: sg.note,
  })

  // RPNI
  const rpni = baseline.rpni
  baselines.push({
    method: 'rpni',
    label: 'RPNI',
    trainFitness: rpni.trainFitness,
    testFitness: rpni.testFitness,
    states: rpni.stateCount,
    transitions: rpni.transitionCount,
    randomAcceptRate: rpni.precision?.randomAcceptanceRate ?? null,
    permutedAcceptRate: rpni.precision?.permutedAcceptanceRate ?? null,
    pmPrecision: null,
    category: 'automata',
    note: rpni.note,
  })

  // PM4Py miners
  const pm4py = baseline.pm4py?.miners || {}
  for (const [name, m] of Object.entries(pm4py) as [string, Record<string, unknown>][]) {
    if (!m.success) continue
    const label = name.charAt(0).toUpperCase() + name.slice(1) + ' Miner'
    baselines.push({
      method: name,
      label,
      trainFitness: m.train_fitness as number,
      testFitness: m.test_fitness as number,
      states: null,
      transitions: (m.transitions as number) || null,
      randomAcceptRate: null,
      permutedAcceptRate: null,
      pmPrecision: (m.test_precision as number) ?? null,
      category: 'process-mining',
      note: `${m.places || '?'}p, ${m.transitions || '?'}t`,
    })
  }

  // AWM
  if (baseline.awm) {
    const awm = baseline.awm
    baselines.push({
      method: 'awm',
      label: 'AWM',
      trainFitness: awm.trainCoverage || 0,
      testFitness: awm.testAvgMatchScore || 0,
      states: null,
      transitions: null,
      randomAcceptRate: null,
      permutedAcceptRate: null,
      pmPrecision: null,
      category: 'workflow',
      note: `${awm.count} workflows (success-only)`,
    })
  }

  if (baseline.awm_all) {
    const awm = baseline.awm_all
    baselines.push({
      method: 'awm_all',
      label: 'AWM-all',
      trainFitness: awm.trainCoverage || 0,
      testFitness: awm.testAvgMatchScore || 0,
      states: null,
      transitions: null,
      randomAcceptRate: null,
      permutedAcceptRate: null,
      pmPrecision: null,
      category: 'workflow',
      note: `${awm.count} workflows (all traces)`,
    })
  }

  // EDSM
  if (baseline.edsm && (baseline.edsm as Record<string, unknown>).success) {
    const edsm = baseline.edsm as Record<string, unknown>
    baselines.push({
      method: 'edsm',
      label: 'EDSM',
      trainFitness: (edsm.train_fitness as number) || 0,
      testFitness: (edsm.test_fitness as number) || 0,
      states: (edsm.state_count as number) || null,
      transitions: (edsm.transition_count as number) || null,
      randomAcceptRate: null,
      permutedAcceptRate: null,
      pmPrecision: null,
      category: 'automata',
      note: `EDSM (AALpy): ${edsm.state_count}S, ${edsm.alphabet_size || '?'} alphabet`,
    })
  }

  // Evolution curve
  const snapshots = evolution.snapshots || evolution.curve || []
  let curve = snapshots.map((p: Record<string, number>) => ({
    n: p.traceCount,
    states: p.states || p.stateCount,
    transitions: p.transitions || p.transitionCount,
    fitness: p.testFitness || p.fitness,
  }))
  const conv = evolution.convergence || {}
  // Fall back to the incremental-convergence-rate run when the evolution file is a stub.
  if (curve.length === 0) {
    const cr = convRateCurve(id)
    if (cr) {
      curve = cr.curve
      conv.fitnessConvergenceAt = conv.fitnessConvergenceAt || cr.fitnessAt
      console.log(`  Convergence curve from convergence-rate run for ${id}`)
    }
  }

  // Per-trace data with FSM traversal
  const trainIds: string[] = baseline.config?.trainIds || []
  const testIds: string[] = baseline.config?.testIds || []
  const perTrace: Record<string, { success: boolean; fullFitness: number; successFitness: number }> = {}
  for (const t of ((downstream as any).failurePrediction?.perTrace || [])) {
    perTrace[t.id] = { success: t.success, fullFitness: t.fullFitness, successFitness: t.successFitness }
  }

  const traces: TraceDetail[] = []
  let replayCount = 0

  function buildTraceDetail(tid: string, split: 'train' | 'test'): TraceDetail {
    const pt = perTrace[tid]
    const raw = traceMap.get(tid)

    if (raw) {
      const replay = replayTrace(transitionFn, raw.messages)
      replayCount++
      return {
        id: tid,
        split,
        success: pt?.success ?? raw.annotations?.success ?? null,
        fullFitness: pt?.fullFitness ?? null,
        successFitness: pt?.successFitness ?? null,
        steps: replay.steps,
        stateSequence: replay.stateSequence,
        fitness: replay.fitness,
        consumed: replay.consumed,
        total: replay.total,
        firstFailIdx: replay.firstFailIdx,
        stateVisits: replay.stateVisits,
      }
    }

    // Fallback: no raw trace data
    return {
      id: tid,
      split,
      success: pt?.success ?? null,
      fullFitness: pt?.fullFitness ?? null,
      successFitness: pt?.successFitness ?? null,
      steps: [],
      stateSequence: [],
      fitness: 0,
      consumed: 0,
      total: 0,
      firstFailIdx: -1,
      stateVisits: {},
    }
  }

  for (const tid of trainIds) traces.push(buildTraceDetail(tid, 'train'))
  for (const tid of testIds) traces.push(buildTraceDetail(tid, 'test'))
  console.log(`  Replayed ${replayCount}/${traces.length} traces through FSM`)

  // Load multi-seed data
  let multiseed: DatasetBundle['multiseed'] = null
  try {
    const ms = loadJSON(`experiments/${id}-multiseed/multiseed-results.json`)
    multiseed = {
      numSeeds: ms.config?.numSeeds || 5,
      ourFSM: ms.aggregated?.ourFSM || { testFitness: { mean: 0, std: 0 }, states: { mean: 0, std: 0 } },
      rpni: ms.aggregated?.rpni || { testFitness: { mean: 0, std: 0 }, states: { mean: 0, std: 0 } },
      awm: ms.aggregated?.awm || { testFitness: { mean: 0, std: 0 } },
    }
    console.log(`  Loaded multi-seed (${multiseed.numSeeds} seeds)`)
  } catch { console.log(`  No multi-seed data for ${id}`) }

  // Build full downstream
  const ds = downstream as Record<string, any>
  const fp = ds.fitnessProfile || null
  const bp = ds.behavioralProfiling || null
  const sa = ds.structuralAnalysis || null
  const df = ds.differentialFSM?.skipped ? null : (ds.differentialFSM || null)
  // Failure prediction is a labeled task; the 3 unlabeled datasets have no trustworthy
  // success labels (gui-odyssey yields a degenerate AUROC 1.000), so null them out to
  // match the paper's 9-labeled-dataset failure-prediction set.
  const UNLABELED = new Set(['gui-odyssey', 'mind2web', 'whoandwhen'])
  const pred = (UNLABELED.has(id) || ds.failurePrediction?.skipped) ? null : (ds.failurePrediction || null)
  const ml = ds.mistakeLocalization || null

  return {
    id,
    meta: META[id],
    baselines,
    traces,
    convergence: {
      stateConvergenceAt: conv.stateConvergenceAt || 0,
      transitionConvergenceAt: conv.transitionConvergenceAt || 0,
      fitnessConvergenceAt: conv.fitnessConvergenceAt || 0,
      curve,
    },
    downstream: {
      fitnessProfile: fp ? {
        mean: fp.mean, std: fp.std, min: fp.min, max: fp.max,
        anomalyCount: fp.anomalyCount, anomalies: fp.anomalies?.map((a: any) => a.id || a) || [],
        histogram: fp.histogram || [],
      } : null,
      behavioralProfiling: bp ? {
        stateVisitDistribution: bp.stateVisitDistribution || [],
        transitionFrequency: bp.transitionFrequency || [],
        deadTransitions: bp.deadTransitions || [],
        transitionCoverage: bp.transitionCoverage ?? 1,
        avgUniqueStatesPerTrace: bp.avgUniqueStatesPerTrace ?? 0,
      } : null,
      structuralAnalysis: sa ? {
        density: sa.density,
        hubs: sa.hubs || [],
        sinks: sa.sinks || [],
        sources: sa.sources || [],
        selfLoops: sa.selfLoops || [],
        nonTrivialSCCs: sa.nonTrivialSCCs || [],
      } : null,
      differentialFSM: df ? {
        successOnlyTransitions: df.successOnlyTransitions || [],
        failureOnlyTransitions: df.failureOnlyTransitions || [],
        sharedTransitions: typeof df.sharedTransitions === 'number'
          ? [] : (df.sharedTransitions || []),
      } : null,
      failurePrediction: pred,
      mistakeLocalization: ml,
      neuralProbe: neuralProbe[id] || null,
      workflowMemory: loadWorkflowMemory(id),
      crossModel: loadCrossModel(id),
      monitor: loadMonitor(id),
      counterfactual: loadCounterfactual(id),
    },
    multiseed,
    graph: {
      states: (graph.states || []).map((s: Record<string, string>) => ({ id: s.id })),
      transitions: (graph.transitions || []).map((t: Record<string, string>) => ({
        source: t.source,
        target: t.target,
        subject: t.subject || '',
      })),
    },
  }
}

// Cap heaviest datasets to keep each shard under 22 MB
// (Cloudflare Pages / Netlify hard limit is 25 MB per static file).
// Traces are sampled deterministically from the full set, preserving train/test split balance.
const SHARD_TRACE_CAP: Record<string, number> = {
  agentnet: 3500,
}

function capTraces(d: DatasetBundle): DatasetBundle {
  const cap = SHARD_TRACE_CAP[d.id]
  if (!cap || d.traces.length <= cap) return d
  const train = d.traces.filter(t => t.split === 'train')
  const test = d.traces.filter(t => t.split === 'test')
  const ratio = train.length / (train.length + test.length)
  const trainKeep = Math.round(cap * ratio)
  const testKeep = cap - trainKeep
  // Deterministic: keep the first N of each split (already shuffled by trainTestSplit seed).
  const truncated = [...train.slice(0, trainKeep), ...test.slice(0, testKeep)]
  return { ...d, traces: truncated }
}

function main() {
  const data = DATASETS.map(buildDataset).map(capTraces)
  const pub = path.join(__dirname, '../public')
  const shardDir = path.join(pub, 'datasets')
  mkdirSync(pub, { recursive: true })
  mkdirSync(shardDir, { recursive: true })

  // Per-dataset shards (lets deploy hosts with ≤25MB-per-file limits accept the build).
  for (const d of data) {
    writeFileSync(path.join(shardDir, `${d.id}.json`), JSON.stringify(d))
  }

  // Slim index for the landing/list views (everything except per-trace replay details).
  const index = data.map(d => ({
    id: d.id, meta: d.meta, baselines: d.baselines,
    convergence: d.convergence, multiseed: d.multiseed, graph: d.graph,
    downstream: {
      ...d.downstream,
      // Drop the largest passthrough blobs from the index to keep it < ~2MB.
      failurePrediction: d.downstream.failurePrediction
        ? { ...d.downstream.failurePrediction, perTrace: undefined }
        : null,
      mistakeLocalization: null,
    },
    traceCount: d.traces.length,
  }))
  writeFileSync(path.join(pub, 'index.json'), JSON.stringify(index))

  // Legacy monolithic file (kept for back-compat with older App.tsx).
  writeFileSync(path.join(pub, 'data.json'), JSON.stringify(data))

  console.log(`\nBuilt data.json + index.json + ${data.length} per-dataset shards`)
  for (const d of data) {
    const withSteps = d.traces.filter(t => t.steps.length > 0).length
    console.log(`  ${d.meta.name}: ${d.baselines.length} baselines, ${withSteps}/${d.traces.length} traces with FSM replay`)
  }
}

main()