chore(compass-collection): ignore mongoType from LLM response CLOUDP-351514

packages/compass-collection/src/components/mock-data-generator-modal/constants.ts

@@ -1,5 +1,5 @@
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import { MockDataGeneratorStep } from './types';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
 
 export const StepButtonLabelMap = {
   [MockDataGeneratorStep.SCHEMA_CONFIRMATION]: 'Confirm',
@@ -12,6 +12,25 @@ export const StepButtonLabelMap = {
 export const DEFAULT_DOCUMENT_COUNT = 1000;
 export const MAX_DOCUMENT_COUNT = 100000;
 
+export const MongoDBFieldTypeValues: MongoDBFieldType[] = [
+  'String',
+  'Number',
+  'Boolean',
+  'Date',
+  'Int32',
+  'Decimal128',
+  'Long',
+  'ObjectId',
+  'RegExp',
+  'Symbol',
+  'MaxKey',
+  'MinKey',
+  'Binary',
+  'Code',
+  'Timestamp',
+  'DBRef',
+];
+
 /**
  * Map of MongoDB types to available Faker v9 methods.
  * Not all Faker methods are included here.

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

@@ -10,11 +10,14 @@ import {
 import sinon from 'sinon';
 import FakerMappingSelector from './faker-mapping-selector';
 import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
-import { MONGO_TYPE_TO_FAKER_METHODS } from './constants';
-import { MongoDBFieldTypeValues } from '@mongodb-js/compass-generative-ai';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
+import {
+  MONGO_TYPE_TO_FAKER_METHODS,
+  MongoDBFieldTypeValues,
+} from './constants';
 import type { FakerArg } from './script-generation-utils';
 
-const mockActiveJsonType = MongoDBFieldTypeValues.String;
+const mockActiveJsonType: MongoDBFieldType = 'String';
 const mockActiveFakerFunction = 'lorem.word';
 const mockActiveFakerArgs: Array<FakerArg> = [];
 const onJsonTypeSelectStub = sinon.stub();
@@ -26,9 +29,6 @@ describe('FakerMappingSelector', () => {
   });
 
   it('should display all MongoDB types in the dropdown', async () => {
-    // Check that all MongoDB types from the constant are present
-    const mongoTypes = Object.keys(MongoDBFieldTypeValues);
-
     render(
       <FakerMappingSelector
         activeJsonType={mockActiveJsonType}
@@ -42,7 +42,7 @@ describe('FakerMappingSelector', () => {
     const jsonTypeSelect = screen.getByLabelText('JSON Type');
     userEvent.click(jsonTypeSelect);
 
-    for (const type of mongoTypes) {
+    for (const type of MongoDBFieldTypeValues) {
       await waitFor(() => {
         expect(screen.getByRole('option', { name: type })).to.exist;
       });

packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx

@@ -12,9 +12,11 @@ import {
 } from '@mongodb-js/compass-components';
 import React, { useMemo } from 'react';
 import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
-import { MongoDBFieldTypeValues } from '@mongodb-js/compass-generative-ai';
-import { MONGO_TYPE_TO_FAKER_METHODS } from './constants';
+import {
+  MONGO_TYPE_TO_FAKER_METHODS,
+  MongoDBFieldTypeValues,
+} from './constants';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
 import type { FakerArg } from './script-generation-utils';
 
 const fieldMappingSelectorsStyles = css({
@@ -93,7 +95,7 @@ const FakerMappingSelector = ({
         value={activeJsonType}
         onChange={(value) => onJsonTypeSelect(value as MongoDBFieldType)}
       >
-        {Object.values(MongoDBFieldTypeValues).map((type) => (
+        {MongoDBFieldTypeValues.map((type) => (
           <Option key={type} value={type}>
             {type}
           </Option>

packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx

@@ -13,13 +13,13 @@ import React from 'react';
 import { connect } from 'react-redux';
 import FieldSelector from './schema-field-selector';
 import FakerMappingSelector from './faker-mapping-selector';
+import { getDefaultFakerMethod } from './script-generation-utils';
 import type {
   FakerSchema,
   FakerFieldMapping,
   MockDataGeneratorState,
 } from './types';
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
-import { getDefaultFakerMethod } from './script-generation-utils';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
 
 const containerStyles = css({
   display: 'flex',

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

@@ -97,7 +97,6 @@ describe('MockDataGeneratorModal', () => {
             fields: [
               {
                 fieldPath: 'name',
-                mongoType: 'String',
                 fakerMethod: 'person.firstName',
                 fakerArgs: [],
               },
@@ -326,25 +325,21 @@ describe('MockDataGeneratorModal', () => {
         fields: [
           {
             fieldPath: 'name',
-            mongoType: 'String',
             fakerMethod: 'person.firstName',
             fakerArgs: [],
           },
           {
             fieldPath: 'age',
-            mongoType: 'Int32',
             fakerMethod: 'number.int',
             fakerArgs: [],
           },
           {
             fieldPath: 'email',
-            mongoType: 'String',
             fakerMethod: 'internet',
             fakerArgs: [],
           },
           {
             fieldPath: 'username',
-            mongoType: 'String',
             fakerMethod: 'noSuchMethod',
             fakerArgs: [],
           },
@@ -439,15 +434,13 @@ describe('MockDataGeneratorModal', () => {
           fields: [
             {
               fieldPath: 'name',
-              mongoType: 'String',
               fakerMethod: 'person.firstName',
               fakerArgs: [],
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'email',
-              mongoType: 'String',
               fakerMethod: 'internet.email',
               fakerArgs: [],
               isArray: false,
@@ -477,15 +470,13 @@ describe('MockDataGeneratorModal', () => {
           fields: [
             {
               fieldPath: 'name',
-              mongoType: 'String',
               fakerMethod: 'person.firstName',
               fakerArgs: [],
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'age',
-              mongoType: 'Int32',
               fakerMethod: 'number.int',
               fakerArgs: [],
               isArray: false,
@@ -553,15 +544,13 @@ describe('MockDataGeneratorModal', () => {
           fields: [
             {
               fieldPath: 'name',
-              mongoType: 'String',
               fakerMethod: 'person.firstName',
               fakerArgs: largeLengthArgs,
               isArray: false,
               probability: 1.0,
             },
             {
               fieldPath: 'age',
-              mongoType: 'Int32',
               fakerMethod: 'number.int',
               fakerArgs: [
                 {
@@ -575,7 +564,6 @@ describe('MockDataGeneratorModal', () => {
             },
             {
               fieldPath: 'username',
-              mongoType: 'String',
               fakerMethod: 'string.alpha',
               // large string
               fakerArgs: ['a'.repeat(1001)],
@@ -584,7 +572,6 @@ describe('MockDataGeneratorModal', () => {
             },
             {
               fieldPath: 'avatar',
-              mongoType: 'String',
               fakerMethod: 'image.url',
               fakerArgs: [
                 {
@@ -653,6 +640,7 @@ describe('MockDataGeneratorModal', () => {
     it('disables the Next button when the faker schema mapping is not confirmed', async () => {
       await renderModal({
         mockServices: mockServicesWithMockDataResponse,
+        schemaAnalysis: mockSchemaAnalysis,
       });
 
       // advance to the schema editor step
@@ -669,6 +657,7 @@ describe('MockDataGeneratorModal', () => {
     it('resets the confirm schema mapping state when the user clicks the back button then goes back to the schema editor step', async () => {
       await renderModal({
         mockServices: mockServicesWithMockDataResponse,
+        schemaAnalysis: mockSchemaAnalysis,
       });
 
       // advance to the schema editor step
@@ -705,6 +694,7 @@ describe('MockDataGeneratorModal', () => {
     it('preserves the confirm schema mapping state when the user clicks the next button then goes back to the schema editor step', async () => {
       await renderModal({
         mockServices: mockServicesWithMockDataResponse,
+        schemaAnalysis: mockSchemaAnalysis,
       });
 
       // advance to the schema editor step

packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts

@@ -1,5 +1,5 @@
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import type { FakerFieldMapping } from './types';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
 import { prettify } from '@mongodb-js/compass-editor';
 import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab';
 import { faker } from '@faker-js/faker/locale/en';

packages/compass-collection/src/components/mock-data-generator-modal/types.ts

@@ -1,6 +1,6 @@
 import type { MockDataSchemaResponse } from '@mongodb-js/compass-generative-ai';
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import type { FakerArg } from './script-generation-utils';
+import type { MongoDBFieldType } from '../../schema-analysis-types';
 
 export enum MockDataGeneratorStep {
   SCHEMA_CONFIRMATION = 'SCHEMA_CONFIRMATION',

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

@@ -818,17 +818,22 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction<void> => {
 /**
  * Transforms LLM array format to keyed object structure.
  * Moves fieldPath from object property to object key.
+ *
+ * @param fakerSchema - The faker schema array from LLM response
+ * @param inputSchema - The schema definition for the LLM input used to carry over `mongoType` data
  */
 function transformFakerSchemaToObject(
-  fakerSchema: LlmFakerMapping[]
+  fakerSchema: LlmFakerMapping[],
+  inputSchema: Record<string, FieldInfo>
 ): FakerSchema {
   const result: FakerSchema = {};
 
   for (const field of fakerSchema) {
     const { fieldPath, ...fieldMapping } = field;
     result[fieldPath] = {
       ...fieldMapping,
-      mongoType: fieldMapping.mongoType,
+      // Note: `validateFakerSchema` already handles fields that are not present in `inputSchema`
+      mongoType: inputSchema[fieldPath]?.type,
     };
   }
 
@@ -968,7 +973,8 @@ export const generateFakerMappings = (): CollectionThunkAction<
 
       // Transform to keyed object structure
       const transformedFakerSchema = transformFakerSchemaToObject(
-        response.fields
+        response.fields,
+        schemaAnalysis.processedSchema
       );
 
       const validatedFakerSchema = validateFakerSchema(

packages/compass-collection/src/schema-analysis-types.ts

@@ -1,5 +1,5 @@
-import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
 import type { Document } from 'mongodb';
+import type { PrimitiveSchemaType } from 'mongodb-schema';
 
 export const SCHEMA_ANALYSIS_STATE_INITIAL = 'initial';
 export const SCHEMA_ANALYSIS_STATE_ANALYZING = 'analyzing';
@@ -42,6 +42,8 @@ export type SampleValue =
   | null
   | undefined;
 
+export type MongoDBFieldType = PrimitiveSchemaType['name'];
+
 /**
  * Schema field information (for LLM processing)
  */

packages/compass-generative-ai/src/atlas-ai-service.spec.ts

@@ -429,13 +429,11 @@ describe('AtlasAiService', function () {
               fields: [
                 {
                   fieldPath: 'name',
-                  mongoType: 'String',
                   fakerMethod: 'person.fullName',
                   fakerArgs: [],
                 },
                 {
                   fieldPath: 'age',
-                  mongoType: 'Int32',
                   fakerMethod: 'number.int',
                   fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
                 },
@@ -464,13 +462,11 @@ describe('AtlasAiService', function () {
               fields: [
                 {
                   fieldPath: 'name',
-                  mongoType: 'String',
                   fakerMethod: 'person.fullName',
                   fakerArgs: [],
                 },
                 {
                   fieldPath: 'age',
-                  mongoType: 'Int32',
                   fakerMethod: 'number.int',
                   fakerArgs: [{ json: '{"min": 18, "max": 122}' }],
                 },
@@ -504,13 +500,11 @@ describe('AtlasAiService', function () {
               fields: [
                 {
                   fieldPath: 'name',
-                  mongoType: 'String',
                   fakerMethod: 'person.fullName',
                   fakerArgs: [],
                 },
                 {
                   fieldPath: 'age',
-                  mongoType: 'Int32',
                   fakerMethod: 'number.int',
                   fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
                 },
@@ -541,13 +535,11 @@ describe('AtlasAiService', function () {
               fields: [
                 {
                   fieldPath: 'name',
-                  mongoType: 'String',
                   fakerMethod: 'person.fullName',
                   fakerArgs: [],
                 },
                 {
                   fieldPath: 'age',
-                  mongoType: 'Int32',
                   fakerMethod: 'number.int',
                   fakerArgs: [{ json: '{"min": 18, "max": 65}' }],
                 },

packages/compass-generative-ai/src/atlas-ai-service.ts

@@ -1,4 +1,4 @@
-import type { PrimitiveSchemaType, SimplifiedSchema } from 'mongodb-schema';
+import type { SimplifiedSchema } from 'mongodb-schema';
 import {
   type PreferencesAccess,
   isAIFeatureEnabled,
@@ -230,37 +230,10 @@ export interface MockDataSchemaRequest {
   signal: AbortSignal;
 }
 
-/**
- * MongoDB schema type
- */
-export type MongoDBFieldType = PrimitiveSchemaType['name'];
-
-export enum MongoDBFieldTypeValues {
-  String = 'String',
-  Number = 'Number',
-  Boolean = 'Boolean',
-  Date = 'Date',
-  Int32 = 'Int32',
-  Decimal128 = 'Decimal128',
-  Long = 'Long',
-  ObjectId = 'ObjectId',
-  RegExp = 'RegExp',
-  Symbol = 'Symbol',
-  MaxKey = 'MaxKey',
-  MinKey = 'MinKey',
-  Binary = 'Binary',
-  Code = 'Code',
-  Timestamp = 'Timestamp',
-  DBRef = 'DBRef',
-}
-
 export const MockDataSchemaResponseShape = z.object({
   fields: z.array(
     z.object({
       fieldPath: z.string(),
-      mongoType: z.custom<MongoDBFieldType>((val) =>
-        Object.values(MongoDBFieldTypeValues).includes(val)
-      ),
       fakerMethod: z.string(),
       fakerArgs: z.array(
         z.union([