fix(genkit-tools/evals): Switch to JSON index in EvalStore (#3196)

Author: ssbushi

genkit-tools/common/src/eval/evaluate.ts

@@ -207,7 +207,7 @@ export async function runEvaluation(params: {
   };
 
   logger.info('Finished evaluation, writing key...');
-  const evalStore = getEvalStore();
+  const evalStore = await getEvalStore();
   await evalStore.save(evalRun);
   return evalRun;
 }

genkit-tools/common/src/eval/index.ts

@@ -23,9 +23,8 @@ export * from './exporter';
 export * from './parser';
 export * from './validate';
 
-export function getEvalStore(): EvalStore {
-  // TODO: This should provide EvalStore, based on tools config.
-  return LocalFileEvalStore.getEvalStore();
+export async function getEvalStore(): Promise<EvalStore> {
+  return await LocalFileEvalStore.getEvalStore();
 }
 
 export function getDatasetStore(): DatasetStore {

genkit-tools/common/src/eval/localFileEvalStore.ts

@@ -15,9 +15,9 @@
  */
 
 import fs from 'fs';
-import { appendFile, readFile, unlink, writeFile } from 'fs/promises';
+import { readFile, unlink, writeFile } from 'fs/promises';
 import path from 'path';
-import * as readline from 'readline';
+import { createInterface } from 'readline';
 import type { ListEvalKeysRequest, ListEvalKeysResponse } from '../types/apis';
 import {
   EvalRunKeySchema,
@@ -32,26 +32,26 @@ import { logger } from '../utils/logger';
  * A local, file-based EvalStore implementation.
  */
 export class LocalFileEvalStore implements EvalStore {
-  private readonly storeRoot;
-  private readonly indexFile;
-  private readonly INDEX_DELIMITER = '\n';
+  private storeRoot: string = '';
+  private indexFile: string = '';
   private static cachedEvalStore: LocalFileEvalStore | null = null;
 
-  private constructor() {
+  private async init() {
     this.storeRoot = this.generateRootPath();
-    this.indexFile = this.getIndexFilePath();
+    this.indexFile = await this.resolveIndexFile();
     fs.mkdirSync(this.storeRoot, { recursive: true });
     if (!fs.existsSync(this.indexFile)) {
-      fs.writeFileSync(path.resolve(this.indexFile), '');
+      fs.writeFileSync(path.resolve(this.indexFile), JSON.stringify({}));
     }
     logger.debug(
       `Initialized local file eval store at root: ${this.storeRoot}`
     );
   }
 
-  static getEvalStore() {
+  static async getEvalStore() {
     if (!this.cachedEvalStore) {
       this.cachedEvalStore = new LocalFileEvalStore();
+      await this.cachedEvalStore.init();
     }
     return this.cachedEvalStore;
   }
@@ -61,7 +61,7 @@ export class LocalFileEvalStore implements EvalStore {
   }
 
   async save(evalRun: EvalRun): Promise<void> {
-    const fileName = this.generateFileName(evalRun.key.evalRunId);
+    const fileName = this.resolveEvalFilename(evalRun.key.evalRunId);
 
     logger.debug(
       `Saving EvalRun ${evalRun.key.evalRunId} to ` +
@@ -72,20 +72,18 @@ export class LocalFileEvalStore implements EvalStore {
       JSON.stringify(evalRun)
     );
 
-    logger.debug(
-      `Save EvalRunKey ${JSON.stringify(evalRun.key)} to ` +
-        path.resolve(this.indexFile)
-    );
-    await appendFile(
+    const index = await this.getEvalsIndex();
+    index[evalRun.key.evalRunId] = evalRun.key;
+    await writeFile(
       path.resolve(this.indexFile),
-      JSON.stringify(evalRun.key) + this.INDEX_DELIMITER
+      JSON.stringify(index, null, 2)
     );
   }
 
   async load(evalRunId: string): Promise<EvalRun | undefined> {
     const filePath = path.resolve(
       this.storeRoot,
-      this.generateFileName(evalRunId)
+      this.resolveEvalFilename(evalRunId)
     );
     if (!fs.existsSync(filePath)) {
       return undefined;
@@ -98,22 +96,11 @@ export class LocalFileEvalStore implements EvalStore {
   async list(
     query?: ListEvalKeysRequest | undefined
   ): Promise<ListEvalKeysResponse> {
-    let keys = await readFile(this.indexFile, 'utf8').then((data) => {
-      if (!data) {
-        return [];
-      }
-      // strip the final carriage return before parsing all lines
-      return data
-        .slice(0, -1)
-        .split(this.INDEX_DELIMITER)
-        .map(this.parseLineToKey);
-    });
-
-    logger.debug(`Found keys: ${JSON.stringify(keys)}`);
+    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);
-      logger.debug(`Filtered keys: ${JSON.stringify(keys)}`);
     }
 
     return {
@@ -124,50 +111,66 @@ export class LocalFileEvalStore implements EvalStore {
   async delete(evalRunId: string): Promise<void> {
     const filePath = path.resolve(
       this.storeRoot,
-      this.generateFileName(evalRunId)
+      this.resolveEvalFilename(evalRunId)
     );
-    if (!fs.existsSync(filePath)) {
-      throw new Error(`Cannot find evalRun with id '${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)
+      );
     }
-    return await unlink(filePath).then(() =>
-      this.deleteEvalRunFromIndex(evalRunId)
-    );
   }
 
-  private generateFileName(evalRunId: string): string {
+  private resolveEvalFilename(evalRunId: string): string {
     return `${evalRunId}.json`;
   }
 
-  private getIndexFilePath(): string {
-    return path.resolve(this.storeRoot, 'index.txt');
+  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 parseLineToKey(key: string): EvalRunKey {
-    return EvalRunKeySchema.parse(JSON.parse(key));
-  }
+  private async processLineByLine(filePath: string) {
+    const fileStream = fs.createReadStream(filePath);
+    const keys: Record<string, EvalRunKey> = {};
 
-  private generateRootPath(): string {
-    return path.resolve(process.cwd(), `.genkit/evals`);
-  }
-
-  private async deleteEvalRunFromIndex(evalRunId: string): Promise<void> {
-    const entries = [];
-    const fileStream = fs.createReadStream(this.getIndexFilePath());
-    const rl = readline.createInterface({
+    const rl = createInterface({
       input: fileStream,
+      crlfDelay: Infinity,
     });
-
     for await (const line of rl) {
-      const entry = EvalRunKeySchema.parse(JSON.parse(line));
-      if (entry.evalRunId !== evalRunId) {
-        entries.push(line);
+      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;
+  }
 
-    await writeFile(
-      this.getIndexFilePath(),
-      // end with delimiter to parse correctly
-      entries.join(this.INDEX_DELIMITER) + this.INDEX_DELIMITER
+  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)
     );
   }
 }

genkit-tools/common/src/server/router.ts

@@ -169,7 +169,8 @@ export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) =>
       .input(apis.ListEvalKeysRequestSchema)
       .output(apis.ListEvalKeysResponseSchema)
       .query(async ({ input }) => {
-        const response = await getEvalStore().list(input);
+        const store = await getEvalStore();
+        const response = await store.list(input);
         return {
           evalRunKeys: response.evalRunKeys,
         };
@@ -182,7 +183,8 @@ export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) =>
       .query(async ({ input }) => {
         const parts = input.name.split('/');
         const evalRunId = parts[1];
-        const evalRun = await getEvalStore().load(evalRunId);
+        const store = await getEvalStore();
+        const evalRun = await store.load(evalRunId);
         if (!evalRun) {
           throw new TRPCError({
             code: 'NOT_FOUND',
@@ -198,7 +200,8 @@ export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) =>
       .mutation(async ({ input }) => {
         const parts = input.name.split('/');
         const evalRunId = parts[1];
-        await getEvalStore().delete(evalRunId);
+        const store = await getEvalStore();
+        await store.delete(evalRunId);
       }),
 
     /** Retrieves all eval datasets */