diff --git a/genkit-tools/common/src/eval/localFileDatasetStore.ts b/genkit-tools/common/src/eval/localFileDatasetStore.ts
index 79f160873..44b5ae90d 100644
--- a/genkit-tools/common/src/eval/localFileDatasetStore.ts
+++ b/genkit-tools/common/src/eval/localFileDatasetStore.ts
@@ -72,6 +72,8 @@ export class LocalFileDatasetStore implements DatasetStore {
     req: CreateDatasetRequest
   ): Promise<DatasetMetadata> {
     const { data, datasetId, schema, targetAction } = req;
+
+    console.log("Schema set to", schema);
     const id = await this.generateDatasetId(datasetId);
     const filePath = path.resolve(this.storeRoot, this.generateFileName(id));
 
@@ -96,6 +98,7 @@ export class LocalFileDatasetStore implements DatasetStore {
       updateTime: now,
     };
 
+    console.log("Dataset metadata set to", schema);
     let metadataMap = await this.getMetadataMap();
     metadataMap[id] = metadata;
 
diff --git a/genkit-tools/common/src/types/eval.ts b/genkit-tools/common/src/types/eval.ts
index bcf3ff1fb..bf804e4b7 100644
--- a/genkit-tools/common/src/types/eval.ts
+++ b/genkit-tools/common/src/types/eval.ts
@@ -15,6 +15,7 @@
  */
 
 import { z } from 'zod';
+import { CustomAnySchema, JSONSchema7Schema } from './action';
 import {
   CreateDatasetRequest,
   ListEvalKeysRequest,
@@ -160,16 +161,10 @@ export interface EvalStore {
   list(query?: ListEvalKeysRequest): Promise<ListEvalKeysResponse>;
 }
 
-export const DatasetSchemaSchema = z.object({
-  inputSchema: z
-    .record(z.any())
-    .describe('Valid JSON Schema for the `input` field of dataset entry.')
-    .optional(),
-  referenceSchema: z
-    .record(z.any())
-    .describe('Valid JSON Schema for the `reference` field of dataset entry.')
-    .optional(),
-});
+export const DatasetSchemaSchema = z
+  .record(z.string(), CustomAnySchema)
+  .describe('Schema objects related to the dataset.')
+  .nullish();
 
 /**
  * Metadata for Dataset objects containing version, create and update time, etc.
diff --git a/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts b/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts
index c658208c8..44a376497 100644
--- a/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts
+++ b/genkit-tools/common/tests/eval/localFileDatasetStore_test.ts
@@ -89,17 +89,30 @@ const CREATE_DATASET_REQUEST = CreateDatasetRequestSchema.parse({
   data: { samples: SAMPLE_DATASET_1_V1 },
 });
 
+const INPUT_SCHEMA = {
+  "type": "string",
+  "$schema": "http://json-schema.org/draft-07/schema#"
+};
+
+const REFERENCE_SCHEMA = {
+  "type": "object",
+  "properties": {
+    "output": {
+      "type": "string"
+    },
+  },
+  "required": [
+    "output",
+  ],
+  "additionalProperties": true,
+  "$schema": "http://json-schema.org/draft-07/schema#"
+};
+
 const CREATE_DATASET_REQUEST_WITH_SCHEMA = CreateDatasetRequestSchema.parse({
   data: { samples: SAMPLE_DATASET_1_V1 },
   schema: {
-    inputSchema: {
-      type: 'string',
-      $schema: 'http://json-schema.org/draft-07/schema#',
-    },
-    referenceSchema: {
-      type: 'number',
-      $schema: 'http://json-schema.org/draft-07/schema#',
-    },
+    "inputSchema": INPUT_SCHEMA,
+    "referenceSchema": REFERENCE_SCHEMA,
   },
   targetAction: '/flow/my-flow',
 });
@@ -235,20 +248,16 @@ describe('localFileDatasetStore', () => {
         ...s,
       }));
 
+      console.log("CREATE_DATASET_REQUEST_WITH_SCHEMA schema", CREATE_DATASET_REQUEST_WITH_SCHEMA.schema);
+
       const datasetMetadata = await DatasetStore.createDataset({
         ...CREATE_DATASET_REQUEST_WITH_SCHEMA,
         datasetId: SAMPLE_DATASET_ID_1,
       });
 
       expect(datasetMetadata.schema).toMatchObject({
-        inputSchema: {
-          type: 'string',
-          $schema: 'http://json-schema.org/draft-07/schema#',
-        },
-        referenceSchema: {
-          type: 'number',
-          $schema: 'http://json-schema.org/draft-07/schema#',
-        },
+        inputSchema: INPUT_SCHEMA,
+        referenceSchema: REFERENCE_SCHEMA,
       });
       expect(datasetMetadata.targetAction).toEqual('/flow/my-flow');
     });
@@ -415,19 +424,15 @@ describe('localFileDatasetStore', () => {
       const datasetMetadata = await DatasetStore.updateDataset({
         datasetId: SAMPLE_DATASET_ID_1,
         schema: {
-          inputSchema: {
-            type: 'string',
-            $schema: 'http://json-schema.org/draft-07/schema#',
-          },
+          inputSchema: INPUT_SCHEMA,
+          referenceSchema: REFERENCE_SCHEMA,
         },
         targetAction: '/flow/my-flow-2',
       });
 
       expect(datasetMetadata.schema).toMatchObject({
-        inputSchema: {
-          type: 'string',
-          $schema: 'http://json-schema.org/draft-07/schema#',
-        },
+        inputSchema: INPUT_SCHEMA,
+        referenceSchema: REFERENCE_SCHEMA,
       });
       expect(datasetMetadata.targetAction).toEqual('/flow/my-flow-2');
     });