feat(evals): Add datasetType; export model input schema

genkit-tools/common/package.json

@@ -24,6 +24,7 @@
     "inquirer": "^8.2.0",
     "js-yaml": "^4.1.0",
     "json-2-csv": "^5.5.1",
+    "json-schema": "^0.4.0",
     "terminate": "^2.6.1",
     "tsx": "^4.19.2",
     "uuid": "^9.0.1",
@@ -33,22 +34,23 @@
     "zod-to-json-schema": "^3.22.4"
   },
   "devDependencies": {
+    "@jest/globals": "^29.7.0",
     "@types/adm-zip": "^0.5.5",
     "@types/body-parser": "^1.19.5",
     "@types/cli-color": "^2.0.6",
     "@types/configstore": "^6.0.2",
     "@types/express": "^4.17.21",
     "@types/inquirer": "^8.1.3",
     "@types/jest": "^29.5.12",
-    "@jest/globals": "^29.7.0",
-    "jest": "^29.7.0",
-    "ts-jest": "^29.1.2",
     "@types/js-yaml": "^4.0.9",
+    "@types/json-schema": "^7.0.15",
     "@types/node": "^20.11.19",
     "@types/uuid": "^9.0.8",
     "genversion": "^3.2.0",
+    "jest": "^29.7.0",
     "npm-run-all": "^4.1.5",
     "rimraf": "^6.0.1",
+    "ts-jest": "^29.1.2",
     "ts-node": "^10.9.2",
     "typescript": "^5.3.3"
   },

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

@@ -92,6 +92,7 @@ export class LocalFileDatasetStore implements DatasetStore {
       targetAction,
       size: dataset.length,
       version: 1,
+      datasetType: req.datasetType,
       createTime: now,
       updateTime: now,
     };
@@ -136,6 +137,7 @@ export class LocalFileDatasetStore implements DatasetStore {
       schema: schema ? schema : prevMetadata.schema,
       targetAction: targetAction ? targetAction : prevMetadata.targetAction,
       version: data ? prevMetadata.version + 1 : prevMetadata.version,
+      datasetType: prevMetadata.datasetType,
       createTime: prevMetadata.createTime,
       updateTime: now,
     };

genkit-tools/common/src/types/apis.ts

@@ -17,6 +17,7 @@
 import { z } from 'zod';
 import {
   DatasetSchemaSchema,
+  DatasetTypeSchema,
   EvalInferenceInputSchema,
   EvalRunKeySchema,
 } from './eval';
@@ -106,6 +107,7 @@ export type GetEvalRunRequest = z.infer<typeof GetEvalRunRequestSchema>;
 export const CreateDatasetRequestSchema = z.object({
   data: EvalInferenceInputSchema,
   datasetId: z.string().optional(),
+  datasetType: DatasetTypeSchema,
   schema: DatasetSchemaSchema.optional(),
   targetAction: z.string().optional(),
 });

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

@@ -14,13 +14,16 @@
  * limitations under the License.
  */
 
+import { JSONSchema7 } from 'json-schema';
 import { z } from 'zod';
+import zodToJsonSchema from 'zod-to-json-schema';
 import {
   CreateDatasetRequest,
   ListEvalKeysRequest,
   ListEvalKeysResponse,
   UpdateDatasetRequest,
 } from './apis';
+import { GenerateRequestSchema } from './model';
 
 /**
  * This file defines schema and types that are used by the Eval store.
@@ -42,6 +45,21 @@ export type EvalInferenceStructuredInput = z.infer<
   typeof EvalInferenceStructuredInputSchema
 >;
 
+/**
+ * Supported datatype when running eval-inference using models
+ */
+export const ModelInferenceInputSchema = z.union([
+  z.string(),
+  GenerateRequestSchema,
+]);
+export type ModelInferenceInput = z.infer<typeof ModelInferenceInputSchema>;
+export const ModelInferenceInputJSONSchema = zodToJsonSchema(
+  ModelInferenceInputSchema,
+  {
+    $refStrategy: 'none',
+    removeAdditionalStrategy: 'strict',
+  }
+) as JSONSchema7;
 /**
  * A set of samples that is ready for inference.
  *
@@ -172,6 +190,10 @@ export const DatasetSchemaSchema = z.object({
     .optional(),
 });
 
+/** Type of dataset, useful for UI niceties. */
+export const DatasetTypeSchema = z.enum(['UNKNOWN', 'FLOW', 'MODEL']);
+export type DatasetType = z.infer<typeof DatasetTypeSchema>;
+
 /**
  * Metadata for Dataset objects containing version, create and update time, etc.
  */
@@ -180,6 +202,7 @@ export const DatasetMetadataSchema = z.object({
   datasetId: z.string(),
   size: z.number(),
   schema: DatasetSchemaSchema.optional(),
+  datasetType: DatasetTypeSchema,
   targetAction: z.string().optional(),
   /** 1 for v1, 2 for v2, etc */
   version: z.number(),

genkit-tools/common/tests/eval/localFileDatasetStore_test.ts

@@ -74,23 +74,27 @@ const SAMPLE_DATASET_METADATA_1_V1 = {
   datasetId: SAMPLE_DATASET_ID_1,
   size: 2,
   version: 1,
+  datasetType: 'UNKNOWN',
   createTime: FAKE_TIME.toString(),
   updateTime: FAKE_TIME.toString(),
 };
 const SAMPLE_DATASET_METADATA_1_V2 = {
   datasetId: SAMPLE_DATASET_ID_1,
   size: 3,
   version: 2,
+  datasetType: 'UNKNOWN',
   createTime: FAKE_TIME.toString(),
   updateTime: FAKE_TIME.toString(),
 };
 
 const CREATE_DATASET_REQUEST = CreateDatasetRequestSchema.parse({
   data: { samples: SAMPLE_DATASET_1_V1 },
+  datasetType: 'UNKNOWN',
 });
 
 const CREATE_DATASET_REQUEST_WITH_SCHEMA = CreateDatasetRequestSchema.parse({
   data: { samples: SAMPLE_DATASET_1_V1 },
+  datasetType: 'UNKNOWN',
   schema: {
     inputSchema: {
       type: 'string',