chore(compass-assistant): add explain plan entry point COMPASS-9606 (#7208)

Author: gagik

package-lock.json

@@ -2,11 +2,11 @@ import React from 'react';
 import { render, screen, userEvent } from '@mongodb-js/testing-library-compass';
 import { AssistantChat } from './assistant-chat';
 import { expect } from 'chai';
-import type { UIMessage } from './@ai-sdk/react/use-chat';
 import { createMockChat } from '../test/utils';
+import type { AssistantMessage } from './compass-assistant-provider';
 
 describe('AssistantChat', function () {
-  const mockMessages: UIMessage[] = [
+  const mockMessages: AssistantMessage[] = [
     {
       id: 'user',
       role: 'user',
@@ -24,7 +24,7 @@ describe('AssistantChat', function () {
     },
   ];
 
-  function renderWithChat(messages: UIMessage[]) {
+  function renderWithChat(messages: AssistantMessage[]) {
     const chat = createMockChat({ messages });
     return {
       result: render(<AssistantChat chat={chat} />),
@@ -181,7 +181,7 @@ describe('AssistantChat', function () {
   });
 
   it('handles messages with multiple text parts', function () {
-    const messagesWithMultipleParts: UIMessage[] = [
+    const messagesWithMultipleParts: AssistantMessage[] = [
       {
         id: '1',
         role: 'assistant',
@@ -198,7 +198,7 @@ describe('AssistantChat', function () {
   });
 
   it('handles messages with mixed part types (filters to text only)', function () {
-    const messagesWithMixedParts: UIMessage[] = [
+    const messagesWithMixedParts: AssistantMessage[] = [
       {
         id: '1',
         role: 'assistant',
@@ -217,4 +217,33 @@ describe('AssistantChat', function () {
       .exist;
     expect(screen.queryByText('This should be filtered out.')).to.not.exist;
   });
+
+  it('displays displayText instead of message parts when displayText is set', function () {
+    const messagesWithDisplayText: AssistantMessage[] = [
+      {
+        id: '1',
+        role: 'assistant',
+        parts: [
+          { type: 'text', text: 'This message part should be ignored.' },
+          { type: 'text', text: 'Another part that should not display.' },
+        ],
+        metadata: {
+          displayText: 'This is the custom display text that should show.',
+        },
+      },
+    ];
+
+    renderWithChat(messagesWithDisplayText);
+
+    // Should display the displayText
+    expect(
+      screen.getByText('This is the custom display text that should show.')
+    ).to.exist;
+
+    // Should NOT display the message parts
+    expect(screen.queryByText('This message part should be ignored.')).to.not
+      .exist;
+    expect(screen.queryByText('Another part that should not display.')).to.not
+      .exist;
+  });
 });

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

@@ -1,10 +1,10 @@
 import React, { useCallback, useState } from 'react';
-import type { UIMessage } from './@ai-sdk/react/use-chat';
+import type { AssistantMessage } from './compass-assistant-provider';
 import type { Chat } from './@ai-sdk/react/chat-react';
 import { useChat } from './@ai-sdk/react/use-chat';
 
 interface AssistantChatProps {
-  chat: Chat<UIMessage>;
+  chat: Chat<AssistantMessage>;
 }
 
 /**
@@ -69,10 +69,12 @@ export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
               whiteSpace: 'pre-wrap',
             }}
           >
-            {message.parts
-              ?.filter((part) => part.type === 'text')
-              .map((part) => part.text)
-              .join('') || ''}
+            {message.metadata?.displayText ||
+              message.parts
+                ?.filter((part) => part.type === 'text')
+                .map((part) => part.text)
+                .join('') ||
+              ''}
           </div>
         ))}
       </div>

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

@@ -8,10 +8,10 @@ import {
 import {
   AssistantProvider,
   CompassAssistantProvider,
+  type AssistantMessage,
 } from './compass-assistant-provider';
 import { expect } from 'chai';
 import sinon from 'sinon';
-import type { UIMessage } from './@ai-sdk/react/use-chat';
 import { Chat } from './@ai-sdk/react/chat-react';
 
 import {
@@ -24,7 +24,7 @@ import { createMockChat } from '../test/utils';
 
 // Test component that renders AssistantProvider with children
 const TestComponent: React.FunctionComponent<{
-  chat: Chat<UIMessage>;
+  chat: Chat<AssistantMessage>;
   autoOpen?: boolean;
 }> = ({ chat, autoOpen }) => {
   return (
@@ -78,7 +78,7 @@ describe('AssistantProvider', function () {
     });
 
     async function renderOpenAssistantDrawer(
-      mockChat: Chat<UIMessage>
+      mockChat: Chat<AssistantMessage>
     ): Promise<ReturnType<typeof render>> {
       const result = render(<TestComponent chat={mockChat} autoOpen={true} />, {
         preferences: { enableAIAssistant: true },
@@ -92,7 +92,7 @@ describe('AssistantProvider', function () {
     }
 
     it('displays messages in the chat feed', async function () {
-      const mockMessages: UIMessage[] = [
+      const mockMessages: AssistantMessage[] = [
         {
           id: '1',
           role: 'user',
@@ -120,7 +120,7 @@ describe('AssistantProvider', function () {
     });
 
     it('handles message sending with custom chat when drawer is open', async function () {
-      const mockChat = new Chat<UIMessage>({
+      const mockChat = new Chat<AssistantMessage>({
         messages: [
           {
             id: 'assistant',
@@ -149,7 +149,7 @@ describe('AssistantProvider', function () {
     });
 
     it('new messages are added to the chat feed when the send button is clicked', async function () {
-      const mockChat = new Chat<UIMessage>({
+      const mockChat = new Chat<AssistantMessage>({
         messages: [
           {
             id: 'assistant',

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

@@ -5,29 +5,73 @@ import { createContext, useContext } from 'react';
 import { registerCompassPlugin } from '@mongodb-js/compass-app-registry';
 import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider';
 import { DocsProviderTransport } from './docs-provider-transport';
+import { useDrawerActions } from '@mongodb-js/compass-components';
+import { buildExplainPlanPrompt } from './prompts';
+import { usePreference } from 'compass-preferences-model/provider';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
-type AssistantContextType = Chat<UIMessage>;
+export type AssistantMessage = UIMessage & {
+  metadata?: {
+    /** The text to display instead of the message text. */
+    displayText?: string;
+  };
+};
+
+type AssistantContextType = Chat<AssistantMessage>;
 
 export const AssistantContext = createContext<AssistantContextType | null>(
   null
 );
 
-type AssistantActionsContextType = unknown;
+type AssistantActionsContextType = {
+  interpretExplainPlan: ({
+    namespace,
+    explainPlan,
+  }: {
+    namespace: string;
+    explainPlan: string;
+  }) => void;
+};
 export const AssistantActionsContext =
-  createContext<AssistantActionsContextType>({});
+  createContext<AssistantActionsContextType>({
+    interpretExplainPlan: () => {},
+  });
 
-export function useAssistantActions(): AssistantActionsContextType {
-  return useContext(AssistantActionsContext);
+export function useAssistantActions(): AssistantActionsContextType & {
+  isAssistantEnabled: boolean;
+} {
+  const isAssistantEnabled = usePreference('enableAIAssistant');
+
+  return {
+    ...useContext(AssistantActionsContext),
+    isAssistantEnabled,
+  };
 }
 
 export const AssistantProvider: React.FunctionComponent<
   PropsWithChildren<{
-    chat: Chat<UIMessage>;
+    chat: Chat<AssistantMessage>;
   }>
 > = ({ chat, children }) => {
-  const assistantActionsContext = useRef<AssistantActionsContextType>();
+  const assistantActionsContext = useRef<AssistantActionsContextType>({
+    interpretExplainPlan: ({ explainPlan }) => {
+      openDrawer(ASSISTANT_DRAWER_ID);
+      const { prompt, displayText } = buildExplainPlanPrompt({
+        explainPlan,
+      });
+      void chat.sendMessage(
+        {
+          text: prompt,
+          metadata: {
+            displayText,
+          },
+        },
+        {}
+      );
+    },
+  });
+  const { openDrawer } = useDrawerActions();
 
   return (
     <AssistantContext.Provider value={chat}>
@@ -45,7 +89,7 @@ export const CompassAssistantProvider = registerCompassPlugin(
       chat,
       children,
     }: PropsWithChildren<{
-      chat?: Chat<UIMessage>;
+      chat?: Chat<AssistantMessage>;
     }>) => {
       if (!chat) {
         throw new Error('Chat was not provided by the state');

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

@@ -1,2 +1,3 @@
 export { CompassAssistantProvider } from './compass-assistant-provider';
 export { CompassAssistantDrawer } from './compass-assistant-drawer';
+export { useAssistantActions } from './compass-assistant-provider';

packages/compass-assistant/src/index.tsx

@@ -0,0 +1,13 @@
+export const buildExplainPlanPrompt = ({
+  explainPlan,
+}: {
+  explainPlan: string;
+}) => {
+  return {
+    prompt: `Given the MongoDB explain plan output below, provide a concise human readable explanation that explains the query execution plan and highlights aspects of the plan that might impact query performance. Respond with as much concision and clarity as possible.
+If a clear optimization should be made, please suggest the optimization and describe how it can be accomplished in MongoDB Compass. Do not advise users to create indexes without weighing the pros and cons. 
+Explain output: 
+${explainPlan}`,
+    displayText: 'Provide an explanation of this explain plan.',
+  };
+};

packages/compass-assistant/src/prompts.ts

@@ -1,13 +1,17 @@
-import type { UIMessage } from 'ai';
 import { Chat } from '../src/@ai-sdk/react/chat-react';
 import sinon from 'sinon';
+import type { AssistantMessage } from '../src/compass-assistant-provider';
 
-export const createMockChat = ({ messages }: { messages: UIMessage[] }) => {
-  const newChat = new Chat<UIMessage>({
+export const createMockChat = ({
+  messages,
+}: {
+  messages: AssistantMessage[];
+}) => {
+  const newChat = new Chat<AssistantMessage>({
     messages,
   });
   sinon.replace(newChat, 'sendMessage', sinon.stub());
-  return newChat as unknown as Chat<UIMessage> & {
+  return newChat as unknown as Chat<AssistantMessage> & {
     sendMessage: sinon.SinonStub;
   };
 };

packages/compass-assistant/test/utils.tsx

@@ -68,6 +68,7 @@
     "xvfb-maybe": "^0.2.1"
   },
   "dependencies": {
+    "@mongodb-js/compass-assistant": "^1.0.0",
     "@mongodb-js/compass-app-registry": "^9.4.19",
     "@mongodb-js/compass-collection": "^4.69.0",
     "@mongodb-js/compass-components": "^1.48.0",

packages/compass-explain-plan/package.json

@@ -10,10 +10,18 @@ import { ExplainPlanModal } from './explain-plan-modal';
 import { Provider } from 'react-redux';
 import { activatePlugin } from '../stores';
 
-function render(props: Partial<ExplainPlanModalProps>) {
+function render(
+  props: Partial<ExplainPlanModalProps>,
+  { preferences }: { preferences: { enableAIAssistant: boolean } } = {
+    preferences: { enableAIAssistant: false },
+  }
+) {
   const { store } = activatePlugin(
     { namespace: 'test.test', isDataLake: false },
-    { dataService: {}, localAppRegistry: {}, preferences: {} } as any,
+    {
+      dataService: {},
+      localAppRegistry: {},
+    } as any,
     { on() {}, cleanup() {} } as any
   );
 
@@ -26,7 +34,8 @@ function render(props: Partial<ExplainPlanModalProps>) {
         onModalClose={() => {}}
         {...props}
       ></ExplainPlanModal>
-    </Provider>
+    </Provider>,
+    { preferences: { enableAIAssistant: preferences?.enableAIAssistant } }
   );
 }
 
@@ -51,4 +60,52 @@ describe('ExplainPlanModal', function () {
     render({ status: 'ready' });
     expect(screen.getByText('Query Performance Summary')).to.exist;
   });
+
+  it('should show "Interpret for me" button when AI assistant is enabled', function () {
+    render(
+      {
+        status: 'ready',
+        explainPlan: {
+          namespace: 'test',
+          usedIndexes: [],
+        } as any,
+      },
+      { preferences: { enableAIAssistant: true } }
+    );
+    expect(screen.getByTestId('interpret-for-me-button')).to.exist;
+    expect(screen.getByTestId('interpret-for-me-button')).to.have.attr(
+      'aria-disabled',
+      'false'
+    );
+  });
+
+  it('should not show "Interpret for me" button when AI assistant is disabled', function () {
+    render(
+      {
+        status: 'ready',
+        explainPlan: {
+          namespace: 'test',
+          usedIndexes: [],
+        } as any,
+      },
+      { preferences: { enableAIAssistant: false } }
+    );
+    expect(screen.queryByTestId('interpret-for-me-button')).to.not.exist;
+  });
+
+  it('should disable the "Interpret for me" button when the status is not ready', function () {
+    render(
+      {
+        status: 'loading',
+        explainPlan: {
+          usedIndexes: [],
+        } as any,
+      },
+      { preferences: { enableAIAssistant: true } }
+    );
+    expect(screen.getByTestId('interpret-for-me-button')).to.have.attr(
+      'aria-disabled',
+      'true'
+    );
+  });
 });