feat(compass-collection): LLM Output Validation - Mock Data Generator CLOUDP-333855

Author: ncarbon

packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx

@@ -16,37 +16,42 @@ import type { CollectionState } from '../../modules/collection-tab';
 import { default as collectionTabReducer } from '../../modules/collection-tab';
 import type { ConnectionInfo } from '@mongodb-js/connection-info';
 import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
+import type { SchemaAnalysisState } from '../../schema-analysis-types';
+
+const defaultSchemaAnalysisState: SchemaAnalysisState = {
+  status: 'complete',
+  processedSchema: {
+    name: {
+      type: 'String',
+      probability: 1.0,
+      sample_values: ['John', 'Jane'],
+    },
+  },
+  sampleDocument: { name: 'John' },
+  schemaMetadata: { maxNestingDepth: 1, validationRules: null },
+};
 
 describe('MockDataGeneratorModal', () => {
   async function renderModal({
     isOpen = true,
     currentStep = MockDataGeneratorStep.SCHEMA_CONFIRMATION,
     enableGenAISampleDocumentPassing = false,
     mockServices = createMockServices(),
+    schemaAnalysis = defaultSchemaAnalysisState,
     connectionInfo,
   }: {
     isOpen?: boolean;
     enableGenAISampleDocumentPassing?: boolean;
     currentStep?: MockDataGeneratorStep;
     mockServices?: any;
     connectionInfo?: ConnectionInfo;
+    schemaAnalysis?: SchemaAnalysisState;
   } = {}) {
     const initialState: CollectionState = {
       workspaceTabId: 'test-workspace-tab-id',
       namespace: 'test.collection',
       metadata: null,
-      schemaAnalysis: {
-        status: 'complete',
-        processedSchema: {
-          name: {
-            type: 'String',
-            probability: 1.0,
-            sample_values: ['John', 'Jane'],
-          },
-        },
-        sampleDocument: { name: 'John' },
-        schemaMetadata: { maxNestingDepth: 1, validationRules: null },
-      },
+      schemaAnalysis,
       fakerSchemaGeneration: {
         status: 'idle',
       },
@@ -284,6 +289,33 @@ describe('MockDataGeneratorModal', () => {
   });
 
   describe('on the schema editor step', () => {
+    const mockSchemaAnalysis: SchemaAnalysisState = {
+      ...defaultSchemaAnalysisState,
+      processedSchema: {
+        name: {
+          type: 'String',
+          probability: 1.0,
+        },
+        age: {
+          type: 'Int32',
+          probability: 1.0,
+        },
+        email: {
+          type: 'String',
+          probability: 1.0,
+        },
+        username: {
+          type: 'String',
+          probability: 1.0,
+        },
+      },
+      sampleDocument: {
+        name: 'Jane',
+        age: 99,
+        email: '[email protected]',
+        username: 'JaneDoe123',
+      },
+    };
     const mockServicesWithMockDataResponse = createMockServices();
     mockServicesWithMockDataResponse.atlasAiService.getMockDataSchema = () =>
       Promise.resolve({
@@ -336,19 +368,22 @@ describe('MockDataGeneratorModal', () => {
                   fields: [],
                 },
               }),
-            1000
+            1
           )
         );
 
-      await renderModal();
+      await renderModal({ mockServices });
 
       // advance to the schema editor step
       userEvent.click(screen.getByText('Confirm'));
       expect(screen.getByTestId('faker-schema-editor-loader')).to.exist;
     });
 
     it('shows the faker schema editor when the faker schema generation is completed', async () => {
-      await renderModal({ mockServices: mockServicesWithMockDataResponse });
+      await renderModal({
+        mockServices: mockServicesWithMockDataResponse,
+        schemaAnalysis: mockSchemaAnalysis,
+      });
 
       // advance to the schema editor step
       userEvent.click(screen.getByText('Confirm'));
@@ -359,7 +394,10 @@ describe('MockDataGeneratorModal', () => {
     });
 
     it('shows correct values for the faker schema editor', async () => {
-      await renderModal({ mockServices: mockServicesWithMockDataResponse });
+      await renderModal({
+        mockServices: mockServicesWithMockDataResponse,
+        schemaAnalysis: mockSchemaAnalysis,
+      });
 
       // advance to the schema editor step
       userEvent.click(screen.getByText('Confirm'));
@@ -402,6 +440,124 @@ describe('MockDataGeneratorModal', () => {
       );
     });
 
+    it('does not show any fields that are not in the input schema', async () => {
+      const mockServices = createMockServices();
+      mockServices.atlasAiService.getMockDataSchema = () =>
+        Promise.resolve({
+          content: {
+            fields: [
+              {
+                fieldPath: 'name',
+                mongoType: 'string',
+                fakerMethod: 'person.firstName',
+                fakerArgs: [],
+                isArray: false,
+                probability: 1.0,
+              },
+              {
+                fieldPath: 'email',
+                mongoType: 'string',
+                fakerMethod: 'internet.email',
+                fakerArgs: [],
+                isArray: false,
+                probability: 1.0,
+              },
+            ],
+          },
+        });
+      await renderModal({
+        mockServices,
+      });
+
+      // advance to the schema editor step
+      userEvent.click(screen.getByText('Confirm'));
+
+      await waitFor(() => {
+        expect(screen.getByTestId('faker-schema-editor')).to.exist;
+      });
+
+      expect(screen.getByText('name')).to.exist;
+      expect(screen.queryByText('email')).to.not.exist;
+    });
+
+    it('shows unmapped fields as "Unrecognized"', async () => {
+      const mockServices = createMockServices();
+      mockServices.atlasAiService.getMockDataSchema = () =>
+        Promise.resolve({
+          content: {
+            fields: [
+              {
+                fieldPath: 'name',
+                mongoType: 'String',
+                fakerMethod: 'person.firstName',
+                fakerArgs: [],
+                isArray: false,
+                probability: 1.0,
+              },
+              {
+                fieldPath: 'age',
+                mongoType: 'Int32',
+                fakerMethod: 'number.int',
+                fakerArgs: [],
+                isArray: false,
+                probability: 1.0,
+              },
+            ],
+          },
+        });
+
+      await renderModal({
+        mockServices,
+        schemaAnalysis: {
+          ...defaultSchemaAnalysisState,
+          processedSchema: {
+            name: {
+              type: 'String',
+              probability: 1.0,
+            },
+            age: {
+              type: 'Int32',
+              probability: 1.0,
+            },
+            type: {
+              type: 'String',
+              probability: 1.0,
+              sample_values: ['cat', 'dog'],
+            },
+          },
+          sampleDocument: { name: 'Peaches', age: 10, type: 'cat' },
+        },
+      });
+
+      // advance to the schema editor step
+      userEvent.click(screen.getByText('Confirm'));
+
+      await waitFor(() => {
+        expect(screen.getByTestId('faker-schema-editor')).to.exist;
+      });
+
+      // select the "name" field
+      userEvent.click(screen.getByText('name'));
+      expect(screen.getByLabelText('JSON Type')).to.have.value('String');
+      expect(screen.getByLabelText('Faker Function')).to.have.value(
+        'person.firstName'
+      );
+
+      // select the "age" field
+      userEvent.click(screen.getByText('age'));
+      expect(screen.getByLabelText('JSON Type')).to.have.value('Int32');
+      expect(screen.getByLabelText('Faker Function')).to.have.value(
+        'number.int'
+      );
+
+      // select the "type" field
+      userEvent.click(screen.getByText('type'));
+      expect(screen.getByLabelText('JSON Type')).to.have.value('String');
+      expect(screen.getByLabelText('Faker Function')).to.have.value(
+        'Unrecognized'
+      );
+    });
+
     it('disables the Next button when the faker schema mapping is not confirmed', async () => {
       await renderModal({
         mockServices: mockServicesWithMockDataResponse,

packages/compass-collection/src/modules/collection-tab.ts

@@ -16,10 +16,7 @@ import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'
 import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
 import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
 import { type PreferencesAccess } from 'compass-preferences-model/provider';
-import type {
-  MockDataSchemaRequest,
-  MockDataSchemaResponse,
-} from '@mongodb-js/compass-generative-ai';
+import type { MockDataSchemaRequest } from '@mongodb-js/compass-generative-ai';
 import { isInternalFieldPath } from 'hadron-document';
 import toNS from 'mongodb-ns';
 import {
@@ -699,32 +696,74 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction<void> => {
   };
 };
 
+/**
+ * Validates a given faker schema against an input schema.
+ *
+ * - Filters out fields from the faker schema that do not exist in the input schema.
+ * - Validates the `fakerMethod` for each field, marking it as unrecognized if invalid.
+ * - Adds any unmapped input schema fields to the result with an unrecognized faker method.
+ *
+ * @param inputSchema - The schema definition for the input, mapping field names to their metadata.
+ * @param fakerSchema - The array of faker schema mappings to validate and map.
+ * @param logger - Logger instance used to log warnings for invalid faker methods.
+ * @returns An array of validated faker schema mappings, including all input schema fields.
+ */
 const validateFakerSchema = (
-  fakerSchema: MockDataSchemaResponse,
+  inputSchema: Record<string, FieldInfo>,
+  fakerSchema: FakerSchemaMapping[],
   logger: Logger
-) => {
-  return fakerSchema.content.fields.map((field) => {
-    const { fakerMethod } = field;
-
-    const [moduleName, methodName, ...rest] = fakerMethod.split('.');
-    if (
-      rest.length > 0 ||
-      typeof (faker as any)[moduleName]?.[methodName] !== 'function'
-    ) {
-      logger.log.warn(
-        mongoLogId(1_001_000_372),
-        'Collection',
-        'Invalid faker method',
-        { fakerMethod }
-      );
-      return {
-        ...field,
-        fakerMethod: UNRECOGNIZED_FAKER_METHOD,
-      };
-    }
+): FakerSchemaMapping[] => {
+  const inputSchemaFields = Object.keys(inputSchema);
+  const validatedFakerSchema = fakerSchema
+    // Drop fields that don't match the input schema structure
+    .filter((field) => inputSchema[field.fieldPath])
+    .map((field) => {
+      const { fakerMethod } = field;
+
+      // validate faker method
+      const methodSegments = fakerMethod.split('.');
+      let methodRef: any = faker;
+      for (const segment of methodSegments) {
+        if (
+          methodRef &&
+          typeof methodRef === 'object' &&
+          segment in methodRef
+        ) {
+          methodRef = methodRef[segment];
+        } else {
+          methodRef = undefined;
+          break;
+        }
+      }
+      if (typeof methodRef !== 'function') {
+        logger.log.warn(
+          mongoLogId(1_001_000_372),
+          'Collection',
+          'Invalid faker method',
+          { fakerMethod }
+        );
+        return {
+          ...field,
+          fakerMethod: UNRECOGNIZED_FAKER_METHOD,
+        };
+      }
 
-    return field;
-  });
+      return field;
+    });
+  const unmappedInputFields = inputSchemaFields.filter(
+    (field) =>
+      !validatedFakerSchema.find(({ fieldPath }) => fieldPath === field)
+  );
+  // Default unmapped input fields to "Unrecognized" faker method
+  const unmappedFields = unmappedInputFields.map((field) => ({
+    fieldPath: field,
+    fakerMethod: UNRECOGNIZED_FAKER_METHOD,
+    mongoType: inputSchema[field].type,
+    fakerArgs: [],
+    probability: 1,
+  }));
+
+  return [...validatedFakerSchema, ...unmappedFields];
 };
 
 export const generateFakerMappings = (): CollectionThunkAction<
@@ -793,7 +832,11 @@ export const generateFakerMappings = (): CollectionThunkAction<
         connectionInfoRef.current
       );
 
-      const validatedFakerSchema = validateFakerSchema(response, logger);
+      const validatedFakerSchema = validateFakerSchema(
+        schemaAnalysis.processedSchema,
+        response.content.fields,
+        logger
+      );
 
       fakerSchemaGenerationAbortControllerRef.current = undefined;
       dispatch({