Skip to content

chore(gen-ai): remove gen ai access check, default always on COMPASS-9117 #6812

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 26, 2025
Merged

chore(gen-ai): remove gen ai access check, default always on COMPASS-9117 #6812

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 0 additions & 6 deletions packages/atlas-service/src/main.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,12 +33,6 @@ describe('CompassAuthServiceMain', function () {
'http://example.com/v1/revoke?client_id=1234abcd': {
ok: true,
},
'http://example.com/unauth/ai/api/v1/hello/': {
ok: true,
json() {
return { features: {} };
},
},
}[url];
});

Expand Down
7 changes: 2 additions & 5 deletions packages/compass-e2e-tests/tests/collection-ai-query.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,12 +121,9 @@ describe('Collection ai query', function () {

// Check that the request was made with the correct parameters.
const requests = getRequests();
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.length).to.equal(1);

const queryRequest = requests[1];
const queryRequest = requests[0];
const queryURL = new URL(
queryRequest.req.url,
`http://${queryRequest.req.headers.host}`
Expand Down
86 changes: 1 addition & 85 deletions packages/compass-generative-ai/src/atlas-ai-service.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,15 +77,13 @@ describe('AtlasAiService', function () {
{
apiURLPreset: 'admin-api',
expectedEndpoints: {
'user-access': 'http://example.com/unauth/ai/api/v1/hello/1234',
'mql-aggregation': `http://example.com/ai/api/v1/mql-aggregation?request_id=abc`,
'mql-query': `http://example.com/ai/api/v1/mql-query?request_id=abc`,
},
},
{
apiURLPreset: 'cloud',
expectedEndpoints: {
'user-access': '/cloud/ai/v1/hello/1234',
'mql-aggregation':
'/cloud/ai/v1/groups/testProject/mql-aggregation?request_id=abc',
'mql-query': '/cloud/ai/v1/groups/testProject/mql-query?request_id=abc',
Expand Down Expand Up @@ -293,7 +291,7 @@ describe('AtlasAiService', function () {
});
});

it('should set the cloudFeatureRolloutAccess true when returned true', async function () {
it('should set the cloudFeatureRolloutAccess true', async function () {
const fetchStub = sandbox.stub().resolves(
makeResponse({
features: {
Expand All @@ -311,94 +309,12 @@ describe('AtlasAiService', function () {

await atlasAiService['setupAIAccess']();

const { args } = fetchStub.firstCall;

expect(fetchStub).to.have.been.calledOnce;

expect(args[0]).to.equal(expectedEndpoints['user-access']);

currentCloudFeatureRolloutAccess =
preferences.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(
makeResponse({
features: {
GEN_AI_COMPASS: {
enabled: false,
},
},
})
);
global.fetch = fetchStub;

let currentCloudFeatureRolloutAccess =
preferences.getPreferences().cloudFeatureRolloutAccess;
expect(currentCloudFeatureRolloutAccess).to.equal(undefined);

await atlasAiService['setupAIAccess']();

const { args } = fetchStub.firstCall;

expect(fetchStub).to.have.been.calledOnce;
expect(args[0]).to.equal(expectedEndpoints['user-access']);

currentCloudFeatureRolloutAccess =
preferences.getPreferences().cloudFeatureRolloutAccess;
expect(currentCloudFeatureRolloutAccess).to.deep.equal({
GEN_AI_COMPASS: false,
});
});

it('should set the cloudFeatureRolloutAccess false when returned null', async function () {
const fetchStub = sandbox.stub().resolves(
makeResponse({
features: null,
})
);
global.fetch = fetchStub;

let currentCloudFeatureRolloutAccess =
preferences.getPreferences().cloudFeatureRolloutAccess;
expect(currentCloudFeatureRolloutAccess).to.equal(undefined);

await atlasAiService['setupAIAccess']();

const { args } = fetchStub.firstCall;

expect(fetchStub).to.have.been.calledOnce;
expect(args[0]).to.equal(expectedEndpoints['user-access']);

currentCloudFeatureRolloutAccess =
preferences.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'));
global.fetch = fetchStub;

let currentCloudFeatureRolloutAccess =
preferences.getPreferences().cloudFeatureRolloutAccess;
expect(currentCloudFeatureRolloutAccess).to.equal(undefined);

await atlasAiService['setupAIAccess']();

const { args } = fetchStub.firstCall;

expect(fetchStub).to.have.been.calledOnce;
expect(args[0]).to.equal(expectedEndpoints['user-access']);

currentCloudFeatureRolloutAccess =
preferences.getPreferences().cloudFeatureRolloutAccess;
expect(currentCloudFeatureRolloutAccess).to.deep.equal(undefined);
});
});
});
}
Expand Down
73 changes: 9 additions & 64 deletions packages/compass-generative-ai/src/atlas-ai-service.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import { AtlasServiceError } from '@mongodb-js/atlas-service/renderer';
import type { ConnectionInfo } from '@mongodb-js/compass-connections/provider';
import type { Document } from 'mongodb';
import type { Logger } from '@mongodb-js/compass-logging';
import { EJSON, UUID } from 'bson';
import { EJSON } from 'bson';
import { signIntoAtlasWithModalPrompt } from './store/atlas-signin-reducer';
import { getStore } from './store/atlas-ai-store';
import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer';
Expand Down Expand Up @@ -198,17 +198,15 @@ const aiURLConfig = {
// Down the line we'd like to only use the admin api, however,
// we cannot currently call that from the Atlas UI. Pending CLOUDP-251201
'admin-api': {
'user-access': (userId: string) => `unauth/ai/api/v1/hello/${userId}`,
aggregation: 'ai/api/v1/mql-aggregation',
query: 'ai/api/v1/mql-query',
},
cloud: {
'user-access': (userId: string) => `ai/v1/hello/${userId}`,
aggregation: (groupId: string) => `ai/v1/groups/${groupId}/mql-aggregation`,
query: (groupId: string) => `ai/v1/groups/${groupId}/mql-query`,
},
} as const;
type AIEndpoint = 'user-access' | 'query' | 'aggregation';
type AIEndpoint = 'query' | 'aggregation';

export class AtlasAiService {
private initPromise: Promise<void> | null = null;
Expand Down Expand Up @@ -242,15 +240,6 @@ export class AtlasAiService {
connectionInfo?: ConnectionInfo
) {
if (this.apiURLPreset === 'cloud') {
if (urlId === 'user-access') {
return this.atlasService.cloudEndpoint(
aiURLConfig[this.apiURLPreset][urlId](
this.preferences.getPreferences().telemetryAtlasUserId ??
new UUID().toString()
)
);
}

const atlasMetadata = connectionInfo?.atlasMetadata;
if (!atlasMetadata) {
throw new Error(
Expand All @@ -262,15 +251,7 @@ export class AtlasAiService {
aiURLConfig[this.apiURLPreset][urlId](atlasMetadata.projectId)
);
}
const urlConfig = aiURLConfig[this.apiURLPreset][urlId];
const urlPath =
typeof urlConfig === 'function'
? urlConfig(
this.preferences.getPreferences().telemetryAtlasUserId ??
new UUID().toString()
)
: urlConfig;

const urlPath = aiURLConfig[this.apiURLPreset][urlId];
return this.atlasService.adminApiEndpoint(urlPath);
}

Expand All @@ -285,50 +266,14 @@ export class AtlasAiService {
}
}

private async getAIFeatureEnablement(): Promise<AIFeatureEnablement> {
const url = this.getUrlForEndpoint('user-access');

const res = await this.atlasService.fetch(url, {
headers: {
Accept: 'application/json',
async setupAIAccess(): Promise<void> {
// We default GEN_AI_ACCESS on for everyone. Down the line if/when
// we add more features with partial rollout, we'll fetch access here.
await this.preferences.savePreferences({
cloudFeatureRolloutAccess: {
GEN_AI_COMPASS: true,
},
});
const body = await res.json();
this.validateAIFeatureEnablementResponse(body);
return body;
}

async setupAIAccess(): Promise<void> {
try {
const featureResponse = await this.getAIFeatureEnablement();

const isAIFeatureEnabled =
!!featureResponse?.features?.GEN_AI_COMPASS?.enabled;

this.logger.log.info(
this.logger.mongoLogId(1_001_000_300),
'AtlasAIService',
'Fetched if the AI feature is enabled',
{
enabled: isAIFeatureEnabled,
featureResponse,
}
);

await this.preferences.savePreferences({
cloudFeatureRolloutAccess: {
GEN_AI_COMPASS: isAIFeatureEnabled,
},
});
} catch (err) {
// Default to what's already in Compass when we can't fetch the preference.
this.logger.log.error(
this.logger.mongoLogId(1_001_000_302),
'AtlasAIService',
'Failed to load if the AI feature is enabled',
{ error: (err as Error).stack }
);
}
}

async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) {
Expand Down
Loading