Skip to content

Commit

Permalink
feat: Add gen ai org and project settings to Compass for CompassWeb u…
Browse files Browse the repository at this point in the history
…sage COMPASS-8377 (#6434)

* limit increased and message improved for PROMPT_TOO_LONG

* updating prompt_too_long error msg

* updating test

* initial ticket changes

* entrypoint changed

* prefProps for atlas org and atlas proj prefs

* function name change

* fixing base values

* changing ui default back

* comment nit

* moving prop defs so defaults get set as part of userpreferences

* updating tests

* updating tests
  • Loading branch information
ruchitharajaghatta authored Nov 7, 2024
1 parent d3ae895 commit 5bc3fe9
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 6 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
import { Preferences, type PreferencesAccess } from './preferences';
import type { UserPreferences } from './preferences-schema';
import { type AllPreferences } from './preferences-schema';
import { InMemoryStorage } from './preferences-in-memory-storage';
import { getActiveUser } from './utils';

export class CompassWebPreferencesAccess implements PreferencesAccess {
private _preferences: Preferences;
constructor(preferencesOverrides?: Partial<AllPreferences>) {
this._preferences = new Preferences({
logger: createNoopLogger(),
preferencesStorage: new InMemoryStorage(preferencesOverrides),
});
}

savePreferences(_attributes: Partial<UserPreferences>) {
// Only allow saving the optInDataExplorerGenAIFeatures preference.
if (
Object.keys(_attributes).length === 1 &&
'optInDataExplorerGenAIFeatures' in _attributes
) {
return Promise.resolve(this._preferences.savePreferences(_attributes));
}
return Promise.resolve(this._preferences.getPreferences());
}

refreshPreferences() {
return Promise.resolve(this._preferences.getPreferences());
}

getPreferences() {
return this._preferences.getPreferences();
}

ensureDefaultConfigurableUserPreferences() {
return this._preferences.ensureDefaultConfigurableUserPreferences();
}

getConfigurableUserPreferences() {
return Promise.resolve(this._preferences.getConfigurableUserPreferences());
}

getPreferenceStates() {
return Promise.resolve(this._preferences.getPreferenceStates());
}

onPreferenceValueChanged<K extends keyof AllPreferences>(
preferenceName: K,
callback: (value: AllPreferences[K]) => void
) {
return (
this._preferences?.onPreferencesChanged?.(
(preferences: Partial<AllPreferences>) => {
if (Object.keys(preferences).includes(preferenceName)) {
return callback((preferences as AllPreferences)[preferenceName]);
}
}
) ??
(() => {
/* no fallback */
})
);
}

createSandbox() {
return Promise.resolve(
new CompassWebPreferencesAccess(this.getPreferences())
);
}

getPreferencesUser() {
return getActiveUser(this);
}
}
74 changes: 72 additions & 2 deletions packages/compass-preferences-model/src/preferences-schema.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ export type UserConfigurablePreferences = PermanentFeatureFlags &
| 'web-sandbox-atlas-local'
| 'web-sandbox-atlas-dev'
| 'web-sandbox-atlas';
optInDataExplorerGenAIFeatures: boolean;
// Features that are enabled by default in Compass, but are disabled in Data
// Explorer
enableExplainPlan: boolean;
Expand Down Expand Up @@ -92,7 +93,9 @@ export type InternalUserPreferences = {

// UserPreferences contains all preferences stored to disk.
export type UserPreferences = UserConfigurablePreferences &
InternalUserPreferences;
InternalUserPreferences &
AtlasOrgPreferences &
AtlasProjectPreferences;

export type CliOnlyPreferences = {
exportConnections?: string;
Expand Down Expand Up @@ -210,6 +213,15 @@ export type StoredPreferencesValidator = ReturnType<

export type StoredPreferences = z.output<StoredPreferencesValidator>;

export type AtlasProjectPreferences = {
enableGenAIFeaturesAtlasProject: boolean;
enableGenAISampleDocumentPassingOnAtlasProject: boolean;
};

export type AtlasOrgPreferences = {
enableGenAIFeaturesAtlasOrg: boolean;
};

// Preference definitions
const featureFlagsProps: Required<{
[K in keyof FeatureFlags]: PreferenceDefinition<K>;
Expand Down Expand Up @@ -461,7 +473,10 @@ export const storedUserPreferencesProps: Required<{
short: 'Enable AI Features',
long: 'Allow the use of AI features in Compass which make requests to 3rd party services.',
},
deriveValue: deriveNetworkTrafficOptionState('enableGenAIFeatures'),
deriveValue: deriveValueFromOtherPreferencesAsLogicalAnd(
'enableGenAIFeatures',
['enableGenAIFeaturesAtlasOrg', 'networkTraffic']
),
validator: z.boolean().default(true),
type: 'boolean',
},
Expand Down Expand Up @@ -679,6 +694,16 @@ export const storedUserPreferencesProps: Required<{
.default('atlas'),
type: 'string',
},
optInDataExplorerGenAIFeatures: {
ui: true,
cli: false,
global: false,
description: {
short: 'User Opt-in for Data Explorer Gen AI Features',
},
validator: z.boolean().default(true),
type: 'boolean',
},

enableAtlasSearchIndexes: {
ui: true,
Expand Down Expand Up @@ -861,6 +886,36 @@ export const storedUserPreferencesProps: Required<{
validator: z.boolean().default(false),
type: 'boolean',
},
enableGenAIFeaturesAtlasProject: {
ui: false,
cli: true,
global: true,
description: {
short: 'Enable Gen AI Features on Atlas Project Level',
},
validator: z.boolean().default(true),
type: 'boolean',
},
enableGenAISampleDocumentPassingOnAtlasProject: {
ui: false,
cli: true,
global: true,
description: {
short: 'Enable Gen AI Sample Document Passing on Atlas Project Level',
},
validator: z.boolean().default(true),
type: 'boolean',
},
enableGenAIFeaturesAtlasOrg: {
ui: false,
cli: true,
global: true,
description: {
short: 'Enable Gen AI Features on Atlas Org Level',
},
validator: z.boolean().default(true),
type: 'boolean',
},

...allFeatureFlagsProps,
};
Expand Down Expand Up @@ -1027,6 +1082,21 @@ function deriveNetworkTrafficOptionState<K extends keyof AllPreferences>(
});
}

/** Helper for deriving value/state for preferences from other preferences */
function deriveValueFromOtherPreferencesAsLogicalAnd<
K extends keyof AllPreferences
>(property: K, preferencesToDeriveFrom: K[]): DeriveValueFunction<boolean> {
return (v, s) => ({
value: v(property) && preferencesToDeriveFrom.every((p) => v(p)),
state:
s(property) ??
(preferencesToDeriveFrom.every((p) => v(p))
? preferencesToDeriveFrom.map((p) => s(p)).filter(Boolean)?.[0] ??
'derived'
: undefined),
});
}

/** Helper for defining how to derive value/state for feature-restricting preferences */
function deriveFeatureRestrictingOptionsState<K extends keyof AllPreferences>(
property: K
Expand Down
8 changes: 6 additions & 2 deletions packages/compass-preferences-model/src/preferences.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,7 @@ describe('Preferences class', function () {
expect(states).to.deep.equal({
trackUsageStatistics: 'set-global',
enableMaps: 'set-cli',
enableGenAIFeatures: 'derived',
...expectedReleasedFeatureFlagsStates,
});
});
Expand Down Expand Up @@ -163,7 +164,6 @@ describe('Preferences class', function () {
enableDevTools: 'set-global',
networkTraffic: 'set-global',
trackUsageStatistics: 'set-global',
enableGenAIFeatures: 'set-global',
enableMaps: 'set-cli',
enableShell: 'set-cli',
readOnly: 'set-global',
Expand All @@ -185,6 +185,7 @@ describe('Preferences class', function () {

expect(states).to.deep.equal({
readOnly: 'set-global',
enableGenAIFeatures: 'derived',
...expectedReleasedFeatureFlagsStates,
});
});
Expand Down Expand Up @@ -227,8 +228,8 @@ describe('Preferences class', function () {
},
{
networkTraffic: false,
enableGenAIFeatures: false,
enableMaps: false,
enableGenAIFeatures: false,
enableFeedbackPanel: false,
trackUsageStatistics: false,
autoUpdates: false,
Expand All @@ -250,6 +251,7 @@ describe('Preferences class', function () {
},
hardcoded: {
networkTraffic: false,
enableGenAIFeatures: false,
},
});
const result = preferences.getPreferences();
Expand Down Expand Up @@ -294,13 +296,15 @@ describe('Preferences class', function () {

expect(mainPreferencesStates).to.deep.equal({
trackUsageStatistics: 'set-global',
enableGenAIFeatures: 'derived',
enableMaps: 'set-cli',
...expectedReleasedFeatureFlagsStates,
});

const sandboxPreferencesStates = sandbox.getPreferenceStates();
expect(sandboxPreferencesStates).to.deep.equal({
enableDevTools: 'derived',
enableGenAIFeatures: 'derived',
trackUsageStatistics: 'set-global',
enableMaps: 'set-cli',
enableShell: 'derived',
Expand Down
1 change: 1 addition & 0 deletions packages/compass-preferences-model/src/provider.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
export * from './react';
export { ReadOnlyPreferenceAccess } from './read-only-preferences-access';
export { CompassWebPreferencesAccess } from './compass-web-preferences-access';
export {
useIsAIFeatureEnabled,
isAIFeatureEnabled,
Expand Down
5 changes: 3 additions & 2 deletions packages/compass-web/src/entrypoint.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ import {
} from '@mongodb-js/compass-databases-collections';
import {
PreferencesProvider,
ReadOnlyPreferenceAccess,
CompassWebPreferencesAccess,
} from 'compass-preferences-model/provider';
import type { AllPreferences } from 'compass-preferences-model/provider';
import FieldStorePlugin from '@mongodb-js/compass-field-store';
Expand Down Expand Up @@ -262,7 +262,7 @@ const CompassWeb = ({
});

const preferencesAccess = useRef(
new ReadOnlyPreferenceAccess({
new CompassWebPreferencesAccess({
maxTimeMS: 10_000,
enableExplainPlan: true,
enableAggregationBuilderRunPipeline: true,
Expand All @@ -279,6 +279,7 @@ const CompassWeb = ({
enableShell: false,
enableCreatingNewConnections: false,
enableGlobalWrites: false,
optInDataExplorerGenAIFeatures: false,
...initialPreferences,
})
);
Expand Down

0 comments on commit 5bc3fe9

Please sign in to comment.