Skip to content

Commit 8c2d893

Browse files
jcobisjcobis
authored
feat(compass-collection): Separate original LLM response from edited faker schema in Mock Data Generator – CLOUDP-350455 (#7437)
* Separate original LLM response from edited faker schema * WIP * WIP * Merge conflict * Readonly
1 parent 0921b1d commit 8c2d893

File tree

6 files changed

+175
-16
lines changed

6 files changed

+175
-16
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
SpinLoaderWithLabel,
1111
} from '@mongodb-js/compass-components';
1212
import React from 'react';
13+
import { connect } from 'react-redux';
1314
import FieldSelector from './schema-field-selector';
1415
import FakerMappingSelector from './faker-mapping-selector';
1516
import type {
@@ -172,7 +173,6 @@ const FakerSchemaEditorScreen = ({
172173
onSchemaConfirmed,
173174
fakerSchemaGenerationState,
174175
}: {
175-
isSchemaConfirmed: boolean;
176176
onSchemaConfirmed: (isConfirmed: boolean) => void;
177177
fakerSchemaGenerationState: MockDataGeneratorState;
178178
}) => {
@@ -200,12 +200,12 @@ const FakerSchemaEditorScreen = ({
200200
)}
201201
{fakerSchemaGenerationState.status === 'completed' && (
202202
<FakerSchemaEditorContent
203-
fakerSchema={fakerSchemaGenerationState.fakerSchema}
203+
fakerSchema={fakerSchemaGenerationState.editedFakerSchema}
204204
onSchemaConfirmed={onSchemaConfirmed}
205205
/>
206206
)}
207207
</div>
208208
);
209209
};
210210

211-
export default FakerSchemaEditorScreen;
211+
export default connect()(FakerSchemaEditorScreen);

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

Lines changed: 60 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -824,7 +824,15 @@ describe('MockDataGeneratorModal', () => {
824824
currentStep: MockDataGeneratorStep.GENERATE_DATA,
825825
fakerSchemaGeneration: {
826826
status: 'completed',
827-
fakerSchema: {
827+
originalLlmResponse: {
828+
name: {
829+
fakerMethod: 'person.firstName',
830+
fakerArgs: [],
831+
probability: 1.0,
832+
mongoType: 'String',
833+
},
834+
},
835+
editedFakerSchema: {
828836
name: {
829837
fakerMethod: 'person.firstName',
830838
fakerArgs: [],
@@ -848,7 +856,15 @@ describe('MockDataGeneratorModal', () => {
848856
currentStep: MockDataGeneratorStep.GENERATE_DATA,
849857
fakerSchemaGeneration: {
850858
status: 'completed',
851-
fakerSchema: {
859+
originalLlmResponse: {
860+
name: {
861+
fakerMethod: 'person.firstName',
862+
fakerArgs: [],
863+
probability: 1.0,
864+
mongoType: 'String',
865+
},
866+
},
867+
editedFakerSchema: {
852868
name: {
853869
fakerMethod: 'person.firstName',
854870
fakerArgs: [],
@@ -872,7 +888,15 @@ describe('MockDataGeneratorModal', () => {
872888
currentStep: MockDataGeneratorStep.GENERATE_DATA,
873889
fakerSchemaGeneration: {
874890
status: 'completed',
875-
fakerSchema: {
891+
originalLlmResponse: {
892+
name: {
893+
fakerMethod: 'person.firstName',
894+
fakerArgs: [],
895+
probability: 1.0,
896+
mongoType: 'String',
897+
},
898+
},
899+
editedFakerSchema: {
876900
name: {
877901
fakerMethod: 'person.firstName',
878902
fakerArgs: [],
@@ -919,7 +943,15 @@ describe('MockDataGeneratorModal', () => {
919943
connectionInfo: atlasConnectionInfo,
920944
fakerSchemaGeneration: {
921945
status: 'completed',
922-
fakerSchema: {
946+
originalLlmResponse: {
947+
name: {
948+
fakerMethod: 'person.firstName',
949+
fakerArgs: [],
950+
probability: 1.0,
951+
mongoType: 'String',
952+
},
953+
},
954+
editedFakerSchema: {
923955
name: {
924956
fakerMethod: 'person.firstName',
925957
fakerArgs: [],
@@ -971,7 +1003,15 @@ describe('MockDataGeneratorModal', () => {
9711003
currentStep: MockDataGeneratorStep.GENERATE_DATA,
9721004
fakerSchemaGeneration: {
9731005
status: 'completed',
974-
fakerSchema: {
1006+
originalLlmResponse: {
1007+
name: {
1008+
fakerMethod: 'person.firstName',
1009+
fakerArgs: [],
1010+
probability: 1.0,
1011+
mongoType: 'String',
1012+
},
1013+
},
1014+
editedFakerSchema: {
9751015
name: {
9761016
fakerMethod: 'person.firstName',
9771017
fakerArgs: [],
@@ -1001,7 +1041,21 @@ describe('MockDataGeneratorModal', () => {
10011041
currentStep: MockDataGeneratorStep.GENERATE_DATA,
10021042
fakerSchemaGeneration: {
10031043
status: 'completed',
1004-
fakerSchema: {
1044+
originalLlmResponse: {
1045+
name: {
1046+
fakerMethod: 'person.firstName',
1047+
fakerArgs: [],
1048+
probability: 1.0,
1049+
mongoType: 'String',
1050+
},
1051+
email: {
1052+
fakerMethod: 'internet.email',
1053+
fakerArgs: [],
1054+
probability: 1.0,
1055+
mongoType: 'String',
1056+
},
1057+
},
1058+
editedFakerSchema: {
10051059
name: {
10061060
fakerMethod: 'person.firstName',
10071061
fakerArgs: [],

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,6 @@ const MockDataGeneratorModal = ({
8383
case MockDataGeneratorStep.SCHEMA_EDITOR:
8484
return (
8585
<FakerSchemaEditorScreen
86-
isSchemaConfirmed={isSchemaConfirmed}
8786
onSchemaConfirmed={setIsSchemaConfirmed}
8887
fakerSchemaGenerationState={fakerSchemaGenerationState}
8988
/>
@@ -100,7 +99,7 @@ const MockDataGeneratorModal = ({
10099
<PreviewScreen
101100
confirmedFakerSchema={
102101
fakerSchemaGenerationState.status === 'completed'
103-
? fakerSchemaGenerationState.fakerSchema
102+
? fakerSchemaGenerationState.editedFakerSchema
104103
: {}
105104
}
106105
/>
@@ -111,7 +110,6 @@ const MockDataGeneratorModal = ({
111110
}, [
112111
currentStep,
113112
fakerSchemaGenerationState,
114-
isSchemaConfirmed,
115113
documentCount,
116114
setDocumentCount,
117115
]);

‎packages/compass-collection/src/components/mock-data-generator-modal/script-screen.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -217,7 +217,7 @@ const mapStateToProps = (state: CollectionState) => {
217217
return {
218218
fakerSchema:
219219
fakerSchemaGeneration.status === 'completed'
220-
? fakerSchemaGeneration.fakerSchema
220+
? fakerSchemaGeneration.editedFakerSchema
221221
: null,
222222
namespace,
223223
arrayLengthMap:

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ type MockDataGeneratorInProgressState = {
2121

2222
type MockDataGeneratorCompletedState = {
2323
status: 'completed';
24-
fakerSchema: FakerSchema;
24+
originalLlmResponse: Readonly<FakerSchema>; // Immutable LLM response
25+
editedFakerSchema: FakerSchema; // User-modified version
2526
requestId: string;
2627
};
2728

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

Lines changed: 108 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@ import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'
1616
import type { experimentationServiceLocator } from '@mongodb-js/compass-telemetry/provider';
1717
import { type Logger, mongoLogId } from '@mongodb-js/compass-logging/provider';
1818
import { type PreferencesAccess } from 'compass-preferences-model/provider';
19-
import type { MockDataSchemaRequest } from '@mongodb-js/compass-generative-ai';
19+
import type {
20+
MockDataSchemaRequest,
21+
MongoDBFieldType,
22+
} from '@mongodb-js/compass-generative-ai';
2023
import { isInternalFieldPath } from 'hadron-document';
2124
import toNS from 'mongodb-ns';
2225
import {
@@ -127,6 +130,8 @@ export enum CollectionActions {
127130
FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted',
128131
FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted',
129132
FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed',
133+
FakerFieldTypeChanged = 'compass-collection/FakerFieldTypeChanged',
134+
FakerFieldMethodChanged = 'compass-collection/FakerFieldMethodChanged',
130135
}
131136

132137
interface CollectionMetadataFetchedAction {
@@ -196,6 +201,18 @@ export interface FakerMappingGenerationFailedAction {
196201
requestId: string;
197202
}
198203

204+
export interface FakerFieldTypeChangedAction {
205+
type: CollectionActions.FakerFieldTypeChanged;
206+
fieldPath: string;
207+
mongoType: MongoDBFieldType;
208+
}
209+
210+
export interface FakerFieldMethodChangedAction {
211+
type: CollectionActions.FakerFieldMethodChanged;
212+
fieldPath: string;
213+
fakerMethod: string;
214+
}
215+
199216
const reducer: Reducer<CollectionState, Action> = (
200217
state = {
201218
// TODO(COMPASS-7782): use hook to get the workspace tab id instead
@@ -457,7 +474,8 @@ const reducer: Reducer<CollectionState, Action> = (
457474
...state,
458475
fakerSchemaGeneration: {
459476
status: 'completed',
460-
fakerSchema: action.fakerSchema,
477+
originalLlmResponse: action.fakerSchema,
478+
editedFakerSchema: action.fakerSchema, // Initially same as LLM response
461479
requestId: action.requestId,
462480
},
463481
};
@@ -487,6 +505,72 @@ const reducer: Reducer<CollectionState, Action> = (
487505
};
488506
}
489507

508+
if (
509+
isAction<FakerFieldTypeChangedAction>(
510+
action,
511+
CollectionActions.FakerFieldTypeChanged
512+
)
513+
) {
514+
if (state.fakerSchemaGeneration.status !== 'completed') {
515+
return state;
516+
}
517+
518+
const { fieldPath, mongoType } = action;
519+
const currentMapping =
520+
state.fakerSchemaGeneration.editedFakerSchema[fieldPath];
521+
522+
if (!currentMapping) {
523+
return state;
524+
}
525+
526+
return {
527+
...state,
528+
fakerSchemaGeneration: {
529+
...state.fakerSchemaGeneration,
530+
editedFakerSchema: {
531+
...state.fakerSchemaGeneration.editedFakerSchema,
532+
[fieldPath]: {
533+
...currentMapping,
534+
mongoType,
535+
},
536+
},
537+
},
538+
};
539+
}
540+
541+
if (
542+
isAction<FakerFieldMethodChangedAction>(
543+
action,
544+
CollectionActions.FakerFieldMethodChanged
545+
)
546+
) {
547+
if (state.fakerSchemaGeneration.status !== 'completed') {
548+
return state;
549+
}
550+
551+
const { fieldPath, fakerMethod } = action;
552+
const currentMapping =
553+
state.fakerSchemaGeneration.editedFakerSchema[fieldPath];
554+
555+
if (!currentMapping) {
556+
return state;
557+
}
558+
559+
return {
560+
...state,
561+
fakerSchemaGeneration: {
562+
...state.fakerSchemaGeneration,
563+
editedFakerSchema: {
564+
...state.fakerSchemaGeneration.editedFakerSchema,
565+
[fieldPath]: {
566+
...currentMapping,
567+
fakerMethod,
568+
},
569+
},
570+
},
571+
};
572+
}
573+
490574
return state;
491575
};
492576

@@ -528,6 +612,28 @@ export const mockDataGeneratorPreviousButtonClicked = (): CollectionThunkAction<
528612
};
529613
};
530614

615+
export const fakerFieldTypeChanged = (
616+
fieldPath: string,
617+
mongoType: MongoDBFieldType
618+
): FakerFieldTypeChangedAction => {
619+
return {
620+
type: CollectionActions.FakerFieldTypeChanged,
621+
fieldPath,
622+
mongoType,
623+
};
624+
};
625+
626+
export const fakerFieldMethodChanged = (
627+
fieldPath: string,
628+
fakerMethod: string
629+
): FakerFieldMethodChangedAction => {
630+
return {
631+
type: CollectionActions.FakerFieldMethodChanged,
632+
fieldPath,
633+
fakerMethod,
634+
};
635+
};
636+
531637
export const selectTab = (
532638
tabName: CollectionSubtab
533639
): CollectionThunkAction<void> => {

0 commit comments

Comments
 (0)