chore: refactor to use atlas service and other feedback changes

Author: gagik

package-lock.json

@@ -67,6 +67,9 @@ export class AtlasService {
   cloudEndpoint(path?: string): string {
     return `${this.config.cloudBaseUrl}${normalizePath(path)}`;
   }
+  assistantApiEndpoint(path?: string): string {
+    return `${this.config.assistantApiBaseUrl}${normalizePath(path)}`;
+  }
   regionalizedCloudEndpoint(
     _atlasMetadata: Pick<AtlasClusterMetadata, 'regionalBaseUrl'>,
     path?: string

packages/atlas-service/src/atlas-service.ts

@@ -116,6 +116,10 @@ export type AtlasServiceConfig = {
    * Atlas Account Portal UI base url
    */
   authPortalUrl: string;
+  /**
+   * Assistant API base url
+   */
+  assistantApiBaseUrl: string;
 };
 
 /**
@@ -139,6 +143,7 @@ const config = {
       issuer: 'https://auth-qa.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account-dev.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
   'atlas-dev': {
     wsBaseUrl: '',
@@ -149,6 +154,7 @@ const config = {
       issuer: 'https://auth-qa.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account-dev.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
   'atlas-qa': {
     wsBaseUrl: '',
@@ -159,6 +165,7 @@ const config = {
       issuer: 'https://auth-qa.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account-qa.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
   atlas: {
     wsBaseUrl: '',
@@ -169,6 +176,7 @@ const config = {
       issuer: 'https://auth.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.corp.mongodb.com/api/v1',
   },
   'web-sandbox-atlas-local': {
     wsBaseUrl: '/ccs',
@@ -189,6 +197,7 @@ const config = {
       issuer: 'https://auth-qa.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account-dev.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
   'web-sandbox-atlas-qa': {
     wsBaseUrl: '/ccs',
@@ -199,6 +208,7 @@ const config = {
       issuer: 'https://auth-qa.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account-dev.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
   'web-sandbox-atlas': {
     wsBaseUrl: '/ccs',
@@ -209,6 +219,7 @@ const config = {
       issuer: 'https://auth.mongodb.com/oauth2/default',
     },
     authPortalUrl: 'https://account.mongodb.com/account/login',
+    assistantApiBaseUrl: 'https://knowledge.staging.corp.mongodb.com/api/v1',
   },
 } as const;
 
@@ -223,6 +234,7 @@ export function getAtlasConfig(
       issuer: process.env.COMPASS_OIDC_ISSUER_OVERRIDE,
     },
     authPortalUrl: process.env.COMPASS_ATLAS_AUTH_PORTAL_URL_OVERRIDE,
+    assistantApiBaseUrl: process.env.COMPASS_ASSISTANT_BASE_URL_OVERRIDE,
   };
   return defaultsDeep(
     envConfig,

packages/atlas-service/src/util.ts

@@ -50,6 +50,7 @@
   },
   "dependencies": {
     "@ai-sdk/openai": "^2.0.4",
+    "@mongodb-js/atlas-service": "^0.54.0",
     "@mongodb-js/compass-components": "^1.30.5",
     "@mongodb-js/compass-logging": "^1.7.10",
     "@mongodb-js/compass-app-registry": "^9.4.18",

packages/compass-assistant/package.json

@@ -8,7 +8,6 @@ import { usePreference } from 'compass-preferences-model/provider';
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
 import { createContext, useContext } from 'react';
-import { docsProviderTransport } from './docs-provider-transport';
 
 type AssistantActions = unknown;
 
@@ -20,21 +19,19 @@ export function useAssistantActions(): AssistantActions {
 
 export const AssistantProvider: React.FunctionComponent<
   PropsWithChildren<{
-    chat?: Chat<UIMessage>;
+    chat: Chat<UIMessage>;
   }>
 > = ({ chat, children }) => {
   const enableAIAssistant = usePreference('enableAIAssistant');
 
   const { messages, sendMessage } = useChat({
-    transport: docsProviderTransport,
     chat,
   });
 
   const contextActions = useRef({});
 
   const handleMessageSend = useCallback(
     (messageBody: string) => {
-      /** Telemetry, etc. */
       void sendMessage({ text: messageBody });
     },
     [sendMessage]
@@ -46,19 +43,14 @@ export const AssistantProvider: React.FunctionComponent<
 
   return (
     <AssistantActionsContext.Provider value={contextActions.current}>
-      {enableAIAssistant && (
-        <DrawerSection
-          id={ASSISTANT_DRAWER_ID}
-          title="MongoDB Assistant"
-          label="MongoDB Assistant"
-          glyph="Sparkle"
-        >
-          <AssistantChat
-            messages={messages}
-            onSendMessage={handleMessageSend}
-          />
-        </DrawerSection>
-      )}
+      <DrawerSection
+        id={ASSISTANT_DRAWER_ID}
+        title="MongoDB Assistant"
+        label="MongoDB Assistant"
+        glyph="Sparkle"
+      >
+        <AssistantChat messages={messages} onSendMessage={handleMessageSend} />
+      </DrawerSection>
       {children}
     </AssistantActionsContext.Provider>
   );

packages/compass-assistant/src/assistant-provider.tsx

@@ -8,24 +8,25 @@ import {
 } from 'ai';
 import { createOpenAI } from '@ai-sdk/openai';
 
-const openai = createOpenAI({
-  baseURL: 'https://knowledge.staging.corp.mongodb.com/api/v1',
-  apiKey: '',
-});
-
 export class DocsProviderTransport implements ChatTransport<UIMessage> {
+  private openai: ReturnType<typeof createOpenAI>;
+
+  constructor({ baseUrl }: { baseUrl: string }) {
+    this.openai = createOpenAI({
+      baseURL: baseUrl,
+      apiKey: '',
+    });
+  }
+
   // eslint-disable-next-line @typescript-eslint/require-await
   async sendMessages({
     messages,
     abortSignal,
   }: Parameters<ChatTransport<UIMessage>['sendMessages']>[0]) {
     const result = streamText({
-      model: openai.responses('mongodb-chat-latest'),
+      model: this.openai.responses('mongodb-chat-latest'),
       messages: convertToModelMessages(messages),
       abortSignal: abortSignal,
-      headers: {
-        origin: 'https://knowledge.staging.corp.mongodb.com',
-      },
     });
 
     return result.toUIMessageStream();
@@ -40,5 +41,3 @@ export class DocsProviderTransport implements ChatTransport<UIMessage> {
     return null;
   }
 }
-
-export const docsProviderTransport = new DocsProviderTransport();

packages/compass-assistant/src/docs-provider-transport.ts

@@ -0,0 +1,43 @@
+import { render } from '@mongodb-js/testing-library-compass';
+import { CompassAssistantProvider } from '.';
+import React from 'react';
+import sinon from 'sinon';
+import { expect } from 'chai';
+import {
+  DrawerAnchor,
+  DrawerContentProvider,
+} from '@mongodb-js/compass-components';
+
+describe('CompassAssistantProvider', function () {
+  beforeEach(function () {
+    process.env.COMPASS_ASSISTANT_USE_ATLAS_SERVICE_URL = 'true';
+  });
+
+  afterEach(function () {
+    delete process.env.COMPASS_ASSISTANT_USE_ATLAS_SERVICE_URL;
+  });
+
+  it('uses the Atlas Service assistantApiEndpoint', function () {
+    const mockAtlasService = {
+      assistantApiEndpoint: sinon
+        .stub()
+        .returns('https://atlas-assistant-api.example.com/api/v1'),
+    };
+
+    const MockedProvider = CompassAssistantProvider.withMockServices({
+      atlasService: mockAtlasService,
+    });
+
+    render(
+      <DrawerContentProvider>
+        <DrawerAnchor />
+        <MockedProvider />
+      </DrawerContentProvider>,
+      {
+        preferences: { enableAIAssistant: true },
+      }
+    );
+
+    expect(mockAtlasService.assistantApiEndpoint.calledOnce).to.be.true;
+  });
+});

packages/compass-assistant/src/index.spec.tsx

@@ -1,22 +1,32 @@
 import { registerCompassPlugin } from '@mongodb-js/compass-app-registry';
 import { AssistantProvider } from './assistant-provider';
 import { Chat } from './@ai-sdk/react/chat-react';
-import { docsProviderTransport } from './docs-provider-transport';
+import { DocsProviderTransport } from './docs-provider-transport';
+import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider';
 
 const CompassAssistantProvider = registerCompassPlugin(
   {
     name: 'CompassAssistant',
     component: AssistantProvider,
-    activate: () => {
-      const chat = new Chat({ transport: docsProviderTransport });
+    activate: (initialProps, { atlasService }) => {
+      // TODO: We will temporarily default to the staging url until the docs
+      // API is deployed to the production environment.
+      const baseUrl = process.env.COMPASS_ASSISTANT_USE_ATLAS_SERVICE_URL
+        ? atlasService.assistantApiEndpoint()
+        : 'https://knowledge.staging.corp.mongodb.com/api/v1';
+      const chat = new Chat({
+        transport: new DocsProviderTransport({
+          baseUrl,
+        }),
+      });
       return {
         store: { state: { chat } },
         deactivate: () => {},
       };
     },
   },
   {
-    transport: () => docsProviderTransport,
+    atlasService: atlasServiceLocator,
   }
 );