Skip to content

Commit

Permalink
feat(atlas-service): ping if ai feature enabled COMPASS-7193 (#4840)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Anemy
Anemy authored Sep 13, 2023
1 parent c5200cb commit 121cce6
Show file tree
Hide file tree
Showing 17 changed files with 343 additions and 52 deletions.
126 changes: 124 additions & 2 deletions packages/atlas-service/src/main.spec.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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];
});

Expand Down Expand Up @@ -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;
Expand All @@ -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();
});
Expand Down Expand Up @@ -407,6 +434,7 @@ describe('AtlasServiceMain', function () {
'getUserInfo',
'introspect',
'revoke',
'getAIFeatureEnablement',
'getAggregationFromUserInput',
'getQueryFromUserInput',
]) {
Expand Down Expand Up @@ -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);
});
});
});
79 changes: 74 additions & 5 deletions packages/atlas-service/src/main.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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');

Expand Down Expand Up @@ -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");
}
}
Expand Down Expand Up @@ -136,6 +144,8 @@ export class AtlasService {

private static currentUser: AtlasUserInfo | null = null;

private static getActiveCompassUser: typeof getActiveUser = getActiveUser;

private static signInPromise: Promise<AtlasUserInfo> | null = null;

private static fetch = (
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -524,6 +535,64 @@ export class AtlasService {
await throwIfNotOk(res);
}

static async getAIFeatureEnablement(): Promise<AIFeatureEnablement> {
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<void> {
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,
Expand Down
18 changes: 18 additions & 0 deletions packages/atlas-service/src/util.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<
Expand Down
6 changes: 3 additions & 3 deletions packages/compass-aggregations/src/components/pipeline-toolbar/index.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -74,7 +74,7 @@ export const PipelineToolbar: React.FunctionComponent<PipelineToolbarProps> = ({
pipelineOutputOption,
}) => {
const darkMode = useDarkMode();
const enableAIExperience = usePreference('enableAIExperience', React);
const isAIFeatureEnabled = useIsAIFeatureEnabled(React);
const [isOptionsVisible, setIsOptionsVisible] = useState(false);
return (
<div
Expand All @@ -99,7 +99,7 @@ export const PipelineToolbar: React.FunctionComponent<PipelineToolbarProps> = ({
<PipelineOptions />
</div>
)}
{enableAIExperience && isBuilderView && <PipelineAI />}
{isAIFeatureEnabled && isBuilderView && <PipelineAI />}
</div>
{isBuilderView ? (
<div className={settingsRowStyles}>
Expand Down
9 changes: 6 additions & 3 deletions ...compass-aggregations/src/components/pipeline-toolbar/pipeline-header/pipeline-actions.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down Expand Up @@ -82,11 +85,11 @@ export const PipelineActions: React.FunctionComponent<PipelineActionsProps> = ({
onCollectionScanInsightActionButtonClick,
}) => {
const showInsights = usePreference('showInsights', React);
const enableAIExperience = usePreference('enableAIExperience', React);
const isAIFeatureEnabled = useIsAIFeatureEnabled(React);

return (
<div className={containerStyles}>
{enableAIExperience && showAIEntry && (
{isAIFeatureEnabled && showAIEntry && (
<AIExperienceEntry onClick={onShowAIInputClick} />
)}
{showInsights && showCollectionScanInsight && (
Expand Down
Loading

0 comments on commit 121cce6

Please sign in to comment.