From 121cce60246567df0da667d593f99b331f2dde41 Mon Sep 17 00:00:00 2001 From: Rhys Date: Wed, 13 Sep 2023 12:33:11 -0400 Subject: [PATCH] feat(atlas-service): ping if ai feature enabled COMPASS-7193 (#4840) --- packages/atlas-service/src/main.spec.ts | 126 +++++++++++++++++- packages/atlas-service/src/main.ts | 79 ++++++++++- packages/atlas-service/src/util.ts | 18 +++ .../src/components/pipeline-toolbar/index.tsx | 6 +- .../pipeline-header/pipeline-actions.tsx | 9 +- .../pipeline-header/pipeline-stages.tsx | 8 +- .../helpers/atlas-service.ts | 25 ++++ .../tests/collection-ai-query.test.ts | 13 +- .../src/feature-flags.ts | 6 +- .../compass-preferences-model/src/index.ts | 1 + .../src/preferences.ts | 20 +++ .../compass-preferences-model/src/react.ts | 2 +- .../compass-preferences-model/src/utils.ts | 13 +- .../src/components/query-bar.spec.tsx | 18 ++- .../src/components/query-bar.tsx | 17 ++- .../src/components/settings/atlas-login.tsx | 6 +- .../components/settings/feature-preview.tsx | 28 ++-- 17 files changed, 343 insertions(+), 52 deletions(-) diff --git a/packages/atlas-service/src/main.spec.ts b/packages/atlas-service/src/main.spec.ts index 3bb5e1ba21f..e7413df0ec3 100644 --- a/packages/atlas-service/src/main.spec.ts +++ b/packages/atlas-service/src/main.spec.ts @@ -3,6 +3,7 @@ import { expect } from 'chai'; import { AtlasService, getTrackingUserInfo, throwIfNotOk } from './main'; import { EventEmitter } from 'events'; import preferencesAccess from 'compass-preferences-model'; +import type { UserPreferences } from 'compass-preferences-model'; import type { AtlasUserConfigStore } from './user-config-store'; import type { AtlasUserInfo } from './util'; @@ -26,6 +27,12 @@ describe('AtlasServiceMain', function () { 'http://example.com/v1/revoke?client_id=1234abcd': { ok: true, }, + 'http://example.com/ai/api/v1/hello/': { + ok: true, + json() { + return { features: {} }; + }, + }, }[url]; }); @@ -63,21 +70,38 @@ describe('AtlasServiceMain', function () { const ipcMain = AtlasService['ipcMain']; const createPlugin = AtlasService['createMongoDBOIDCPlugin']; const userStore = AtlasService['atlasUserConfigStore']; + const getActiveCompassUser = AtlasService['getActiveCompassUser']; - beforeEach(function () { + let cloudFeatureRolloutAccess: UserPreferences['cloudFeatureRolloutAccess']; + + beforeEach(async function () { AtlasService['ipcMain'] = { handle: sandbox.stub() }; AtlasService['fetch'] = mockFetch as any; AtlasService['createMongoDBOIDCPlugin'] = () => mockOidcPlugin; AtlasService['atlasUserConfigStore'] = mockUserConfigStore as unknown as AtlasUserConfigStore; + AtlasService['getActiveCompassUser'] = () => + Promise.resolve({ + id: 'test', + createdAt: new Date(), + lastUsed: new Date(), + }); AtlasService['config'] = defaultConfig; AtlasService['setupPlugin'](); AtlasService['attachOidcPluginLoggerEvents'](); + + cloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + await preferencesAccess.savePreferences({ + cloudFeatureRolloutAccess: { + GEN_AI_COMPASS: true, + }, + }); }); - afterEach(function () { + afterEach(async function () { AtlasService['fetch'] = fetch; AtlasService['atlasUserConfigStore'] = userStore; AtlasService['ipcMain'] = ipcMain; @@ -86,6 +110,9 @@ describe('AtlasServiceMain', function () { AtlasService['oidcPluginLogger'].removeAllListeners(); AtlasService['signInPromise'] = null; AtlasService['currentUser'] = null; + AtlasService['getActiveCompassUser'] = getActiveCompassUser; + + await preferencesAccess.savePreferences({ cloudFeatureRolloutAccess }); sandbox.resetHistory(); }); @@ -407,6 +434,7 @@ describe('AtlasServiceMain', function () { 'getUserInfo', 'introspect', 'revoke', + 'getAIFeatureEnablement', 'getAggregationFromUserInput', 'getQueryFromUserInput', ]) { @@ -469,4 +497,98 @@ describe('AtlasServiceMain', function () { }); }); }); + + describe('setupAIAccess', function () { + beforeEach(async function () { + await preferencesAccess.savePreferences({ + cloudFeatureRolloutAccess: undefined, + }); + }); + + it('should set the cloudFeatureRolloutAccess true when returned true', async function () { + const fetchStub = sandbox.stub().resolves({ + ok: true, + json() { + return Promise.resolve({ + features: { + GEN_AI_COMPASS: { + enabled: true, + }, + }, + }); + }, + }); + AtlasService['fetch'] = fetchStub; + + let currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.equal(undefined); + + await AtlasService.setupAIAccess(); + + const { args } = fetchStub.getCall(0); + + expect(AtlasService['fetch']).to.have.been.calledOnce; + expect(args[0]).to.eq(`http://example.com/ai/api/v1/hello/test`); + + currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.deep.equal({ + GEN_AI_COMPASS: true, + }); + }); + + it('should set the cloudFeatureRolloutAccess false when returned false', async function () { + const fetchStub = sandbox.stub().resolves({ + ok: true, + json() { + return Promise.resolve({ + features: { + GEN_AI_COMPASS: { + enabled: false, + }, + }, + }); + }, + }); + AtlasService['fetch'] = fetchStub; + + let currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.equal(undefined); + + await AtlasService.setupAIAccess(); + + const { args } = fetchStub.getCall(0); + + expect(AtlasService['fetch']).to.have.been.calledOnce; + expect(args[0]).to.eq(`http://example.com/ai/api/v1/hello/test`); + + currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.deep.equal({ + GEN_AI_COMPASS: false, + }); + }); + + it('should not set the cloudFeatureRolloutAccess false when returned false', async function () { + const fetchStub = sandbox.stub().throws(new Error('error')); + AtlasService['fetch'] = fetchStub; + + let currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.equal(undefined); + + await AtlasService.setupAIAccess(); + + const { args } = fetchStub.getCall(0); + + expect(AtlasService['fetch']).to.have.been.calledOnce; + expect(args[0]).to.eq(`http://example.com/ai/api/v1/hello/test`); + + currentCloudFeatureRolloutAccess = + preferencesAccess.getPreferences().cloudFeatureRolloutAccess; + expect(currentCloudFeatureRolloutAccess).to.deep.equal(undefined); + }); + }); }); diff --git a/packages/atlas-service/src/main.ts b/packages/atlas-service/src/main.ts index 074d982cd9d..5abca83a154 100644 --- a/packages/atlas-service/src/main.ts +++ b/packages/atlas-service/src/main.ts @@ -18,11 +18,16 @@ import type { Document } from 'mongodb'; import type { AtlasUserConfig, AIAggregation, + AIFeatureEnablement, AIQuery, IntrospectInfo, AtlasUserInfo, } from './util'; -import { validateAIQueryResponse, validateAIAggregationResponse } from './util'; +import { + validateAIQueryResponse, + validateAIAggregationResponse, + validateAIFeatureEnablementResponse, +} from './util'; import { broadcast, ipcExpose, @@ -36,6 +41,7 @@ import preferences from 'compass-preferences-model'; import { SecretStore, SECRET_STORE_KEY } from './secret-store'; import { AtlasUserConfigStore } from './user-config-store'; import { OidcPluginLogger } from './oidc-plugin-logger'; +import { getActiveUser } from 'compass-preferences-model'; const { log, track } = createLoggerAndTelemetry('COMPASS-ATLAS-SERVICE'); @@ -87,12 +93,14 @@ export async function throwIfNotOk( } function throwIfAINotEnabled(atlasService: typeof AtlasService) { + if (!preferences.getPreferences().cloudFeatureRolloutAccess?.GEN_AI_COMPASS) { + throw new Error( + "Compass' AI functionality is not currently enabled. Please try again later." + ); + } // Only throw if we actually have userInfo / logged in. Otherwise allow // request to fall through so that we can get a proper network error - if ( - atlasService['currentUser'] && - atlasService['currentUser'].enabledAIFeature === false - ) { + if (atlasService['currentUser']?.enabledAIFeature === false) { throw new Error("Can't use AI before accepting terms and conditions"); } } @@ -136,6 +144,8 @@ export class AtlasService { private static currentUser: AtlasUserInfo | null = null; + private static getActiveCompassUser: typeof getActiveUser = getActiveUser; + private static signInPromise: Promise | null = null; private static fetch = ( @@ -230,6 +240,7 @@ export class AtlasService { ); const serializedState = await this.secretStore.getItem(SECRET_STORE_KEY); this.setupPlugin(serializedState); + await this.setupAIAccess(); // Whether or not we got the state, try requesting user info. If there was // no serialized state returned, this will just fail quickly. If there was // some state, we will prepare the service state for user interactions by @@ -524,6 +535,64 @@ export class AtlasService { await throwIfNotOk(res); } + static async getAIFeatureEnablement(): Promise { + throwIfNetworkTrafficDisabled(); + + const userId = (await this.getActiveCompassUser()).id; + + const res = await this.fetch( + `${this.config.atlasApiBaseUrl}/ai/api/v1/hello/${userId}` + ); + + await throwIfNotOk(res); + + const body = await res.json(); + + validateAIFeatureEnablementResponse(body); + + return body; + } + + static async setupAIAccess(): Promise { + log.info( + mongoLogId(1_001_000_227), + 'AtlasService', + 'Fetching if the AI feature is enabled' + ); + + try { + throwIfNetworkTrafficDisabled(); + + const featureResponse = await this.getAIFeatureEnablement(); + + const isAIFeatureEnabled = + !!featureResponse.features.GEN_AI_COMPASS?.enabled; + + log.info( + mongoLogId(1_001_000_229), + 'AtlasService', + 'Fetched if the AI feature is enabled', + { + enabled: isAIFeatureEnabled, + } + ); + + await preferences.savePreferences({ + cloudFeatureRolloutAccess: { + GEN_AI_COMPASS: isAIFeatureEnabled, + }, + }); + } catch (err) { + // Default to what's already in Compass when we can't fetch the preference. + log.error( + mongoLogId(1_001_000_243), + 'AtlasService', + 'Failed to load if the AI feature is enabled', + { error: (err as Error).stack } + ); + } + } + static async getAggregationFromUserInput({ signal, userInput, diff --git a/packages/atlas-service/src/util.ts b/packages/atlas-service/src/util.ts index 55c7debf523..7979d109af4 100644 --- a/packages/atlas-service/src/util.ts +++ b/packages/atlas-service/src/util.ts @@ -72,6 +72,24 @@ export function validateAIAggregationResponse( } } +export type AIFeatureEnablement = { + features: { + [featureName: string]: { + enabled: boolean; + }; + }; +}; + +export function validateAIFeatureEnablementResponse( + response: any +): asserts response is AIFeatureEnablement { + const { features } = response; + + if (typeof features !== 'object' || features === null) { + throw new Error('Unexpected response: expected features to be an object'); + } +} + export type AIQuery = { content: { query: Record< diff --git a/packages/compass-aggregations/src/components/pipeline-toolbar/index.tsx b/packages/compass-aggregations/src/components/pipeline-toolbar/index.tsx index c03a9d6736f..ee9948c4539 100644 --- a/packages/compass-aggregations/src/components/pipeline-toolbar/index.tsx +++ b/packages/compass-aggregations/src/components/pipeline-toolbar/index.tsx @@ -7,7 +7,7 @@ import { useDarkMode, } from '@mongodb-js/compass-components'; import { connect } from 'react-redux'; -import { usePreference } from 'compass-preferences-model'; +import { useIsAIFeatureEnabled } from 'compass-preferences-model'; import PipelineHeader from './pipeline-header'; import PipelineOptions from './pipeline-options'; @@ -74,7 +74,7 @@ export const PipelineToolbar: React.FunctionComponent = ({ pipelineOutputOption, }) => { const darkMode = useDarkMode(); - const enableAIExperience = usePreference('enableAIExperience', React); + const isAIFeatureEnabled = useIsAIFeatureEnabled(React); const [isOptionsVisible, setIsOptionsVisible] = useState(false); return (
= ({
)} - {enableAIExperience && isBuilderView && } + {isAIFeatureEnabled && isBuilderView && } {isBuilderView ? (
diff --git a/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-actions.tsx b/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-actions.tsx index dd7f9cbc465..77b2b9982dc 100644 --- a/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-actions.tsx +++ b/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-actions.tsx @@ -22,7 +22,10 @@ import { } from '../../../modules/pipeline-builder/builder-helpers'; import { isOutputStage } from '../../../utils/stage'; import { openCreateIndexModal } from '../../../modules/insights'; -import { usePreference } from 'compass-preferences-model'; +import { + usePreference, + useIsAIFeatureEnabled, +} from 'compass-preferences-model'; import { showInput as showAIInput } from '../../../modules/pipeline-builder/pipeline-ai'; const containerStyles = css({ @@ -82,11 +85,11 @@ export const PipelineActions: React.FunctionComponent = ({ onCollectionScanInsightActionButtonClick, }) => { const showInsights = usePreference('showInsights', React); - const enableAIExperience = usePreference('enableAIExperience', React); + const isAIFeatureEnabled = useIsAIFeatureEnabled(React); return (
- {enableAIExperience && showAIEntry && ( + {isAIFeatureEnabled && showAIEntry && ( )} {showInsights && showCollectionScanInsight && ( diff --git a/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-stages.tsx b/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-stages.tsx index fd83f8fc5c7..d1413b8eb9f 100644 --- a/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-stages.tsx +++ b/packages/compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-stages.tsx @@ -11,7 +11,7 @@ import { Button, Icon, } from '@mongodb-js/compass-components'; -import { usePreference } from 'compass-preferences-model'; +import { useIsAIFeatureEnabled } from 'compass-preferences-model'; import type { RootState } from '../../../modules'; import { editPipeline } from '../../../modules/workspace'; @@ -61,7 +61,7 @@ export const PipelineStages: React.FunctionComponent = ({ onEditPipelineClick, onShowAIInputClick, }) => { - const enableAIExperience = usePreference('enableAIExperience', React); + const isAIFeatureEnabled = useIsAIFeatureEnabled(React); return (
@@ -70,7 +70,7 @@ export const PipelineStages: React.FunctionComponent = ({ Your pipeline is currently empty. {showAddNewStage && ( <> - {enableAIExperience && showAIEntry ? ( + {isAIFeatureEnabled && showAIEntry ? ( <>{nbsp}Need help getting started? ) : ( <> @@ -88,7 +88,7 @@ export const PipelineStages: React.FunctionComponent = ({ )} )} - {enableAIExperience && showAIEntry && ( + {isAIFeatureEnabled && showAIEntry && ( <> {nbsp} diff --git a/packages/compass-e2e-tests/helpers/atlas-service.ts b/packages/compass-e2e-tests/helpers/atlas-service.ts index c96bbcd88c6..cad0687736c 100644 --- a/packages/compass-e2e-tests/helpers/atlas-service.ts +++ b/packages/compass-e2e-tests/helpers/atlas-service.ts @@ -7,6 +7,23 @@ export type MockAtlasServerResponse = { body: any; }; +function aiFeatureEnableResponse( + req: http.IncomingMessage, + res: http.ServerResponse +) { + // Get request to hello service. + res.setHeader('Content-Type', 'application/json'); + return res.end( + JSON.stringify({ + features: { + GEN_AI_COMPASS: { + enabled: true, + }, + }, + }) + ); +} + export async function startMockAtlasServiceServer( { response: _response, @@ -44,6 +61,14 @@ export async function startMockAtlasServiceServer( let response = _response; const server = http .createServer((req, res) => { + if (req.method === 'GET') { + requests.push({ + req, + content: null, + }); + return aiFeatureEnableResponse(req, res); + } + let body = ''; req .setEncoding('utf8') diff --git a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts index c958184b87b..beab005bdb6 100644 --- a/packages/compass-e2e-tests/tests/collection-ai-query.test.ts +++ b/packages/compass-e2e-tests/tests/collection-ai-query.test.ts @@ -110,11 +110,14 @@ describe('Collection ai query', function () { // Check that the request was made with the correct parameters. const requests = getRequests(); - expect(requests.length).to.equal(1); - expect(requests[0].content.userInput).to.equal(testUserInput); - expect(requests[0].content.collectionName).to.equal('numbers'); - expect(requests[0].content.databaseName).to.equal('test'); - expect(requests[0].content.schema).to.exist; + expect(requests.length).to.equal(2); + const lastPathRegex = /[^/]*$/; + const userId = lastPathRegex.exec(requests[0].req.url)?.[0]; + expect((userId?.match(/-/g) || []).length).to.equal(4); // Is uuid like. + expect(requests[1].content.userInput).to.equal(testUserInput); + expect(requests[1].content.collectionName).to.equal('numbers'); + expect(requests[1].content.databaseName).to.equal('test'); + expect(requests[1].content.schema).to.exist; // Run it and check that the correct documents are shown. await browser.runFind('Documents', true); diff --git a/packages/compass-preferences-model/src/feature-flags.ts b/packages/compass-preferences-model/src/feature-flags.ts index f6c12c3a842..38a8a382f55 100644 --- a/packages/compass-preferences-model/src/feature-flags.ts +++ b/packages/compass-preferences-model/src/feature-flags.ts @@ -31,10 +31,10 @@ export const featureFlags: Required<{ * Epic: COMPASS-6866 */ enableAIExperience: { - stage: 'development', + stage: 'preview', description: { - short: 'AI Query Generator', - long: 'Use AI to generate queries with a natural language text input on the query bar. Do not use this feature with sensitive data.', + short: 'Compass AI Features', + long: 'Use AI to generate queries and aggregations with a natural language text. Do not use this feature with sensitive data.', }, }, diff --git a/packages/compass-preferences-model/src/index.ts b/packages/compass-preferences-model/src/index.ts index d4b20d54553..9c7dc753777 100644 --- a/packages/compass-preferences-model/src/index.ts +++ b/packages/compass-preferences-model/src/index.ts @@ -28,6 +28,7 @@ export { capMaxTimeMSAtPreferenceLimit, setupPreferencesAndUser, getActiveUser, + useIsAIFeatureEnabled, } from './utils'; export type { User } from './storage'; diff --git a/packages/compass-preferences-model/src/preferences.ts b/packages/compass-preferences-model/src/preferences.ts index cffc742c6e6..b5b2c47c82f 100644 --- a/packages/compass-preferences-model/src/preferences.ts +++ b/packages/compass-preferences-model/src/preferences.ts @@ -54,6 +54,9 @@ export type InternalUserPreferences = { // by users. showedNetworkOptIn: boolean; // Has the settings dialog been shown before. id: string; + cloudFeatureRolloutAccess?: { + GEN_AI_COMPASS?: boolean; + }; lastKnownVersion: string; currentUserId?: string; telemetryAnonymousId?: string; @@ -349,6 +352,23 @@ export const storedUserPreferencesProps: Required<{ validator: z.string().uuid().optional(), type: 'string', }, + /** + * Enable/disable the AI services. This is currently set + * in the atlas-service initialization where we make a request to the + * ai endpoint to check what's enabled for the user (incremental rollout). + */ + cloudFeatureRolloutAccess: { + ui: false, + cli: false, + global: false, + description: null, + validator: z + .object({ + GEN_AI_COMPASS: z.boolean().optional(), + }) + .optional(), + type: 'object', + }, /** * Master switch to disable all network traffic * and make Compass behave like Isolated edition always, diff --git a/packages/compass-preferences-model/src/react.ts b/packages/compass-preferences-model/src/react.ts index 94a051ad898..117d0095b8d 100644 --- a/packages/compass-preferences-model/src/react.ts +++ b/packages/compass-preferences-model/src/react.ts @@ -1,7 +1,7 @@ import preferencesAccess from './'; import type { AllPreferences } from './'; -interface ReactHooks { +export interface ReactHooks { useState: (initialValue: T) => [T, (newValue: T) => void]; useEffect: ( effectFn: () => () => void, diff --git a/packages/compass-preferences-model/src/utils.ts b/packages/compass-preferences-model/src/utils.ts index 794c81bd94a..6edc649bc86 100644 --- a/packages/compass-preferences-model/src/utils.ts +++ b/packages/compass-preferences-model/src/utils.ts @@ -1,7 +1,8 @@ -import preferences, { preferencesAccess } from '.'; +import preferences, { preferencesAccess, usePreference } from '.'; import type { ParsedGlobalPreferencesResult } from '.'; import { setupPreferences } from './setup-preferences'; import { UserStorage } from './storage'; +import type { ReactHooks } from './react'; export async function setupPreferencesAndUser( globalPreferences: ParsedGlobalPreferencesResult @@ -39,3 +40,13 @@ export function capMaxTimeMSAtPreferenceLimit(value: T): T | number { } return value; } + +export function useIsAIFeatureEnabled(React: ReactHooks) { + const enableAIExperience = usePreference('enableAIExperience', React); + const isAIFeatureEnabled = usePreference( + 'cloudFeatureRolloutAccess', + React + )?.GEN_AI_COMPASS; + + return enableAIExperience && isAIFeatureEnabled; +} diff --git a/packages/compass-query-bar/src/components/query-bar.spec.tsx b/packages/compass-query-bar/src/components/query-bar.spec.tsx index 3dfc854b71b..6b39e6f7464 100644 --- a/packages/compass-query-bar/src/components/query-bar.spec.tsx +++ b/packages/compass-query-bar/src/components/query-bar.spec.tsx @@ -124,9 +124,12 @@ describe('QueryBar Component', function () { beforeEach(function () { sandbox = sinon.createSandbox(); - sandbox - .stub(preferencesAccess, 'getPreferences') - .returns({ enableAIExperience: true } as any); + sandbox.stub(preferencesAccess, 'getPreferences').returns({ + enableAIExperience: true, + cloudFeatureRolloutAccess: { + GEN_AI_COMPASS: true, + }, + } as any); }); afterEach(function () { @@ -173,9 +176,12 @@ describe('QueryBar Component', function () { beforeEach(function () { sandbox = sinon.createSandbox(); - sandbox - .stub(preferencesAccess, 'getPreferences') - .returns({ enableAIExperience: false } as any); + sandbox.stub(preferencesAccess, 'getPreferences').returns({ + enableAIExperience: false, + cloudFeatureRolloutAccess: { + GEN_AI_COMPASS: true, + }, + } as any); renderQueryBar({ queryOptionsLayout: ['filter'], }); diff --git a/packages/compass-query-bar/src/components/query-bar.tsx b/packages/compass-query-bar/src/components/query-bar.tsx index 7d2230bdb7c..c88d40c6ab4 100644 --- a/packages/compass-query-bar/src/components/query-bar.tsx +++ b/packages/compass-query-bar/src/components/query-bar.tsx @@ -15,7 +15,10 @@ import { GuideCue, } from '@mongodb-js/compass-components'; import { connect } from 'react-redux'; -import { usePreference } from 'compass-preferences-model'; +import { + usePreference, + useIsAIFeatureEnabled, +} from 'compass-preferences-model'; import type { Signal } from '@mongodb-js/compass-components'; import { @@ -175,7 +178,7 @@ export const QueryBar: React.FunctionComponent = ({ }) => { const darkMode = useDarkMode(); const newExplainPlan = usePreference('newExplainPlan', React); - const enableAIQuery = usePreference('enableAIExperience', React); + const isAIFeatureEnabled = useIsAIFeatureEnabled(React); const onFormSubmit = useCallback( (evt: React.FormEvent) => { @@ -188,7 +191,7 @@ export const QueryBar: React.FunctionComponent = ({ const filterQueryOptionId = 'query-bar-option-input-filter'; const filterPlaceholder = useMemo(() => { - return enableAIQuery && !isAIInputVisible + return isAIFeatureEnabled && !isAIInputVisible ? createAIPlaceholderHTMLPlaceholder({ onClickAI: () => { onShowAIInputClick(); @@ -198,7 +201,7 @@ export const QueryBar: React.FunctionComponent = ({ }) : placeholders?.filter; }, [ - enableAIQuery, + isAIFeatureEnabled, isAIInputVisible, darkMode, placeholders?.filter, @@ -206,13 +209,13 @@ export const QueryBar: React.FunctionComponent = ({ ]); const showAskAIButton = useMemo(() => { - if (!enableAIQuery || isAIInputVisible) { + if (isAIInputVisible || !isAIFeatureEnabled) { return false; } // See if there is content in the filter. return filterHasContent; - }, [enableAIQuery, isAIInputVisible, filterHasContent]); + }, [isAIFeatureEnabled, isAIInputVisible, filterHasContent]); return (
= ({ ))}
)} - {enableAIQuery && ( + {isAIFeatureEnabled && (
{ diff --git a/packages/compass-settings/src/components/settings/atlas-login.tsx b/packages/compass-settings/src/components/settings/atlas-login.tsx index 766c32c55f1..448e589e967 100644 --- a/packages/compass-settings/src/components/settings/atlas-login.tsx +++ b/packages/compass-settings/src/components/settings/atlas-login.tsx @@ -110,7 +110,7 @@ export const AtlasLoginSettings: React.FunctionComponent<{ color={darkMode ? 'white' : 'black'} height={18} > - AI Query + AI Query and Aggregation Preview
@@ -166,8 +166,8 @@ export const AtlasLoginSettings: React.FunctionComponent<{ htmlFor="use-ai-toggle" className={atlasLoginToggleControlLabelStyles} > - Use AI to generate queries with a natural language text input on the - query bar + Use AI to generate queries and aggregations with a natural language + text input on the query bar and aggregation page.
{isSignedIn && !isAIFeatureEnabled && ( diff --git a/packages/compass-settings/src/components/settings/feature-preview.tsx b/packages/compass-settings/src/components/settings/feature-preview.tsx index 78e314919fd..c458a916754 100644 --- a/packages/compass-settings/src/components/settings/feature-preview.tsx +++ b/packages/compass-settings/src/components/settings/feature-preview.tsx @@ -5,6 +5,7 @@ import { featureFlags, withPreferences, } from 'compass-preferences-model'; +import type { UserPreferences } from 'compass-preferences-model'; import { connect } from 'react-redux'; import type { RootState } from '../../stores'; import { ConnectedAtlasLoginSettings } from './atlas-login'; @@ -83,14 +84,23 @@ export const FeaturePreviewSettings: React.FunctionComponent<{ }; export default withPreferences( - connect((state: RootState, ownProps: { enableAIExperience?: boolean }) => { - return { - showAtlasLoginSettings: - state.settings.settings.enableAIExperience || - ['authenticated', 'in-progress'].includes(state.atlasLogin.status) || - ownProps.enableAIExperience, - }; - })(FeaturePreviewSettings), - ['enableAIExperience'], + connect( + ( + state: RootState, + ownProps: { + enableAIExperience?: boolean; + cloudFeatureRolloutAccess?: UserPreferences['cloudFeatureRolloutAccess']; + } + ) => { + return { + showAtlasLoginSettings: + state.settings.settings.enableAIExperience || + ['authenticated', 'in-progress'].includes(state.atlasLogin.status) || + ownProps.enableAIExperience || + ownProps.cloudFeatureRolloutAccess?.GEN_AI_COMPASS, + }; + } + )(FeaturePreviewSettings), + ['enableAIExperience', 'cloudFeatureRolloutAccess'], React );