From 67af1a99f46d3ae28a6cf20f3df32f32d16b4f4f Mon Sep 17 00:00:00 2001 From: Kevin Pamaran Date: Thu, 16 Oct 2025 16:26:01 -0400 Subject: [PATCH 1/4] Ignore mongoType from LLM response --- .../mock-data-generator-modal/constants.ts | 21 +++++++++++++- .../faker-mapping-selector.spec.tsx | 13 ++++----- .../faker-mapping-selector.tsx | 10 ++++--- .../faker-schema-editor-screen.tsx | 4 +-- .../mock-data-generator-modal.spec.tsx | 16 ++-------- .../script-generation-utils.ts | 2 +- .../mock-data-generator-modal/types.ts | 2 +- .../src/modules/collection-tab.ts | 11 +++++-- .../src/schema-analysis-types.ts | 4 ++- .../src/atlas-ai-service.spec.ts | 8 ----- .../src/atlas-ai-service.ts | 29 +------------------ packages/compass-generative-ai/src/index.ts | 3 -- 12 files changed, 51 insertions(+), 72 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts b/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts index c5935479ee6..6f7f9d7bbb4 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/constants.ts +++ b/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. diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx index 84d7c129603..5c918b56964 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx @@ -10,11 +10,13 @@ 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 { + MONGO_TYPE_TO_FAKER_METHODS, + MongoDBFieldTypeValues, +} from './constants'; import type { FakerArg } from './script-generation-utils'; -const mockActiveJsonType = MongoDBFieldTypeValues.String; +const mockActiveJsonType = 'String'; const mockActiveFakerFunction = 'lorem.word'; const mockActiveFakerArgs: Array = []; const onJsonTypeSelectStub = sinon.stub(); @@ -26,9 +28,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( { 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; }); diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx index ecff308c314..aa08248262b 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.tsx +++ b/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) => ( diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx index 1d7661a12b1..d86bf702b5e 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-schema-editor-screen.tsx @@ -12,13 +12,13 @@ import { import React from 'react'; 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', diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx index 96492f25e3f..d454ec735c3 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/mock-data-generator-modal.spec.tsx +++ b/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,7 +434,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], isArray: false, @@ -447,7 +441,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'email', - mongoType: 'String', fakerMethod: 'internet.email', fakerArgs: [], isArray: false, @@ -477,7 +470,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: [], isArray: false, @@ -485,7 +477,6 @@ describe('MockDataGeneratorModal', () => { }, { fieldPath: 'age', - mongoType: 'Int32', fakerMethod: 'number.int', fakerArgs: [], isArray: false, @@ -553,7 +544,6 @@ describe('MockDataGeneratorModal', () => { fields: [ { fieldPath: 'name', - mongoType: 'String', fakerMethod: 'person.firstName', fakerArgs: largeLengthArgs, isArray: false, @@ -561,7 +551,6 @@ describe('MockDataGeneratorModal', () => { }, { 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 diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts b/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts index 7179a08ad38..7d26ef9825b 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/script-generation-utils.ts +++ b/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'; diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts b/packages/compass-collection/src/components/mock-data-generator-modal/types.ts index af5150fbd55..eabc0f04061 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/types.ts +++ b/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', diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index a1bfe9866e5..8369b3803d2 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -712,9 +712,12 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction => { /** * Transforms LLM array format to keyed object structure. * Moves fieldPath from object property to object key. + * + * @param inputSchema - schema definition for the LLM input used to carry over `mongoType` data */ function transformFakerSchemaToObject( - fakerSchema: LlmFakerMapping[] + fakerSchema: LlmFakerMapping[], + inputSchema: Record ): FakerSchema { const result: FakerSchema = {}; @@ -722,7 +725,7 @@ function transformFakerSchemaToObject( const { fieldPath, ...fieldMapping } = field; result[fieldPath] = { ...fieldMapping, - mongoType: fieldMapping.mongoType, + mongoType: inputSchema[fieldPath]?.type ?? 'String', }; } @@ -862,7 +865,9 @@ export const generateFakerMappings = (): CollectionThunkAction< // Transform to keyed object structure const transformedFakerSchema = transformFakerSchemaToObject( - response.fields + // do a rename + response.fields, + schemaAnalysis.processedSchema ); const validatedFakerSchema = validateFakerSchema( diff --git a/packages/compass-collection/src/schema-analysis-types.ts b/packages/compass-collection/src/schema-analysis-types.ts index c975f136130..b1705b41cc6 100644 --- a/packages/compass-collection/src/schema-analysis-types.ts +++ b/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) */ diff --git a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts b/packages/compass-generative-ai/src/atlas-ai-service.spec.ts index 008a6ad772a..dbcb793f307 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.spec.ts +++ b/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}' }], }, diff --git a/packages/compass-generative-ai/src/atlas-ai-service.ts b/packages/compass-generative-ai/src/atlas-ai-service.ts index cb19dae3d1a..262be30797a 100644 --- a/packages/compass-generative-ai/src/atlas-ai-service.ts +++ b/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((val) => - Object.values(MongoDBFieldTypeValues).includes(val) - ), fakerMethod: z.string(), fakerArgs: z.array( z.union([ diff --git a/packages/compass-generative-ai/src/index.ts b/packages/compass-generative-ai/src/index.ts index 0089fe592d8..525d5b12008 100644 --- a/packages/compass-generative-ai/src/index.ts +++ b/packages/compass-generative-ai/src/index.ts @@ -35,7 +35,4 @@ export type { MockDataSchemaRequest, MockDataSchemaRawField, MockDataSchemaResponse, - MongoDBFieldType, } from './atlas-ai-service'; - -export { MongoDBFieldTypeValues } from './atlas-ai-service'; From 6123d0a8fa8c2c8372469ae128d57e0b7e827b19 Mon Sep 17 00:00:00 2001 From: Kevin Pamaran Date: Thu, 16 Oct 2025 16:41:14 -0400 Subject: [PATCH 2/4] nit --- .../faker-mapping-selector.spec.tsx | 3 ++- packages/compass-collection/src/modules/collection-tab.ts | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx index 5c918b56964..edae8dedc29 100644 --- a/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx +++ b/packages/compass-collection/src/components/mock-data-generator-modal/faker-mapping-selector.spec.tsx @@ -10,13 +10,14 @@ import { import sinon from 'sinon'; import FakerMappingSelector from './faker-mapping-selector'; import { UNRECOGNIZED_FAKER_METHOD } from '../../modules/collection-tab'; +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 = 'String'; +const mockActiveJsonType: MongoDBFieldType = 'String'; const mockActiveFakerFunction = 'lorem.word'; const mockActiveFakerArgs: Array = []; const onJsonTypeSelectStub = sinon.stub(); diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 8369b3803d2..d91a0c90ebb 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -713,7 +713,8 @@ export const cancelSchemaAnalysis = (): CollectionThunkAction => { * Transforms LLM array format to keyed object structure. * Moves fieldPath from object property to object key. * - * @param inputSchema - schema definition for the LLM input used to carry over `mongoType` data + * @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[], @@ -725,6 +726,7 @@ function transformFakerSchemaToObject( const { fieldPath, ...fieldMapping } = field; result[fieldPath] = { ...fieldMapping, + // Note: fallback to String would only occur if LLM response renames the input field path mongoType: inputSchema[fieldPath]?.type ?? 'String', }; } @@ -865,7 +867,6 @@ export const generateFakerMappings = (): CollectionThunkAction< // Transform to keyed object structure const transformedFakerSchema = transformFakerSchemaToObject( - // do a rename response.fields, schemaAnalysis.processedSchema ); From 93cf62fd70c132b8205236376a52f998c44ba682 Mon Sep 17 00:00:00 2001 From: Kevin Pamaran Date: Fri, 17 Oct 2025 15:13:02 -0400 Subject: [PATCH 3/4] update fallback and note --- packages/compass-collection/src/modules/collection-tab.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index d91a0c90ebb..571558fec16 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/packages/compass-collection/src/modules/collection-tab.ts @@ -726,8 +726,8 @@ function transformFakerSchemaToObject( const { fieldPath, ...fieldMapping } = field; result[fieldPath] = { ...fieldMapping, - // Note: fallback to String would only occur if LLM response renames the input field path - mongoType: inputSchema[fieldPath]?.type ?? 'String', + // Note: `validateFakerSchema` already handles fields that are not present in `inputSchema` + mongoType: inputSchema[fieldPath]?.type, }; } From 85a92f108a70381044b454d8be360e752104dd7e Mon Sep 17 00:00:00 2001 From: Kevin Pamaran Date: Tue, 21 Oct 2025 12:15:49 -0400 Subject: [PATCH 4/4] fix import --- packages/compass-collection/src/modules/collection-tab.ts | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/packages/compass-collection/src/modules/collection-tab.ts b/packages/compass-collection/src/modules/collection-tab.ts index 8a86788f53b..569360e9841 100644 --- a/packages/compass-collection/src/modules/collection-tab.ts +++ b/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, - MongoDBFieldType, -} 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 { @@ -30,6 +27,7 @@ import { type SchemaAnalysisError, type SchemaAnalysisState, type FieldInfo, + type MongoDBFieldType, } from '../schema-analysis-types'; import { calculateSchemaDepth } from '../calculate-schema-depth'; import {