feat: Add gen ai org and project settings to Compass for CompassWeb usage 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

Author: ruchitharajaghatta

packages/compass-preferences-model/src/compass-web-preferences-access.ts

@@ -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);
+  }
+}

packages/compass-preferences-model/src/preferences-schema.ts

@@ -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;
@@ -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;
@@ -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>;
@@ -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',
   },
@@ -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,
@@ -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,
 };
@@ -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

packages/compass-preferences-model/src/preferences.spec.ts

@@ -132,6 +132,7 @@ describe('Preferences class', function () {
     expect(states).to.deep.equal({
       trackUsageStatistics: 'set-global',
       enableMaps: 'set-cli',
+      enableGenAIFeatures: 'derived',
       ...expectedReleasedFeatureFlagsStates,
     });
   });
@@ -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',
@@ -185,6 +185,7 @@ describe('Preferences class', function () {
 
     expect(states).to.deep.equal({
       readOnly: 'set-global',
+      enableGenAIFeatures: 'derived',
       ...expectedReleasedFeatureFlagsStates,
     });
   });
@@ -227,8 +228,8 @@ describe('Preferences class', function () {
       },
       {
         networkTraffic: false,
-        enableGenAIFeatures: false,
         enableMaps: false,
+        enableGenAIFeatures: false,
         enableFeedbackPanel: false,
         trackUsageStatistics: false,
         autoUpdates: false,
@@ -250,6 +251,7 @@ describe('Preferences class', function () {
       },
       hardcoded: {
         networkTraffic: false,
+        enableGenAIFeatures: false,
       },
     });
     const result = preferences.getPreferences();
@@ -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',

packages/compass-preferences-model/src/provider.ts

@@ -1,5 +1,6 @@
 export * from './react';
 export { ReadOnlyPreferenceAccess } from './read-only-preferences-access';
+export { CompassWebPreferencesAccess } from './compass-web-preferences-access';
 export {
   useIsAIFeatureEnabled,
   isAIFeatureEnabled,

packages/compass-web/src/entrypoint.tsx

@@ -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';
@@ -262,7 +262,7 @@ const CompassWeb = ({
   });
 
   const preferencesAccess = useRef(
-    new ReadOnlyPreferenceAccess({
+    new CompassWebPreferencesAccess({
       maxTimeMS: 10_000,
       enableExplainPlan: true,
       enableAggregationBuilderRunPipeline: true,
@@ -279,6 +279,7 @@ const CompassWeb = ({
       enableShell: false,
       enableCreatingNewConnections: false,
       enableGlobalWrites: false,
+      optInDataExplorerGenAIFeatures: false,
       ...initialPreferences,
     })
   );