/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { writeFile } from 'fs/promises';
import { json2csv } from 'json-2-csv';
import { EvalRun } from '../types/eval';
import { logger } from '../utils/logger';
const SUPPORTED_FORMATS: Record<string, EvalExporter> = {
csv: toCsv,
json: toJson,
};
/**
* Interface for exporting evaluation to file.
*/
export interface EvalExporter {
(evalRun: EvalRun, filePath: string): Promise<void>;
}
/**
* Export an evalRun to csv.
*/
export async function toCsv(evalRun: EvalRun, filePath: string) {
// manually unpack the metrics into columns because we don't always want to pull nested fields into columns.
const unpackedCases = unpackMetricsToColumns(evalRun);
const csvRecords = json2csv(unpackedCases, {
emptyFieldValue: '',
expandNestedObjects: false,
});
logger.info(`Writing csv results to '${filePath}'...`);
await writeFile(filePath, csvRecords);
}
/**
* Export an evalRun to json.
*/
export async function toJson(evalRun: EvalRun, filePath: string) {
logger.info(`Writing json results to '${filePath}'...`);
await writeFile(filePath, JSON.stringify(evalRun.results, undefined, ' '));
}
function unpackMetricsToColumns(evalRun: EvalRun): Record<string, any>[] {
return evalRun.results.map((result) => {
let record: Record<string, any> = {
...result,
};
// remove metrics so that we can unnest them
delete record['metrics'];
result.metrics?.forEach((metric) => {
record[`${metric.evaluator}_score`] = metric.score;
record[`${metric.evaluator}_rationale`] = metric.rationale;
record[`${metric.evaluator}_error`] = metric.error;
record[`${metric.evaluator}_traceId`] = metric.traceId;
record[`${metric.evaluator}_spanId`] = metric.spanId;
});
return record;
});
}
export function getExporterForString(outputFormat: string): EvalExporter {
if (!(outputFormat in SUPPORTED_FORMATS)) {
logger.info(
`Encountered unrecognized output format ${outputFormat}. Defaulting to json.`
);
return toJson;
}
return SUPPORTED_FORMATS[outputFormat];
}