Skip to content

Commit 18dcb0b

Browse files
jcobisjcobis
committed
Separate original LLM response from edited faker schema
1 parent 27e95a3 commit 18dcb0b

File tree

5 files changed

+115
-15
lines changed

5 files changed

+115
-15
lines changed

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

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,12 @@ 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 { FakerSchema, MockDataGeneratorState } from './types';
1617
import type { MongoDBFieldType } from '@mongodb-js/compass-generative-ai';
18+
import { updateEditedFakerSchema } from '../../modules/collection-tab';
1719

1820
const containerStyles = css({
1921
display: 'flex',
@@ -52,9 +54,11 @@ const schemaEditorLoaderStyles = css({
5254
const FakerSchemaEditorContent = ({
5355
fakerSchema,
5456
onSchemaConfirmed,
57+
onSchemaChange,
5558
}: {
5659
fakerSchema: FakerSchema;
5760
onSchemaConfirmed: (isConfirmed: boolean) => void;
61+
onSchemaChange: (updatedSchema: FakerSchema) => void;
5862
}) => {
5963
const [fakerSchemaFormValues, setFakerSchemaFormValues] =
6064
React.useState<FakerSchema>(fakerSchema);
@@ -73,27 +77,31 @@ const FakerSchemaEditorContent = ({
7377
const onJsonTypeSelect = (newJsonType: MongoDBFieldType) => {
7478
const currentMapping = fakerSchemaFormValues[activeField];
7579
if (currentMapping) {
76-
setFakerSchemaFormValues({
80+
const updatedSchema = {
7781
...fakerSchemaFormValues,
7882
[activeField]: {
7983
...currentMapping,
8084
mongoType: newJsonType,
8185
},
82-
});
86+
};
87+
setFakerSchemaFormValues(updatedSchema);
88+
onSchemaChange(updatedSchema);
8389
resetIsSchemaConfirmed();
8490
}
8591
};
8692

8793
const onFakerFunctionSelect = (newFakerFunction: string) => {
8894
const currentMapping = fakerSchemaFormValues[activeField];
8995
if (currentMapping) {
90-
setFakerSchemaFormValues({
96+
const updatedSchema = {
9197
...fakerSchemaFormValues,
9298
[activeField]: {
9399
...currentMapping,
94100
fakerMethod: newFakerFunction,
95101
},
96-
});
102+
};
103+
setFakerSchemaFormValues(updatedSchema);
104+
onSchemaChange(updatedSchema);
97105
resetIsSchemaConfirmed();
98106
}
99107
};
@@ -131,10 +139,12 @@ const FakerSchemaEditorContent = ({
131139
const FakerSchemaEditorScreen = ({
132140
onSchemaConfirmed,
133141
fakerSchemaGenerationState,
142+
dispatch,
134143
}: {
135144
isSchemaConfirmed: boolean;
136145
onSchemaConfirmed: (isConfirmed: boolean) => void;
137146
fakerSchemaGenerationState: MockDataGeneratorState;
147+
dispatch: (action: any) => void;
138148
}) => {
139149
return (
140150
<div data-testid="faker-schema-editor" className={containerStyles}>
@@ -160,12 +170,15 @@ const FakerSchemaEditorScreen = ({
160170
)}
161171
{fakerSchemaGenerationState.status === 'completed' && (
162172
<FakerSchemaEditorContent
163-
fakerSchema={fakerSchemaGenerationState.fakerSchema}
173+
fakerSchema={fakerSchemaGenerationState.editedFakerSchema}
164174
onSchemaConfirmed={onSchemaConfirmed}
175+
onSchemaChange={(updatedSchema) =>
176+
dispatch(updateEditedFakerSchema(updatedSchema))
177+
}
165178
/>
166179
)}
167180
</div>
168181
);
169182
};
170183

171-
export default FakerSchemaEditorScreen;
184+
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/script-screen.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ const mapStateToProps = (state: CollectionState) => {
219219
return {
220220
fakerSchema:
221221
fakerSchemaGeneration.status === 'completed'
222-
? fakerSchemaGeneration.fakerSchema
222+
? fakerSchemaGeneration.editedFakerSchema
223223
: null,
224224
namespace,
225225
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: 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: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ export enum CollectionActions {
127127
FakerMappingGenerationStarted = 'compass-collection/FakerMappingGenerationStarted',
128128
FakerMappingGenerationCompleted = 'compass-collection/FakerMappingGenerationCompleted',
129129
FakerMappingGenerationFailed = 'compass-collection/FakerMappingGenerationFailed',
130+
FakerSchemaEdited = 'compass-collection/FakerSchemaEdited',
130131
}
131132

132133
interface CollectionMetadataFetchedAction {
@@ -196,6 +197,11 @@ export interface FakerMappingGenerationFailedAction {
196197
requestId: string;
197198
}
198199

200+
export interface FakerSchemaEditedAction {
201+
type: CollectionActions.FakerSchemaEdited;
202+
editedFakerSchema: FakerSchema;
203+
}
204+
199205
const reducer: Reducer<CollectionState, Action> = (
200206
state = {
201207
// TODO(COMPASS-7782): use hook to get the workspace tab id instead
@@ -457,7 +463,8 @@ const reducer: Reducer<CollectionState, Action> = (
457463
...state,
458464
fakerSchemaGeneration: {
459465
status: 'completed',
460-
fakerSchema: action.fakerSchema,
466+
originalLlmResponse: action.fakerSchema,
467+
editedFakerSchema: action.fakerSchema, // Initially same as LLM response
461468
requestId: action.requestId,
462469
},
463470
};
@@ -487,6 +494,25 @@ const reducer: Reducer<CollectionState, Action> = (
487494
};
488495
}
489496

497+
if (
498+
isAction<FakerSchemaEditedAction>(
499+
action,
500+
CollectionActions.FakerSchemaEdited
501+
)
502+
) {
503+
if (state.fakerSchemaGeneration.status !== 'completed') {
504+
return state;
505+
}
506+
507+
return {
508+
...state,
509+
fakerSchemaGeneration: {
510+
...state.fakerSchemaGeneration,
511+
editedFakerSchema: action.editedFakerSchema,
512+
},
513+
};
514+
}
515+
490516
return state;
491517
};
492518

@@ -528,6 +554,12 @@ export const mockDataGeneratorPreviousButtonClicked = (): CollectionThunkAction<
528554
};
529555
};
530556

557+
export const updateEditedFakerSchema = (
558+
editedFakerSchema: FakerSchema
559+
): FakerSchemaEditedAction => {
560+
return { type: CollectionActions.FakerSchemaEdited, editedFakerSchema };
561+
};
562+
531563
export const selectTab = (
532564
tabName: CollectionSubtab
533565
): CollectionThunkAction<void> => {

0 commit comments

Comments
 (0)