feat(compass-assistant): add a system prompt to precede user prompts with the basic UI context COMPASS-10140 (#7629)

* simple assistant global state implementation

* a basic system context prompt

* remove the non-genuine connections hack for now

* lint

* move currentActiveConnections to ConnectionsComponent

* use workspace-info, add tests

* test that it added the sytem message

* don't take context messages into account for confirmations

* typeguard..

* more typeguards

* clear the state when unmounting

* check for existence of namespace rather

* set collection tab related context in compass-collection rather

* only send the context if it changed

* TODO ticket number

---------

Co-authored-by: Sergey Petushkov <[email protected]>

Author: lerouxb

package-lock.json

@@ -61,8 +61,10 @@
     "@mongodb-js/compass-logging": "^1.7.25",
     "@mongodb-js/compass-telemetry": "^1.19.5",
     "@mongodb-js/connection-info": "^0.24.0",
+    "@mongodb-js/workspace-info": "^1.0.0",
     "ai": "^5.0.26",
     "compass-preferences-model": "^2.66.3",
+    "mongodb-collection-model": "^5.37.0",
     "mongodb-connection-string-url": "^3.0.1",
     "react": "^17.0.2",
     "throttleit": "^2.1.0",

packages/compass-assistant/package.json

@@ -0,0 +1,77 @@
+import type { ConnectionInfo } from '@mongodb-js/connection-info';
+import type {
+  WorkspaceTab,
+  CollectionSubtab,
+} from '@mongodb-js/workspace-info';
+import type { CollectionMetadata } from 'mongodb-collection-model';
+import React, { useEffect } from 'react';
+
+export type GlobalState = {
+  activeConnections: ConnectionInfo[];
+  activeWorkspace: WorkspaceTab | null;
+  activeCollectionMetadata: CollectionMetadata | null;
+  currentQuery: object | null;
+  currentAggregation: object | null;
+  activeCollectionSubTab: CollectionSubtab | null;
+};
+
+const INITIAL_STATE: GlobalState = {
+  activeConnections: [],
+  activeWorkspace: null,
+  activeCollectionMetadata: null,
+  currentQuery: null,
+  currentAggregation: null,
+  activeCollectionSubTab: null,
+};
+
+const AssistantGlobalStateContext = React.createContext<GlobalState>({
+  ...INITIAL_STATE,
+});
+
+const AssistantGlobalSetStateContext = React.createContext<
+  React.Dispatch<React.SetStateAction<GlobalState>>
+>(() => undefined);
+
+export const AssistantGlobalStateProvider: React.FunctionComponent = ({
+  children,
+}) => {
+  const [globalState, setGlobalState] = React.useState({ ...INITIAL_STATE });
+  return (
+    <AssistantGlobalStateContext.Provider value={globalState}>
+      <AssistantGlobalSetStateContext.Provider value={setGlobalState}>
+        {children}
+      </AssistantGlobalSetStateContext.Provider>
+    </AssistantGlobalStateContext.Provider>
+  );
+};
+
+export function useSyncAssistantGlobalState<T extends keyof GlobalState>(
+  stateKey: T,
+  newState: GlobalState[T]
+) {
+  const setState = React.useContext(AssistantGlobalSetStateContext);
+  useEffect(() => {
+    setState((prevState) => {
+      const state = {
+        ...prevState,
+        [stateKey]: newState,
+      };
+
+      // Get rid of some non-sensical states incase the user switches away from
+      // a collection tab to something that is not a collection tab.
+      // activeConnections and activeWorkspace will get updated no matter
+      // how/where the user navigates because those concepts are always
+      // "present" in a way that an active collection is not.
+      if (state.activeWorkspace?.type !== 'Collection') {
+        state.activeCollectionMetadata = null;
+        state.activeCollectionSubTab = null;
+      }
+
+      return state;
+    });
+  }, [newState, setState, stateKey]);
+}
+
+export function useAssistantGlobalState() {
+  return React.useContext(AssistantGlobalStateContext);
+}

packages/compass-assistant/src/assistant-global-state.tsx

@@ -414,20 +414,48 @@ describe('CompassAssistantProvider', function () {
 
       await renderOpenAssistantDrawer({ chat: mockChat });
 
-      userEvent.type(
-        screen.getByPlaceholderText('Ask a question'),
-        'Hello assistant!'
-      );
-      userEvent.click(screen.getByLabelText('Send message'));
+      for (let i = 0; i < 2; i++) {
+        userEvent.type(
+          screen.getByPlaceholderText('Ask a question'),
+          `Hello assistant! (${i})`
+        );
+        userEvent.click(screen.getByLabelText('Send message'));
 
-      await waitFor(() => {
-        expect(sendMessageSpy.calledOnce).to.be.true;
-        expect(sendMessageSpy.firstCall.args[0]).to.deep.include({
-          text: 'Hello assistant!',
+        await waitFor(() => {
+          expect(sendMessageSpy.callCount).to.equal(i + 1);
+          expect(sendMessageSpy.getCall(i).args[0]).to.deep.include({
+            text: `Hello assistant! (${i})`,
+          });
+
+          expect(screen.getByText(`Hello assistant! (${i})`)).to.exist;
         });
+      }
 
-        expect(screen.getByText('Hello assistant!')).to.exist;
-      });
+      const contextMessages = mockChat.messages.filter(
+        (message) => message.metadata?.isSystemContext
+      );
+
+      for (const contextMessage of contextMessages) {
+        // just clear it up so we can deep compare
+        contextMessage.id = 'system-context';
+      }
+
+      // it only sent one
+      expect(contextMessages).to.deep.equal([
+        {
+          id: 'system-context',
+          role: 'system',
+          metadata: {
+            isSystemContext: true,
+          },
+          parts: [
+            {
+              type: 'text',
+              text: 'The user does not have any tabs open.',
+            },
+          ],
+        },
+      ]);
     });
 
     it('will not send new messages if the user does not opt in', async function () {

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

@@ -13,11 +13,13 @@ import {
 } from '@mongodb-js/atlas-service/provider';
 import { DocsProviderTransport } from './docs-provider-transport';
 import {
+  useCurrentValueRef,
   useDrawerActions,
   useInitialValue,
 } from '@mongodb-js/compass-components';
 import {
   buildConnectionErrorPrompt,
+  buildContextPrompt,
   buildExplainPlanPrompt,
   buildProactiveInsightsPrompt,
   type EntryPointMessage,
@@ -31,7 +33,7 @@ import {
   createLoggerLocator,
   type Logger,
 } from '@mongodb-js/compass-logging/provider';
-import type { ConnectionInfo } from '@mongodb-js/connection-info';
+import { type ConnectionInfo } from '@mongodb-js/connection-info';
 import {
   telemetryLocator,
   type TrackFunction,
@@ -41,10 +43,15 @@ import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider'
 import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
 import { buildConversationInstructionsPrompt } from './prompts';
 import { createOpenAI } from '@ai-sdk/openai';
+import {
+  AssistantGlobalStateProvider,
+  useAssistantGlobalState,
+} from './assistant-global-state';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
 export type AssistantMessage = UIMessage & {
+  role?: 'user' | 'assistant' | 'system';
   metadata?: {
     /** The text to display instead of the message text. */
     displayText?: string;
@@ -63,6 +70,11 @@ export type AssistantMessage = UIMessage & {
     instructions?: string;
     /** Excludes history if this message is the last message being sent */
     sendWithoutHistory?: boolean;
+    /** Whether to send the current context along with the message if the context changed */
+    sendContext?: boolean;
+
+    /** Whether this is a message to the model that we don't want to display to the user*/
+    isSystemContext?: boolean;
   };
 };
 
@@ -173,6 +185,16 @@ export type CompassAssistantService = {
   getIsAssistantEnabled: () => boolean;
 };
 
+// Type guard to check if activeWorkspace has a connectionId property
+function hasConnectionId(obj: unknown): obj is { connectionId: string } {
+  return (
+    typeof obj === 'object' &&
+    obj !== null &&
+    'connectionId' in obj &&
+    typeof (obj as any).connectionId === 'string'
+  );
+}
+
 export const AssistantProvider: React.FunctionComponent<
   PropsWithChildren<{
     appNameForPrompt: string;
@@ -183,12 +205,23 @@ export const AssistantProvider: React.FunctionComponent<
   const { openDrawer } = useDrawerActions();
   const track = useTelemetry();
 
+  const assistantGlobalStateRef = useCurrentValueRef(useAssistantGlobalState());
+
+  const lastContextPromptRef = useRef<string | null>(null);
+
   const ensureOptInAndSend = useInitialValue(() => {
     return async function (
       message: SendMessage,
       options: SendOptions,
       callback: () => void
     ) {
+      const {
+        activeWorkspace,
+        activeConnections,
+        activeCollectionMetadata,
+        activeCollectionSubTab,
+      } = assistantGlobalStateRef.current;
+
       try {
         await atlasAiService.ensureAiFeatureAccess();
       } catch {
@@ -204,6 +237,36 @@ export const AssistantProvider: React.FunctionComponent<
         await chat.stop();
       }
 
+      const activeConnection =
+        activeConnections.find((connInfo) => {
+          return (
+            hasConnectionId(activeWorkspace) &&
+            connInfo.id === activeWorkspace.connectionId
+          );
+        }) ?? null;
+
+      const contextPrompt = buildContextPrompt({
+        activeWorkspace,
+        activeConnection,
+        activeCollectionMetadata,
+        activeCollectionSubTab,
+      });
+
+      // use just the text so we have a stable reference to compare against
+      const contextPromptText =
+        contextPrompt.parts[0].type === 'text'
+          ? contextPrompt.parts[0].text
+          : '';
+
+      const shouldSendContextPrompt =
+        message?.metadata?.sendContext &&
+        (!lastContextPromptRef.current ||
+          lastContextPromptRef.current !== contextPromptText);
+      if (shouldSendContextPrompt) {
+        lastContextPromptRef.current = contextPromptText;
+        chat.messages = [...chat.messages, contextPrompt];
+      }
+
       await chat.sendMessage(message, options);
     };
   });
@@ -224,6 +287,7 @@ export const AssistantProvider: React.FunctionComponent<
             metadata: {
               ...metadata,
               source: entryPointName,
+              sendContext: true,
             },
           },
           {},
@@ -285,13 +349,15 @@ export const CompassAssistantProvider = registerCompassPlugin(
         throw new Error('atlasAiService was not provided by the state');
       }
       return (
-        <AssistantProvider
-          appNameForPrompt={appNameForPrompt}
-          chat={chat}
-          atlasAiService={atlasAiService}
-        >
-          {children}
-        </AssistantProvider>
+        <AssistantGlobalStateProvider>
+          <AssistantProvider
+            appNameForPrompt={appNameForPrompt}
+            chat={chat}
+            atlasAiService={atlasAiService}
+          >
+            {children}
+          </AssistantProvider>
+        </AssistantGlobalStateProvider>
       );
     },
     activate: (

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

@@ -263,11 +263,15 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
       const trimmedMessageBody = messageBody.trim();
       if (trimmedMessageBody) {
         await chat.stop();
-        void ensureOptInAndSend?.({ text: trimmedMessageBody }, {}, () => {
-          track('Assistant Prompt Submitted', {
-            user_input_length: trimmedMessageBody.length,
-          });
-        });
+        void ensureOptInAndSend?.(
+          { text: trimmedMessageBody, metadata: { sendContext: true } },
+          {},
+          () => {
+            track('Assistant Prompt Submitted', {
+              user_input_length: trimmedMessageBody.length,
+            });
+          }
+        );
       }
     },
     [track, ensureOptInAndSend, chat]
@@ -357,6 +361,10 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
     [ensureOptInAndSend, setMessages, track]
   );
 
+  const visibleMessages = messages.filter(
+    (message) => !message.metadata?.isSystemContext
+  );
+
   return (
     <div
       data-testid="assistant-chat"
@@ -374,7 +382,7 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
             ref={messagesContainerRef}
           >
             <div className={messagesWrapStyles}>
-              {messages.map((message, index) => {
+              {visibleMessages.map((message, index) => {
                 const { id, role, metadata, parts } = message;
                 const seenTitles = new Set<string>();
                 const sources = [];
@@ -395,7 +403,7 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
                 }
                 if (metadata?.confirmation) {
                   const { description, state } = metadata.confirmation;
-                  const isLastMessage = index === messages.length - 1;
+                  const isLastMessage = index === visibleMessages.length - 1;
 
                   return (
                     <ConfirmationMessage

packages/compass-assistant/src/components/assistant-chat.tsx

@@ -7,3 +7,4 @@ export {
 export type { CompassAssistantService } from './compass-assistant-provider';
 export type { ProactiveInsightsContext, EntryPointMessage } from './prompts';
 export { APP_NAMES_FOR_PROMPT } from './prompts';
+export { useSyncAssistantGlobalState } from './assistant-global-state';