diff --git a/genkit-tools/common/src/eval/index.ts b/genkit-tools/common/src/eval/index.ts index 28f2374b9..51b157ae3 100644 --- a/genkit-tools/common/src/eval/index.ts +++ b/genkit-tools/common/src/eval/index.ts @@ -14,7 +14,8 @@ * limitations under the License. */ -import { EvalStore } from '../types/eval'; +import { DatasetStore, EvalStore } from '../types/eval'; +import { LocalFileDatasetStore } from './localFileDatasetStore'; import { LocalFileEvalStore } from './localFileEvalStore'; export { EvalFlowInput, EvalFlowInputSchema } from '../types/eval'; export * from './exporter'; @@ -24,3 +25,7 @@ export function getEvalStore(): EvalStore { // TODO: This should provide EvalStore, based on tools config. return LocalFileEvalStore.getEvalStore(); } + +export function getDatasetStore(): DatasetStore { + return LocalFileDatasetStore.getDatasetStore(); +} diff --git a/genkit-tools/common/src/eval/localFileDatasetStore.ts b/genkit-tools/common/src/eval/localFileDatasetStore.ts new file mode 100644 index 000000000..c5a1948f2 --- /dev/null +++ b/genkit-tools/common/src/eval/localFileDatasetStore.ts @@ -0,0 +1,222 @@ +/** + * 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 crypto from 'crypto'; +import fs from 'fs'; +import { readFile, rm, writeFile } from 'fs/promises'; +import path from 'path'; +import { v4 as uuidv4 } from 'uuid'; +import { CreateDatasetRequest, UpdateDatasetRequest } from '../types/apis'; +import { + Dataset, + DatasetMetadata, + DatasetStore, + EvalFlowInputSchema, +} from '../types/eval'; +import { logger } from '../utils/logger'; + +/** + * A local, file-based DatasetStore implementation. + */ +export class LocalFileDatasetStore implements DatasetStore { + private readonly storeRoot; + private readonly indexFile; + private readonly INDEX_DELIMITER = '\n'; + private static cachedDatasetStore: LocalFileDatasetStore | null = null; + + private constructor(storeRoot: string) { + this.storeRoot = storeRoot; + this.indexFile = this.getIndexFilePath(); + fs.mkdirSync(this.storeRoot, { recursive: true }); + if (!fs.existsSync(this.indexFile)) { + fs.writeFileSync(path.resolve(this.indexFile), ''); + } + logger.info( + `Initialized local file dataset store at root: ${this.storeRoot}` + ); + } + + static getDatasetStore() { + if (!this.cachedDatasetStore) { + this.cachedDatasetStore = new LocalFileDatasetStore( + this.generateRootPath() + ); + } + return this.cachedDatasetStore; + } + + static reset() { + this.cachedDatasetStore = null; + } + + async createDataset(req: CreateDatasetRequest): Promise { + return this.createDatasetInternal(req.data, req.displayName); + } + + private async createDatasetInternal( + data: Dataset, + displayName?: string + ): Promise { + const datasetId = this.generateDatasetId(); + const filePath = path.resolve( + this.storeRoot, + this.generateFileName(datasetId) + ); + + if (fs.existsSync(filePath)) { + logger.error(`Dataset already exists at ` + filePath); + throw new Error( + `Create dataset failed: file already exists at {$filePath}` + ); + } + + logger.info(`Saving Dataset to ` + filePath); + await writeFile(filePath, JSON.stringify(data)); + + const now = new Date().toString(); + const metadata = { + datasetId, + size: Array.isArray(data) ? data.length : data.samples.length, + version: 1, + displayName: displayName, + createTime: now, + updateTime: now, + }; + + let metadataMap = await this.getMetadataMap(); + metadataMap[datasetId] = metadata; + + logger.debug( + `Saving DatasetMetadata for ID ${datasetId} to ` + + path.resolve(this.indexFile) + ); + + await writeFile(path.resolve(this.indexFile), JSON.stringify(metadataMap)); + return metadata; + } + + async updateDataset(req: UpdateDatasetRequest): Promise { + const datasetId = req.datasetId; + const filePath = path.resolve( + this.storeRoot, + this.generateFileName(datasetId) + ); + if (!fs.existsSync(filePath)) { + throw new Error(`Update dataset failed: dataset not found`); + } + + let metadataMap = await this.getMetadataMap(); + const prevMetadata = metadataMap[datasetId]; + if (!prevMetadata) { + throw new Error(`Update dataset failed: dataset metadata not found`); + } + + logger.info(`Updating Dataset at ` + filePath); + await writeFile(filePath, JSON.stringify(req.patch)); + + const now = new Date().toString(); + const newMetadata = { + datasetId: datasetId, + size: Array.isArray(req.patch) + ? req.patch.length + : req.patch.samples.length, + version: prevMetadata.version + 1, + displayName: req.displayName, + createTime: prevMetadata.createTime, + updateTime: now, + }; + + logger.debug( + `Updating DatasetMetadata for ID ${datasetId} at ` + + path.resolve(this.indexFile) + ); + // Replace the metadata object in the metadata map + metadataMap[datasetId] = newMetadata; + await writeFile(path.resolve(this.indexFile), JSON.stringify(metadataMap)); + + return newMetadata; + } + + async getDataset(datasetId: string): Promise { + const filePath = path.resolve( + this.storeRoot, + this.generateFileName(datasetId) + ); + if (!fs.existsSync(filePath)) { + throw new Error(`Dataset not found for dataset ID {$id}`); + } + return await readFile(filePath, 'utf8').then((data) => + EvalFlowInputSchema.parse(JSON.parse(data)) + ); + } + + async listDatasets(): Promise { + return this.getMetadataMap().then((metadataMap) => { + let metadatas = []; + + for (var key in metadataMap) { + metadatas.push(metadataMap[key]); + } + return metadatas; + }); + } + + async deleteDataset(datasetId: string): Promise { + const filePath = path.resolve( + this.storeRoot, + this.generateFileName(datasetId) + ); + await rm(filePath); + + let metadataMap = await this.getMetadataMap(); + delete metadataMap[datasetId]; + + logger.debug( + `Deleting DatasetMetadata for ID ${datasetId} in ` + + path.resolve(this.indexFile) + ); + await writeFile(path.resolve(this.indexFile), JSON.stringify(metadataMap)); + } + + private static generateRootPath(): string { + const rootHash = crypto + .createHash('md5') + .update(process.cwd() || 'unknown') + .digest('hex'); + return path.resolve(process.cwd(), `.genkit/${rootHash}/datasets`); + } + + private generateDatasetId(): string { + return uuidv4(); + } + + private generateFileName(datasetId: string): string { + return `${datasetId}.json`; + } + + private getIndexFilePath(): string { + return path.resolve(this.storeRoot, 'index.json'); + } + + private async getMetadataMap(): Promise { + if (!fs.existsSync(this.indexFile)) { + return Promise.resolve({} as any); + } + return await readFile(path.resolve(this.indexFile), 'utf8').then((data) => + JSON.parse(data) + ); + } +} diff --git a/genkit-tools/common/src/server/router.ts b/genkit-tools/common/src/server/router.ts index 5f1859db8..65f85c6f1 100644 --- a/genkit-tools/common/src/server/router.ts +++ b/genkit-tools/common/src/server/router.ts @@ -14,7 +14,8 @@ * limitations under the License. */ import { initTRPC, TRPCError } from '@trpc/server'; -import { getEvalStore } from '../eval'; +import { z } from 'zod'; +import { getDatasetStore, getEvalStore } from '../eval'; import { Runner } from '../runner/runner'; import { GenkitToolsError } from '../runner/types'; import { Action } from '../types/action'; @@ -202,6 +203,51 @@ export const TOOLS_SERVER_ROUTER = (runner: Runner) => return evalRun; }), + /** Retrieves all eval datasets */ + listDatasets: loggedProcedure + .input(z.void()) + .output(z.array(evals.DatasetMetadataSchema)) + .query(async () => { + const response = await getDatasetStore().listDatasets(); + return response; + }), + + /** Retrieves an existing dataset */ + getDataset: loggedProcedure + .input(z.string()) + .output(evals.EvalFlowInputSchema) + .query(async ({ input }) => { + const response = await getDatasetStore().getDataset(input); + return response; + }), + + /** Creates a new dataset */ + createDataset: loggedProcedure + .input(apis.CreateDatasetRequestSchema) + .output(evals.DatasetMetadataSchema) + .query(async ({ input }) => { + const response = await getDatasetStore().createDataset(input); + return response; + }), + + /** Updates an exsting dataset */ + updateDataset: loggedProcedure + .input(apis.UpdateDatasetRequestSchema) + .output(evals.DatasetMetadataSchema) + .query(async ({ input }) => { + const response = await getDatasetStore().updateDataset(input); + return response; + }), + + /** Deletes an exsting dataset */ + deleteDataset: loggedProcedure + .input(z.string()) + .output(z.void()) + .query(async ({ input }) => { + const response = await getDatasetStore().deleteDataset(input); + return response; + }), + /** Send a screen view analytics event */ sendPageView: t.procedure .input(apis.PageViewSchema) diff --git a/genkit-tools/common/src/types/apis.ts b/genkit-tools/common/src/types/apis.ts index 6ed3e4961..e355a0841 100644 --- a/genkit-tools/common/src/types/apis.ts +++ b/genkit-tools/common/src/types/apis.ts @@ -15,7 +15,7 @@ */ import { z } from 'zod'; -import { EvalRunKeySchema } from './eval'; +import { EvalFlowInputSchema, EvalRunKeySchema } from './eval'; import { FlowStateSchema } from './flow'; import { GenerationCommonConfigSchema, @@ -132,3 +132,18 @@ export const GetEvalRunRequestSchema = z.object({ name: z.string(), }); export type GetEvalRunRequest = z.infer; + +export const CreateDatasetRequestSchema = z.object({ + data: EvalFlowInputSchema, + displayName: z.string().optional(), +}); + +export type CreateDatasetRequest = z.infer; + +export const UpdateDatasetRequestSchema = z.object({ + /** Supports upsert */ + patch: EvalFlowInputSchema, + datasetId: z.string(), + displayName: z.string().optional(), +}); +export type UpdateDatasetRequest = z.infer; diff --git a/genkit-tools/common/src/types/eval.ts b/genkit-tools/common/src/types/eval.ts index d9631d13e..c43fa4fc7 100644 --- a/genkit-tools/common/src/types/eval.ts +++ b/genkit-tools/common/src/types/eval.ts @@ -15,7 +15,12 @@ */ import { z } from 'zod'; -import { ListEvalKeysRequest, ListEvalKeysResponse } from './apis'; +import { + CreateDatasetRequest, + ListEvalKeysRequest, + ListEvalKeysResponse, + UpdateDatasetRequest, +} from './apis'; /** * This file defines schema and types that are used by the Eval store. @@ -145,25 +150,17 @@ export interface EvalStore { /** * Metadata for Dataset objects containing version, create and update time, etc. */ -export const DatasetMetadaSchema = z.object({ +export const DatasetMetadataSchema = z.object({ /** autogenerated */ datasetId: z.string(), size: z.number(), - uri: z.string(), /** 1 for v1, 2 for v2, etc */ version: z.number(), displayName: z.string().optional(), createTime: z.string(), updateTime: z.string(), }); -export type DatasetMetadata = z.infer; - -export const UpdateDatasetRequestSchema = z.object({ - /** Supports upsert */ - patch: EvalFlowInputSchema, - displayName: z.string().optional(), -}); -export type UpdateDatasetRequest = z.infer; +export type DatasetMetadata = z.infer; /** * Eval dataset store persistence interface. @@ -171,10 +168,10 @@ export type UpdateDatasetRequest = z.infer; export interface DatasetStore { /** * Create new dataset with the given data - * @param data data containing eval flow inputs + * @param req create requeest with the data * @returns dataset metadata */ - createDataset(data: Dataset): Promise; + createDataset(req: CreateDatasetRequest): Promise; /** * Update dataset @@ -185,10 +182,10 @@ export interface DatasetStore { /** * Get existing dataset - * @param id the ID of the dataset + * @param datasetId the ID of the dataset * @returns dataset ready for inference */ - getDataset(id: string): Promise; + getDataset(datasetId: string): Promise; /** * List all existing datasets @@ -198,7 +195,7 @@ export interface DatasetStore { /** * Delete existing dataset - * @param id the ID of the dataset + * @param datasetId the ID of the dataset */ - deleteDataset(id: string): Promise; + deleteDataset(datasetId: string): Promise; } diff --git a/genkit-tools/common/src/types/model.ts b/genkit-tools/common/src/types/model.ts index 3bfc2f444..efc0f2db4 100644 --- a/genkit-tools/common/src/types/model.ts +++ b/genkit-tools/common/src/types/model.ts @@ -14,7 +14,7 @@ * limitations under the License. */ import { z } from 'zod'; -import { DocumentDataSchema } from './document.js'; +import { DocumentDataSchema } from './document'; // // IMPORTANT: Keep this file in sync with genkit/ai/src/model.ts! diff --git a/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts b/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts new file mode 100644 index 000000000..d3f1691fa --- /dev/null +++ b/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts @@ -0,0 +1,290 @@ +/** + * 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 { + afterEach, + beforeEach, + describe, + expect, + it, + jest, +} from '@jest/globals'; +import fs from 'fs'; +import { LocalFileDatasetStore } from '../../src/eval/localFileDatasetStore'; +import { + CreateDatasetRequestSchema, + UpdateDatasetRequestSchema, +} from '../../src/types/apis'; +import { DatasetStore } from '../../src/types/eval'; + +const FAKE_TIME = new Date('2024-02-03T12:05:33.243Z'); + +const SAMPLE_DATASET_1_V1 = { + samples: [ + { + input: 'Cats are evil', + reference: 'Sorry no reference', + }, + { + input: 'Dogs are beautiful', + }, + ], +}; + +const SAMPLE_DATASET_1_V2 = { + samples: [ + { + input: 'Cats are evil', + reference: 'Sorry no reference', + }, + { + input: 'Dogs are angels', + }, + { + input: 'Dogs are also super cute', + }, + ], +}; + +const SAMPLE_DATASET_ID_1 = '12345678'; +const SAMPLE_DATASET_NAME_1 = 'dataset-1'; + +const SAMPLE_DATASET_METADATA_1_V1 = { + datasetId: SAMPLE_DATASET_ID_1, + size: 2, + version: 1, + displayName: SAMPLE_DATASET_NAME_1, + createTime: FAKE_TIME.toString(), + updateTime: FAKE_TIME.toString(), +}; +const SAMPLE_DATASET_METADATA_1_V2 = { + datasetId: SAMPLE_DATASET_ID_1, + size: 3, + version: 2, + displayName: SAMPLE_DATASET_NAME_1, + createTime: FAKE_TIME.toString(), + updateTime: FAKE_TIME.toString(), +}; + +const CREATE_DATASET_REQUEST = CreateDatasetRequestSchema.parse({ + data: SAMPLE_DATASET_1_V1, + displayName: SAMPLE_DATASET_NAME_1, +}); + +const UPDATE_DATASET_REQUEST = UpdateDatasetRequestSchema.parse({ + patch: SAMPLE_DATASET_1_V2, + datasetId: SAMPLE_DATASET_ID_1, + displayName: SAMPLE_DATASET_NAME_1, +}); + +const SAMPLE_DATASET_ID_2 = '22345678'; +const SAMPLE_DATASET_NAME_2 = 'dataset-2'; + +const SAMPLE_DATASET_METADATA_2 = { + datasetId: SAMPLE_DATASET_ID_2, + size: 5, + version: 1, + displayName: SAMPLE_DATASET_NAME_2, + createTime: FAKE_TIME.toString(), + updateTime: FAKE_TIME.toString(), +}; + +jest.mock('crypto', () => { + return { + createHash: jest.fn().mockReturnThis(), + update: jest.fn().mockReturnThis(), + digest: jest.fn(() => 'store-root'), + }; +}); + +jest.mock('uuid', () => ({ + v4: () => SAMPLE_DATASET_ID_1, +})); + +jest.useFakeTimers({ advanceTimers: true }); +jest.setSystemTime(FAKE_TIME); + +describe('localFileDatasetStore', () => { + let DatasetStore: DatasetStore; + + beforeEach(() => { + LocalFileDatasetStore.reset(); + DatasetStore = LocalFileDatasetStore.getDatasetStore() as DatasetStore; + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + describe('createDataset', () => { + it('writes and updates index for new dataset', async () => { + fs.promises.writeFile = jest.fn(async () => Promise.resolve(undefined)); + fs.promises.appendFile = jest.fn(async () => Promise.resolve(undefined)); + // For index file reads + fs.promises.readFile = jest.fn(async () => + Promise.resolve(JSON.stringify({}) as any) + ); + + const datasetMetadata = await DatasetStore.createDataset( + CREATE_DATASET_REQUEST + ); + + expect(fs.promises.writeFile).toHaveBeenCalledTimes(2); + expect(fs.promises.writeFile).toHaveBeenNthCalledWith( + 1, + expect.stringContaining('datasets/12345678.json'), + JSON.stringify(CREATE_DATASET_REQUEST.data) + ); + const metadataMap = { + [SAMPLE_DATASET_ID_1]: SAMPLE_DATASET_METADATA_1_V1, + }; + expect(fs.promises.writeFile).toHaveBeenNthCalledWith( + 2, + expect.stringContaining('datasets/index.json'), + JSON.stringify(metadataMap) + ); + expect(datasetMetadata).toMatchObject(SAMPLE_DATASET_METADATA_1_V1); + }); + + it('fails request if dataset already exists', async () => { + fs.existsSync = jest.fn(() => true); + + expect(async () => { + await DatasetStore.createDataset(CREATE_DATASET_REQUEST); + }).rejects.toThrow(); + + expect(fs.promises.writeFile).toBeCalledTimes(0); + }); + }); + + describe('updateDataset', () => { + it('succeeds for existing dataset', async () => { + fs.existsSync = jest.fn(() => true); + let metadataMap = { + [SAMPLE_DATASET_ID_1]: SAMPLE_DATASET_METADATA_1_V1, + [SAMPLE_DATASET_ID_2]: SAMPLE_DATASET_METADATA_2, + }; + // For index file reads + fs.promises.readFile = jest.fn(async () => + Promise.resolve(JSON.stringify(metadataMap) as any) + ); + fs.promises.writeFile = jest.fn(async () => Promise.resolve(undefined)); + fs.promises.appendFile = jest.fn(async () => Promise.resolve(undefined)); + + const datasetMetadata = await DatasetStore.updateDataset( + UPDATE_DATASET_REQUEST + ); + + expect(fs.promises.writeFile).toHaveBeenCalledTimes(2); + expect(fs.promises.writeFile).toHaveBeenNthCalledWith( + 1, + expect.stringContaining('datasets/12345678.json'), + JSON.stringify(SAMPLE_DATASET_1_V2) + ); + const updatedMetadataMap = { + [SAMPLE_DATASET_ID_1]: SAMPLE_DATASET_METADATA_1_V2, + [SAMPLE_DATASET_ID_2]: SAMPLE_DATASET_METADATA_2, + }; + expect(fs.promises.writeFile).toHaveBeenNthCalledWith( + 2, + expect.stringContaining('datasets/index.json'), + JSON.stringify(updatedMetadataMap) + ); + expect(datasetMetadata).toMatchObject(SAMPLE_DATASET_METADATA_1_V2); + }); + + it('fails for non existing dataset', async () => { + fs.existsSync = jest.fn(() => false); + + expect(async () => { + await DatasetStore.updateDataset(UPDATE_DATASET_REQUEST); + }).rejects.toThrow(); + + expect(fs.promises.writeFile).toBeCalledTimes(0); + }); + }); + + describe('listDatasets', () => { + it('succeeds for zero datasets', async () => { + fs.existsSync = jest.fn(() => false); + + const metadatas = await DatasetStore.listDatasets(); + + expect(metadatas).toMatchObject([]); + }); + + it('succeeds for existing datasets', async () => { + fs.existsSync = jest.fn(() => true); + const metadataMap = { + [SAMPLE_DATASET_ID_1]: SAMPLE_DATASET_METADATA_1_V1, + [SAMPLE_DATASET_ID_2]: SAMPLE_DATASET_METADATA_2, + }; + fs.promises.readFile = jest.fn(async () => + Promise.resolve(JSON.stringify(metadataMap) as any) + ); + + const metadatas = await DatasetStore.listDatasets(); + + expect(metadatas).toMatchObject([ + SAMPLE_DATASET_METADATA_1_V1, + SAMPLE_DATASET_METADATA_2, + ]); + }); + }); + + describe('getDataset', () => { + it('succeeds for existing dataset', async () => { + fs.existsSync = jest.fn(() => true); + fs.promises.readFile = jest.fn(async () => + Promise.resolve(JSON.stringify(SAMPLE_DATASET_1_V1) as any) + ); + + const fetchedDataset = await DatasetStore.getDataset(SAMPLE_DATASET_ID_1); + + expect(fetchedDataset).toMatchObject(SAMPLE_DATASET_1_V1); + }); + + it('fails for non existing dataset', async () => { + // TODO: Implement this. + }); + }); + + describe('deleteDataset', () => { + it('deletes dataset and updates index', async () => { + fs.promises.rm = jest.fn(async () => Promise.resolve()); + let metadataMap = { + [SAMPLE_DATASET_ID_1]: SAMPLE_DATASET_METADATA_1_V1, + [SAMPLE_DATASET_ID_2]: SAMPLE_DATASET_METADATA_2, + }; + fs.promises.readFile = jest.fn(async () => + Promise.resolve(JSON.stringify(metadataMap) as any) + ); + + await DatasetStore.deleteDataset(SAMPLE_DATASET_ID_1); + + expect(fs.promises.rm).toHaveBeenCalledWith( + expect.stringContaining('datasets/12345678.json') + ); + let updatedMetadataMap = { + [SAMPLE_DATASET_ID_2]: SAMPLE_DATASET_METADATA_2, + }; + expect(fs.promises.writeFile).toHaveBeenCalledWith( + expect.stringContaining('datasets/index.json'), + JSON.stringify(updatedMetadataMap) + ); + }); + }); +});