localFileEvalStore.ts•5.13 kB
/**
* 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 fs from 'fs';
import { readFile, unlink, writeFile } from 'fs/promises';
import path from 'path';
import { createInterface } from 'readline';
import type { ListEvalKeysRequest, ListEvalKeysResponse } from '../types/apis';
import {
EvalRunKeySchema,
EvalRunSchema,
type EvalRun,
type EvalRunKey,
type EvalStore,
} from '../types/eval';
import { logger } from '../utils/logger';
/**
* A local, file-based EvalStore implementation.
*/
export class LocalFileEvalStore implements EvalStore {
private storeRoot: string = '';
private indexFile: string = '';
private static cachedEvalStore: LocalFileEvalStore | null = null;
private async init() {
this.storeRoot = this.generateRootPath();
this.indexFile = await this.resolveIndexFile();
fs.mkdirSync(this.storeRoot, { recursive: true });
if (!fs.existsSync(this.indexFile)) {
fs.writeFileSync(path.resolve(this.indexFile), JSON.stringify({}));
}
logger.debug(
`Initialized local file eval store at root: ${this.storeRoot}`
);
}
static async getEvalStore() {
if (!this.cachedEvalStore) {
this.cachedEvalStore = new LocalFileEvalStore();
await this.cachedEvalStore.init();
}
return this.cachedEvalStore;
}
static reset() {
this.cachedEvalStore = null;
}
async save(evalRun: EvalRun): Promise<void> {
const fileName = this.resolveEvalFilename(evalRun.key.evalRunId);
logger.debug(
`Saving EvalRun ${evalRun.key.evalRunId} to ` +
path.resolve(this.storeRoot, fileName)
);
await writeFile(
path.resolve(this.storeRoot, fileName),
JSON.stringify(evalRun)
);
const index = await this.getEvalsIndex();
index[evalRun.key.evalRunId] = evalRun.key;
await writeFile(
path.resolve(this.indexFile),
JSON.stringify(index, null, 2)
);
}
async load(evalRunId: string): Promise<EvalRun | undefined> {
const filePath = path.resolve(
this.storeRoot,
this.resolveEvalFilename(evalRunId)
);
if (!fs.existsSync(filePath)) {
return undefined;
}
return await readFile(filePath, 'utf8').then((data) =>
EvalRunSchema.parse(JSON.parse(data))
);
}
async list(
query?: ListEvalKeysRequest | undefined
): Promise<ListEvalKeysResponse> {
logger.debug(`Listing keys for filter: ${JSON.stringify(query)}`);
let keys = await this.getEvalsIndex().then((index) => Object.values(index));
if (query?.filter?.actionRef) {
keys = keys.filter((key) => key.actionRef === query?.filter?.actionRef);
}
return {
evalRunKeys: keys,
};
}
async delete(evalRunId: string): Promise<void> {
const filePath = path.resolve(
this.storeRoot,
this.resolveEvalFilename(evalRunId)
);
if (fs.existsSync(filePath)) {
await unlink(filePath);
const index = await this.getEvalsIndex();
delete index[evalRunId];
await writeFile(
path.resolve(this.indexFile),
JSON.stringify(index, null, 2)
);
}
}
private resolveEvalFilename(evalRunId: string): string {
return `${evalRunId}.json`;
}
private async resolveIndexFile(): Promise<string> {
const txtPath = path.resolve(this.storeRoot, 'index.txt');
const jsonPath = path.resolve(this.storeRoot, 'index.json');
if (fs.existsSync(txtPath)) {
// Copy over index, delete txt file
const keys = await this.processLineByLine(txtPath);
await writeFile(path.resolve(jsonPath), JSON.stringify(keys, null, 2));
await unlink(txtPath);
}
return jsonPath;
}
private async processLineByLine(filePath: string) {
const fileStream = fs.createReadStream(filePath);
const keys: Record<string, EvalRunKey> = {};
const rl = createInterface({
input: fileStream,
crlfDelay: Infinity,
});
for await (const line of rl) {
try {
const entry = JSON.parse(line);
const runKey = EvalRunKeySchema.parse(entry);
keys[runKey.evalRunId] = runKey;
} catch (e) {
logger.debug(`Error parsing ${line}:\n`, JSON.stringify(e));
}
}
return keys;
}
private generateRootPath(): string {
return path.resolve(process.cwd(), '.genkit', 'evals');
}
private async getEvalsIndex(): Promise<Record<string, EvalRunKey>> {
if (!fs.existsSync(this.indexFile)) {
return Promise.resolve({} as any);
}
return await readFile(path.resolve(this.indexFile), 'utf8').then((data) =>
JSON.parse(data)
);
}
}