src/lib/server/hf-bucket/tools-snapshot.ts

import { fetchHfBucketText, HfBucketError } from './client';
import { parseBucketDrilldownIndex } from './details-payload';
import { fetchBucketRankingSnapshot, verifyBucketFileHash } from './ranking-snapshot';
import {
	parseBucketGroupedBars,
	parseBucketHeatmap,
	parseBucketPareto,
	parseBucketPerformance,
	parseBucketRadar,
	parseBucketRobustness,
	parseBucketScatter
} from './visualizations-payload';
import type {
	BucketMetricMap,
	BucketVisualizationPayloadByKey,
	BucketRankingSnapshot,
	BucketVisualizationSnapshot
} from './types';

function parseJson(text: string, path: string) {
	try {
		return JSON.parse(text) as unknown;
	} catch (cause) {
		throw new HfBucketError(`Bucket file "${path}" is not valid JSON.`, { cause });
	}
}

function duplicateKeys(items: string[]) {
	return items.filter((item, index) => items.indexOf(item) !== index);
}

function assertKnownId(ids: ReadonlySet<string>, id: string, label: string) {
	if (!ids.has(id)) throw new HfBucketError(`${label} references unknown id "${id}".`);
}

function verifyMetricMap(
	map: BucketMetricMap,
	rowIds: ReadonlySet<string>,
	modelIds: ReadonlySet<string>,
	label: string
) {
	for (const [rowId, row] of Object.entries(map)) {
		assertKnownId(rowIds, rowId, `${label} row`);
		for (const modelId of Object.keys(row)) {
			assertKnownId(modelIds, modelId, `${label} model`);
		}
	}
}

function verifyVisualizationSnapshot(snapshot: BucketVisualizationSnapshot) {
	const modelIds = new Set(snapshot.catalog.models.map((model) => model.model_id));
	const groupIds = new Set(snapshot.catalog.groups.map((group) => group.group_id));
	const datasetIds = new Set(snapshot.catalog.datasets.map((dataset) => dataset.dataset_id));

	for (const modelId of snapshot.radar.model_ids)
		assertKnownId(modelIds, modelId, 'radar.model_ids');
	for (const modelId of snapshot.radar.default_model_ids) {
		assertKnownId(modelIds, modelId, 'radar.default_model_ids');
	}
	for (const groupId of snapshot.radar.group_ids)
		assertKnownId(groupIds, groupId, 'radar.group_ids');
	verifyMetricMap(snapshot.radar.values, groupIds, modelIds, 'radar.values');
	verifyMetricMap(snapshot.radar.f1_values, groupIds, modelIds, 'radar.f1_values');

	if (duplicateKeys(snapshot.scatter.row_ids).length) {
		throw new HfBucketError('scatter.row_ids contains duplicate values.');
	}
	for (const point of snapshot.scatter.points) {
		assertKnownId(modelIds, point.model_id, 'scatter.points');
	}
	for (const [rowId, row] of Object.entries(snapshot.scatter.values)) {
		if (!snapshot.scatter.row_ids.includes(rowId as (typeof snapshot.scatter.row_ids)[number])) {
			throw new HfBucketError(`scatter.values references unknown row "${rowId}".`);
		}
		for (const modelId of Object.keys(row)) assertKnownId(modelIds, modelId, 'scatter.values');
	}

	const unknownHeatmapRowIds = snapshot.heatmap.row_ids.filter(
		(datasetId) => !datasetIds.has(datasetId)
	);
	if (unknownHeatmapRowIds.length) {
		console.warn(
			'[hf-bucket] Ignoring heatmap rows that reference unknown dataset ids.',
			unknownHeatmapRowIds
		);
	}
	for (const modelId of snapshot.heatmap.column_ids) {
		assertKnownId(modelIds, modelId, 'heatmap.column_ids');
	}
	const heatmapRowIds = new Set(snapshot.heatmap.row_ids);
	for (const [rowId, row] of Object.entries(snapshot.heatmap.values)) {
		if (!heatmapRowIds.has(rowId)) {
			throw new HfBucketError(`heatmap.values references unknown row "${rowId}".`);
		}
		if (!datasetIds.has(rowId)) continue;
		for (const modelId of Object.keys(row)) assertKnownId(modelIds, modelId, 'heatmap.values');
	}

	for (const groupId of snapshot.grouped_bars.group_ids) {
		assertKnownId(groupIds, groupId, 'grouped_bars.group_ids');
	}
	for (const model of snapshot.grouped_bars.models) {
		assertKnownId(modelIds, model.model_id, 'grouped_bars.models');
		for (const groupId of Object.keys(model.groups)) {
			assertKnownId(groupIds, groupId, 'grouped_bars.models.groups');
		}
	}

	for (const modelId of Object.keys(snapshot.pareto.by_model_id)) {
		assertKnownId(modelIds, modelId, 'pareto.by_model_id');
	}
	for (const row of snapshot.performance.rows) {
		assertKnownId(modelIds, row.model_id, 'performance.rows');
	}
	for (const row of snapshot.robustness.rows) {
		assertKnownId(modelIds, row.model_id, 'robustness.rows');
	}

	for (const row of snapshot.drilldownIndex.heatmap) {
		assertKnownId(modelIds, row.model_id, 'drilldown_index.heatmap.model_id');
		assertKnownId(groupIds, row.group_id, 'drilldown_index.heatmap.group_id');
		for (const datasetId of row.dataset_ids) {
			assertKnownId(datasetIds, datasetId, 'drilldown_index.heatmap.dataset_ids');
		}
	}
}

async function fetchVisualizationText(
	snapshot: Awaited<ReturnType<typeof fetchBucketRankingSnapshot>>,
	key: keyof BucketVisualizationPayloadByKey | 'drilldown_index'
) {
	const path = snapshot.manifest.files[key];
	const text = await fetchHfBucketText(path);
	verifyBucketFileHash(snapshot.manifest, key, text);
	return { path, text };
}

export async function fetchBucketVisualizationSnapshot(
	rankingSnapshot?: BucketRankingSnapshot
): Promise<BucketVisualizationSnapshot> {
	rankingSnapshot ??= await fetchBucketRankingSnapshot();
	const drilldownIndexText = await fetchVisualizationText(rankingSnapshot, 'drilldown_index');
	const radarText = await fetchVisualizationText(rankingSnapshot, 'radar');
	const scatterText = await fetchVisualizationText(rankingSnapshot, 'scatter');
	const heatmapText = await fetchVisualizationText(rankingSnapshot, 'heatmap');
	const groupedBarsText = await fetchVisualizationText(rankingSnapshot, 'grouped_bars');
	const paretoText = await fetchVisualizationText(rankingSnapshot, 'pareto');
	const performanceText = await fetchVisualizationText(rankingSnapshot, 'performance');
	const robustnessText = await fetchVisualizationText(rankingSnapshot, 'robustness');

	let snapshot: BucketVisualizationSnapshot;
	try {
		snapshot = {
			...rankingSnapshot,
			drilldownIndex: parseBucketDrilldownIndex(
				parseJson(drilldownIndexText.text, drilldownIndexText.path)
			),
			radar: parseBucketRadar(parseJson(radarText.text, radarText.path)),
			scatter: parseBucketScatter(parseJson(scatterText.text, scatterText.path)),
			heatmap: parseBucketHeatmap(parseJson(heatmapText.text, heatmapText.path)),
			grouped_bars: parseBucketGroupedBars(parseJson(groupedBarsText.text, groupedBarsText.path)),
			pareto: parseBucketPareto(parseJson(paretoText.text, paretoText.path)),
			performance: parseBucketPerformance(parseJson(performanceText.text, performanceText.path)),
			robustness: parseBucketRobustness(parseJson(robustnessText.text, robustnessText.path))
		};
	} catch (cause) {
		if (cause instanceof HfBucketError) throw cause;
		throw new HfBucketError('Visualization bucket payload failed validation.', { cause });
	}

	verifyVisualizationSnapshot(snapshot);
	return snapshot;
}