diff --git a/packages/compass-assistant/src/compass-assistant-drawer.tsx b/packages/compass-assistant/src/compass-assistant-drawer.tsx index de7a37257d3..c61af06e546 100644 --- a/packages/compass-assistant/src/compass-assistant-drawer.tsx +++ b/packages/compass-assistant/src/compass-assistant-drawer.tsx @@ -8,7 +8,7 @@ import { showConfirmation, spacing, } from '@mongodb-js/compass-components'; -import { AssistantChat } from './assistant-chat'; +import { AssistantChat } from './components/assistant-chat'; import { ASSISTANT_DRAWER_ID, AssistantActionsContext, diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index b72d95c2878..8bffc1a99c7 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -29,6 +29,7 @@ import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; 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'; export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer'; @@ -40,6 +41,11 @@ export type AssistantMessage = UIMessage & { * Used for warning messages in cases like using non-genuine MongoDB. */ isPermanent?: boolean; + /** Information for confirmation messages. */ + confirmation?: { + description: string; + state: 'confirmed' | 'rejected' | 'pending'; + }; }; }; @@ -172,9 +178,12 @@ export const AssistantProvider: React.FunctionComponent< return; } - const { prompt, displayText } = builder(props); + const { prompt, metadata } = builder(props); void assistantActionsContext.current.ensureOptInAndSend( - { text: prompt, metadata: { displayText } }, + { + text: prompt, + metadata, + }, {}, () => { openDrawer(ASSISTANT_DRAWER_ID); @@ -185,17 +194,17 @@ export const AssistantProvider: React.FunctionComponent< } ); }; - }); + }).current; const assistantActionsContext = useRef({ - interpretExplainPlan: createEntryPointHandler.current( + interpretExplainPlan: createEntryPointHandler( 'explain plan', buildExplainPlanPrompt ), - interpretConnectionError: createEntryPointHandler.current( + interpretConnectionError: createEntryPointHandler( 'connection error', buildConnectionErrorPrompt ), - tellMoreAboutInsight: createEntryPointHandler.current( + tellMoreAboutInsight: createEntryPointHandler( 'performance insights', buildProactiveInsightsPrompt ), @@ -220,6 +229,10 @@ export const AssistantProvider: React.FunctionComponent< // place to do tracking. callback(); + if (chat.status === 'streaming') { + await chat.stop(); + } + await chat.sendMessage(message, options); }, }); @@ -267,10 +280,13 @@ export const CompassAssistantProvider = registerCompassPlugin( initialProps.chat ?? new Chat({ transport: new DocsProviderTransport({ - baseUrl: atlasService.assistantApiEndpoint(), instructions: buildConversationInstructionsPrompt({ target: initialProps.appNameForPrompt, }), + model: createOpenAI({ + baseURL: atlasService.assistantApiEndpoint(), + apiKey: '', + }).responses('mongodb-chat-latest'), }), onError: (err: Error) => { logger.log.error( diff --git a/packages/compass-assistant/src/assistant-chat.spec.tsx b/packages/compass-assistant/src/components/assistant-chat.spec.tsx similarity index 67% rename from packages/compass-assistant/src/assistant-chat.spec.tsx rename to packages/compass-assistant/src/components/assistant-chat.spec.tsx index 00687a84f19..ad33c41f369 100644 --- a/packages/compass-assistant/src/assistant-chat.spec.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.spec.tsx @@ -8,13 +8,14 @@ import { } from '@mongodb-js/testing-library-compass'; import { AssistantChat } from './assistant-chat'; import { expect } from 'chai'; -import { createMockChat } from '../test/utils'; +import { createMockChat } from '../../test/utils'; import type { ConnectionInfo } from '@mongodb-js/connection-info'; import { AssistantActionsContext, type AssistantMessage, -} from './compass-assistant-provider'; +} from '../compass-assistant-provider'; import sinon from 'sinon'; +import type { TextPart } from 'ai'; describe('AssistantChat', function () { const mockMessages: AssistantMessage[] = [ @@ -533,6 +534,237 @@ describe('AssistantChat', function () { }); }); + describe('messages with confirmation', function () { + let mockConfirmationMessage: AssistantMessage; + + beforeEach(function () { + mockConfirmationMessage = { + id: 'confirmation-test', + role: 'assistant', + parts: [{ type: 'text', text: 'This is a confirmation message.' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Are you sure you want to proceed with this action?', + }, + }, + }; + }); + + it('renders confirmation message when message has confirmation metadata', function () { + renderWithChat([mockConfirmationMessage]); + + expect(screen.getByText('Please confirm your request')).to.exist; + expect( + screen.getByText('Are you sure you want to proceed with this action?') + ).to.exist; + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + }); + + it('does not render regular message content when confirmation metadata exists', function () { + renderWithChat([mockConfirmationMessage]); + + // Should not show the message text content when confirmation is present + expect(screen.queryByText('This is a confirmation message.')).to.not + .exist; + }); + + it('shows confirmation as pending when it is the last message', function () { + renderWithChat([mockConfirmationMessage]); + + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + expect(screen.queryByText('Request confirmed')).to.not.exist; + expect(screen.queryByText('Request cancelled')).to.not.exist; + }); + + it('shows confirmation as rejected when it is not the last message', function () { + const messages: AssistantMessage[] = [ + mockConfirmationMessage, + { + id: 'newer-message', + role: 'user' as const, + parts: [{ type: 'text', text: 'Another message' }], + }, + ]; + + renderWithChat(messages); + + // The confirmation message (first one) should show as rejected since it's not the last + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + expect(screen.getByText('Request cancelled')).to.exist; + }); + + it('adds new confirmed message when confirmation is confirmed', function () { + const { chat, ensureOptInAndSendStub } = renderWithChat([ + mockConfirmationMessage, + ]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // Should add a new message without confirmation metadata + expect(chat.messages).to.have.length(2); + const newMessage = chat.messages[1]; + expect(newMessage.id).to.equal('confirmation-test-confirmed'); + expect(newMessage.metadata?.confirmation).to.be.undefined; + expect(newMessage.parts).to.deep.equal(mockConfirmationMessage.parts); + + // Should call ensureOptInAndSend to send the new message + expect(ensureOptInAndSendStub.calledOnce).to.be.true; + }); + + it('updates confirmation state to confirmed and adds a new message when confirm button is clicked', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // Original message should have updated confirmation state + const originalMessage = chat.messages[0]; + expect(originalMessage.metadata?.confirmation?.state).to.equal( + 'confirmed' + ); + + expect(chat.messages).to.have.length(2); + + expect( + screen.getByText((mockConfirmationMessage.parts[0] as TextPart).text) + ).to.exist; + }); + + it('updates confirmation state to rejected and does not add a new message when cancel button is clicked', function () { + const { chat, ensureOptInAndSendStub } = renderWithChat([ + mockConfirmationMessage, + ]); + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + // Original message should have updated confirmation state + const originalMessage = chat.messages[0]; + expect(originalMessage.metadata?.confirmation?.state).to.equal( + 'rejected' + ); + + // Should not add a new message + expect(chat.messages).to.have.length(1); + + // Should not call ensureOptInAndSend + expect(ensureOptInAndSendStub.notCalled).to.be.true; + }); + + it('shows confirmed status after confirmation is confirmed', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + // Verify buttons are initially present + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // The state update should be immediate - check the chat messages + const updatedMessage = chat.messages[0]; + expect(updatedMessage.metadata?.confirmation?.state).to.equal( + 'confirmed' + ); + }); + + it('shows cancelled status after confirmation is rejected', function () { + const { chat } = renderWithChat([mockConfirmationMessage]); + + // Verify buttons are initially present + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + // The state update should be immediate - check the chat messages + const updatedMessage = chat.messages[0]; + expect(updatedMessage.metadata?.confirmation?.state).to.equal('rejected'); + }); + + it('handles multiple confirmation messages correctly', function () { + const confirmationMessage1: AssistantMessage = { + id: 'confirmation-1', + role: 'assistant', + parts: [{ type: 'text', text: 'First confirmation' }], + metadata: { + confirmation: { + state: 'pending', + description: 'First confirmation description', + }, + }, + }; + + const confirmationMessage2: AssistantMessage = { + id: 'confirmation-2', + role: 'assistant', + parts: [{ type: 'text', text: 'Second confirmation' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Second confirmation description', + }, + }, + }; + + renderWithChat([confirmationMessage1, confirmationMessage2]); + + expect(screen.getAllByText('Request cancelled')).to.have.length(1); + + expect(screen.getAllByText('Confirm')).to.have.length(1); + expect(screen.getAllByText('Cancel')).to.have.length(1); + expect(screen.getByText('Second confirmation description')).to.exist; + }); + + it('preserves other metadata when creating confirmed message', function () { + const messageWithExtraMetadata: AssistantMessage = { + id: 'confirmation-with-metadata', + role: 'assistant', + parts: [{ type: 'text', text: 'Message with extra metadata' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirmation description', + }, + displayText: 'Custom display text', + isPermanent: true, + }, + }; + + const { chat } = renderWithChat([messageWithExtraMetadata]); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + // New confirmed message should preserve other metadata + const newMessage = chat.messages[1]; + expect(newMessage.metadata?.displayText).to.equal('Custom display text'); + expect(newMessage.metadata?.isPermanent).to.equal(true); + expect(newMessage.metadata?.confirmation).to.be.undefined; + }); + + it('does not render confirmation component for regular messages', function () { + const regularMessage: AssistantMessage = { + id: 'regular', + role: 'assistant', + parts: [{ type: 'text', text: 'This is a regular message' }], + }; + + renderWithChat([regularMessage]); + + expect(screen.queryByText('Please confirm your request')).to.not.exist; + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + expect(screen.getByText('This is a regular message')).to.exist; + }); + }); + describe('related sources', function () { it('displays related resources links for assistant messages that include them', async function () { renderWithChat(mockMessages); diff --git a/packages/compass-assistant/src/assistant-chat.tsx b/packages/compass-assistant/src/components/assistant-chat.tsx similarity index 65% rename from packages/compass-assistant/src/assistant-chat.tsx rename to packages/compass-assistant/src/components/assistant-chat.tsx index 5934c3fea7c..4d9d32ccef9 100644 --- a/packages/compass-assistant/src/assistant-chat.tsx +++ b/packages/compass-assistant/src/components/assistant-chat.tsx @@ -1,8 +1,8 @@ import React, { useCallback, useEffect, useContext } from 'react'; -import type { AssistantMessage } from './compass-assistant-provider'; -import { AssistantActionsContext } from './compass-assistant-provider'; -import type { Chat } from './@ai-sdk/react/chat-react'; -import { useChat } from './@ai-sdk/react/use-chat'; +import type { AssistantMessage } from '../compass-assistant-provider'; +import { AssistantActionsContext } from '../compass-assistant-provider'; +import type { Chat } from '../@ai-sdk/react/chat-react'; +import { useChat } from '../@ai-sdk/react/use-chat'; import { LgChatChatWindow, LgChatLeafygreenChatProvider, @@ -18,8 +18,9 @@ import { LgChatChatDisclaimer, Link, } from '@mongodb-js/compass-components'; +import { ConfirmationMessage } from './confirmation-message'; import { useTelemetry } from '@mongodb-js/compass-telemetry/provider'; -import { NON_GENUINE_WARNING_MESSAGE } from './preset-messages'; +import { NON_GENUINE_WARNING_MESSAGE } from '../preset-messages'; const { DisclaimerText } = LgChatChatDisclaimer; const { ChatWindow } = LgChatChatWindow; @@ -179,25 +180,6 @@ export const AssistantChat: React.FunctionComponent = ({ } }, [hasNonGenuineConnections, chat, setMessages]); - // Transform AI SDK messages to LeafyGreen chat format - const lgMessages = messages.map((message) => ({ - id: message.id, - messageBody: - message.metadata?.displayText || - message.parts - ?.filter((part) => part.type === 'text') - .map((part) => part.text) - .join(''), - isSender: message.role === 'user', - sources: message.parts - .filter((part) => part.type === 'source-url') - .map((part) => ({ - children: part.title || 'Documentation Link', - href: part.url, - variant: 'Docs', - })), - })); - const handleMessageSend = useCallback( (messageBody: string) => { const trimmedMessageBody = messageBody.trim(); @@ -242,6 +224,52 @@ export const AssistantChat: React.FunctionComponent = ({ [track] ); + const handleConfirmation = useCallback( + ( + confirmedMessage: AssistantMessage, + newState: 'confirmed' | 'rejected' + ) => { + setMessages((messages) => { + const newMessages: AssistantMessage[] = messages.map((message) => { + if ( + message.id === confirmedMessage.id && + message.metadata?.confirmation + ) { + return { + ...message, + metadata: { + ...message.metadata, + confirmation: { + ...message.metadata.confirmation, + state: newState, + }, + }, + }; + } + return message; + }); + + // If confirmed, add a new message with the same content but without confirmation metadata + if (newState === 'confirmed') { + newMessages.push({ + ...confirmedMessage, + id: `${confirmedMessage.id}-confirmed`, + metadata: { + ...confirmedMessage.metadata, + confirmation: undefined, + }, + }); + } + return newMessages; + }); + if (newState === 'confirmed') { + // Force the new message request to be sent + void ensureOptInAndSend?.(undefined, {}, () => {}); + } + }, + [ensureOptInAndSend, setMessages] + ); + return (
= ({ className={messageFeedFixesStyles} >
- {lgMessages.map((messageFields) => ( - - {messageFields.isSender === false && ( - { + const { id, role, metadata, parts } = message; + const sources = parts + .filter((part) => part.type === 'source-url') + .map((part) => ({ + children: part.title || 'Documentation Link', + href: part.url, + variant: 'Docs', + })); + if (metadata?.confirmation) { + const { description, state } = metadata.confirmation; + const isLastMessage = index === messages.length - 1; + + return ( + handleConfirmation(message, 'confirmed')} + onReject={() => handleConfirmation(message, 'rejected')} /> - )} - {messageFields.sources.length > 0 && ( - - )} - - ))} + ); + } + + const displayText = + message.metadata?.displayText || + message.parts + ?.filter((part) => part.type === 'text') + .map((part) => part.text) + .join(''); + + const isSender = role === 'user'; + + return ( + + {isSender === false && ( + + )} + {sources.length > 0 && } + + ); + })}
This feature is powered by generative AI. See our{' '} @@ -297,7 +364,7 @@ export const AssistantChat: React.FunctionComponent = ({
)} - {lgMessages.length === 0 && ( + {messages.length === 0 && (

Welcome to your MongoDB Assistant.

Ask any question about MongoDB to receive expert guidance and diff --git a/packages/compass-assistant/src/components/confirmation-message.spec.tsx b/packages/compass-assistant/src/components/confirmation-message.spec.tsx new file mode 100644 index 00000000000..d9b822afa65 --- /dev/null +++ b/packages/compass-assistant/src/components/confirmation-message.spec.tsx @@ -0,0 +1,118 @@ +import React from 'react'; +import { render, screen, userEvent } from '@mongodb-js/testing-library-compass'; +import { ConfirmationMessage } from './confirmation-message'; +import { expect } from 'chai'; +import sinon from 'sinon'; + +describe('ConfirmationMessage', function () { + const defaultProps = { + state: 'pending' as const, + title: 'Test Confirmation', + description: 'Are you sure you want to proceed with this action?', + onConfirm: () => {}, + onReject: () => {}, + }; + + it('renders title and description', function () { + render(); + + expect(screen.getByText(defaultProps.title)).to.exist; + expect(screen.getByText(defaultProps.description)).to.exist; + }); + + describe('pending state', function () { + it('shows confirm and cancel buttons', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + expect(screen.getByText('Confirm')).to.exist; + expect(screen.getByText('Cancel')).to.exist; + }); + + it('calls onConfirm when confirm button is clicked', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + const confirmButton = screen.getByText('Confirm'); + userEvent.click(confirmButton); + + expect(onConfirm.calledOnce).to.be.true; + expect(onReject.notCalled).to.be.true; + }); + + it('calls onReject when cancel button is clicked', function () { + const onConfirm = sinon.stub(); + const onReject = sinon.stub(); + + render( + + ); + + const cancelButton = screen.getByText('Cancel'); + userEvent.click(cancelButton); + + expect(onReject.calledOnce).to.be.true; + expect(onConfirm.notCalled).to.be.true; + }); + + it('does not show status when in pending state', function () { + render( + + ); + + expect(screen.queryByText('Request confirmed')).to.not.exist; + expect(screen.queryByText('Request cancelled')).to.not.exist; + }); + }); + + describe('confirmed and rejected states', function () { + it('shows confirmed status with checkmark icon', function () { + render(); + + expect(screen.getByText('Request confirmed')).to.exist; + // sic from the icon library + expect(screen.getByLabelText('Checkmark With Circle Icon')).to.exist; + + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + }); + + it('shows cancelled status', function () { + render(); + + expect(screen.getByText('Request cancelled')).to.exist; + // sic from the icon library + expect(screen.getByLabelText('XWith Circle Icon')).to.exist; + + expect(screen.queryByText('Confirm')).to.not.exist; + expect(screen.queryByText('Cancel')).to.not.exist; + }); + }); +}); diff --git a/packages/compass-assistant/src/components/confirmation-message.tsx b/packages/compass-assistant/src/components/confirmation-message.tsx new file mode 100644 index 00000000000..d6929beeac9 --- /dev/null +++ b/packages/compass-assistant/src/components/confirmation-message.tsx @@ -0,0 +1,103 @@ +import React from 'react'; +import { + Icon, + Body, + Button, + ButtonVariant, + spacing, + css, + palette, + useDarkMode, +} from '@mongodb-js/compass-components'; + +const confirmationMessageStyles = css({ + padding: spacing[300], + borderRadius: spacing[200], + backgroundColor: palette.gray.light3, + border: `1px solid ${palette.gray.light2}`, +}); + +const confirmationStatusStyles = css({ + display: 'flex', + alignItems: 'center', + gap: spacing[100], + marginTop: spacing[200], +}); + +const confirmationStatusTextStyles = css({ + color: palette.gray.dark1, +}); + +const confirmationTitleStyles = css({ + marginBottom: spacing[150], + fontWeight: 600, +}); + +const buttonGroupStyles = css({ + display: 'flex', + gap: spacing[200], + marginTop: spacing[300], + '> button': { + width: '100%', + }, +}); + +interface ConfirmationMessageProps { + state: 'confirmed' | 'rejected' | 'pending'; + title: string; + description: string; + onConfirm: () => void; + onReject: () => void; +} + +export const ConfirmationMessage: React.FunctionComponent< + ConfirmationMessageProps +> = ({ state, title, description, onConfirm, onReject }) => { + const darkMode = useDarkMode(); + + return ( +
+ {title} + + {description} + + {state === 'pending' && ( +
+ + +
+ )} + {state !== 'pending' && ( +
+ + + Request {state === 'confirmed' ? 'confirmed' : 'cancelled'} + +
+ )} +
+ ); +}; diff --git a/packages/compass-assistant/src/docs-provider-transport.spec.ts b/packages/compass-assistant/src/docs-provider-transport.spec.ts new file mode 100644 index 00000000000..2143d429971 --- /dev/null +++ b/packages/compass-assistant/src/docs-provider-transport.spec.ts @@ -0,0 +1,202 @@ +import { expect } from 'chai'; +import sinon from 'sinon'; +import { + DocsProviderTransport, + shouldExcludeMessage, +} from './docs-provider-transport'; +import type { AssistantMessage } from './compass-assistant-provider'; +import { MockLanguageModelV2 } from 'ai/test'; +import type { UIMessageChunk } from 'ai'; +import { waitFor } from '@mongodb-js/testing-library-compass'; + +describe('DocsProviderTransport', function () { + describe('shouldExcludeMessage', function () { + it('returns false for messages without confirmation metadata', function () { + const message: AssistantMessage = { + id: 'test-1', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + }; + + expect(shouldExcludeMessage(message)).to.be.false; + }); + + it('returns true for confirmation messages', function () { + const message: AssistantMessage = { + id: 'test-5', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirm this action', + }, + }, + }; + + expect(shouldExcludeMessage(message)).to.be.true; + }); + }); + + describe('sending messages', function () { + let mockModel: MockLanguageModelV2; + let doStream: sinon.SinonStub; + let transport: DocsProviderTransport; + let abortController: AbortController; + let sendMessages: ( + params: Partial[0]> + ) => Promise>; + + beforeEach(function () { + // Mock the OpenAI client + doStream = sinon.stub().returns({ + stream: DocsProviderTransport.emptyStream, + request: { + body: { + messages: [], + }, + }, + }); + mockModel = new MockLanguageModelV2({ + doStream, + }); + abortController = new AbortController(); + transport = new DocsProviderTransport({ + instructions: 'Test instructions for MongoDB assistance', + model: mockModel, + }); + sendMessages = (params) => + transport.sendMessages({ + trigger: 'submit-message', + chatId: 'test-chat', + messageId: undefined, + abortSignal: abortController.signal, + messages: [], + ...params, + }); + }); + + afterEach(function () { + sinon.restore(); + }); + + describe('sendMessages', function () { + const userMessage: AssistantMessage = { + id: 'included1', + role: 'user', + parts: [{ type: 'text', text: 'User message' }], + }; + const confirmationPendingMessage: AssistantMessage = { + id: 'test', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'pending', + description: 'Confirm this action', + }, + }, + }; + const confirmationConfirmedMessage: AssistantMessage = { + id: 'test', + role: 'assistant', + parts: [{ type: 'text', text: 'Response' }], + metadata: { + confirmation: { + state: 'confirmed', + description: 'Confirmed action', + }, + }, + }; + it('returns empty stream when last message should be excluded', async function () { + const messages: AssistantMessage[] = [ + userMessage, + confirmationConfirmedMessage, + ]; + + const result = await sendMessages({ + messages, + }); + + expect(result).to.equal(DocsProviderTransport.emptyStream); + expect(mockModel.doStreamCalls).to.be.empty; + }); + + it('returns empty stream when all messages are filtered out', async function () { + const messages: AssistantMessage[] = [ + confirmationPendingMessage, + confirmationConfirmedMessage, + ]; + + const result = await sendMessages({ + messages, + }); + + expect(result).to.equal(DocsProviderTransport.emptyStream); + expect(mockModel.doStreamCalls).to.be.empty; + }); + + it('sends filtered messages to AI when valid messages exist', async function () { + await sendMessages({ + messages: [userMessage], + }); + + await waitFor(() => { + expect(doStream).to.have.been.calledOnce; + expect(doStream.firstCall.args[0]).to.deep.include({ + prompt: [ + { + role: 'user', + providerOptions: undefined, + content: [ + { + type: 'text', + text: 'User message', + providerOptions: undefined, + }, + ], + }, + ], + }); + }); + }); + + it('sends only valid messages when confirmation required messages exist', async function () { + await sendMessages({ + messages: [ + confirmationConfirmedMessage, + confirmationPendingMessage, + userMessage, + ], + }); + + await waitFor(() => { + expect(doStream).to.have.been.calledOnce; + expect(doStream.firstCall.args[0]).to.deep.include({ + prompt: [ + { + role: 'user', + providerOptions: undefined, + content: [ + { + type: 'text', + text: 'User message', + providerOptions: undefined, + }, + ], + }, + ], + }); + }); + }); + }); + + // We currently do not support reconnecting to streams but we may want to in the future + describe('reconnectToStream', function () { + it('always returns null', async function () { + const result = await transport.reconnectToStream(); + expect(result).to.be.null; + }); + }); + }); +}); diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index 1cdc5bc3c1a..9a7ae78200e 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -1,37 +1,64 @@ import { type ChatTransport, - type UIMessage, + type LanguageModel, type UIMessageChunk, convertToModelMessages, streamText, } from 'ai'; -import { createOpenAI } from '@ai-sdk/openai'; +import type { AssistantMessage } from './compass-assistant-provider'; -export class DocsProviderTransport implements ChatTransport { - private openai: ReturnType; +/** Returns true if the message should be excluded from being sent to the assistant API. */ +export function shouldExcludeMessage({ metadata }: AssistantMessage) { + if (metadata?.confirmation) { + return true; + } + return false; +} + +export class DocsProviderTransport implements ChatTransport { + private model: LanguageModel; private instructions: string; constructor({ - baseUrl, instructions, + model, }: { - baseUrl: string; instructions: string; + model: LanguageModel; }) { - this.openai = createOpenAI({ - baseURL: baseUrl, - apiKey: '', - }); this.instructions = instructions; + this.model = model; } + static emptyStream = new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + sendMessages({ messages, abortSignal, - }: Parameters['sendMessages']>[0]) { + }: Parameters['sendMessages']>[0]) { + // If the most recent message is a message that is meant to be excluded + // then we do not need to send this request to the assistant API as it's likely + // redundant otherwise. + if (shouldExcludeMessage(messages[messages.length - 1])) { + return Promise.resolve(DocsProviderTransport.emptyStream); + } + + const filteredMessages = messages.filter( + (message) => !shouldExcludeMessage(message) + ); + + // If no messages remain after filtering, return an empty stream + if (filteredMessages.length === 0) { + return Promise.resolve(DocsProviderTransport.emptyStream); + } + const result = streamText({ - model: this.openai.responses('mongodb-chat-latest'), - messages: convertToModelMessages(messages), + model: this.model, + messages: convertToModelMessages(filteredMessages), abortSignal: abortSignal, providerOptions: { openai: { diff --git a/packages/compass-assistant/src/prompts.ts b/packages/compass-assistant/src/prompts.ts index 067ea2a29b9..de7be205c6e 100644 --- a/packages/compass-assistant/src/prompts.ts +++ b/packages/compass-assistant/src/prompts.ts @@ -1,5 +1,11 @@ import type { ConnectionInfo } from '@mongodb-js/connection-info'; import { redactConnectionString } from 'mongodb-connection-string-url'; +import type { AssistantMessage } from './compass-assistant-provider'; + +export type EntryPointMessage = { + prompt: string; + metadata: AssistantMessage['metadata']; +}; export const APP_NAMES_FOR_PROMPT = { Compass: 'MongoDB Compass', @@ -38,11 +44,6 @@ Always call the 'search_content' tool when asked a technical question that would `; }; -export type EntryPointMessage = { - prompt: string; - displayText?: string; -}; - export type ExplainPlanContext = { explainPlan: string; }; @@ -116,7 +117,14 @@ Tell the user if indexes need to be created or modified to enable any recommenda ${explainPlan} `, - displayText: 'Interpret this explain plan output for me.', + metadata: { + displayText: 'Interpret this explain plan output for me.', + confirmation: { + description: + 'Explain plan metadata, including the original query, may be used to process your request', + state: 'pending', + }, + }, }; }; @@ -136,8 +144,6 @@ export const buildProactiveInsightsPrompt = ( switch (context.id) { case 'aggregation-executed-without-index': { return { - displayText: - 'Help me understand the performance impact of running aggregations without an index.', prompt: `The given MongoDB aggregation was executed without an index. Provide a concise human readable explanation that explains why it might degrade performance to not use an index. Please suggest whether an existing index can be used to improve the performance of this query, or if a new index must be created, and describe how it can be accomplished in MongoDB Compass. Do not advise users to create indexes without weighing the pros and cons. @@ -147,12 +153,14 @@ Respond with as much concision and clarity as possible. ${context.stages.join('\n')} `, + metadata: { + displayText: + 'Help me understand the performance impact of running aggregations without an index.', + }, }; } case 'query-executed-without-index': return { - displayText: - 'Help me understand the performance impact of running queries without an index.', prompt: `The given MongoDB query was executed without an index. Provide a concise human readable explanation that explains why it might degrade performance to not use an index. Please suggest whether an existing index can be used to improve the performance of this query, or if a new index must be created, and describe how it can be accomplished in MongoDB Compass. Do not advise users to create indexes without weighing the pros and cons. @@ -162,6 +170,10 @@ Respond with as much concision and clarity as possible. ${context.query} `, + metadata: { + displayText: + 'Help me understand the performance impact of running queries without an index.', + }, }; } }; @@ -190,7 +202,9 @@ ${connectionString} Error message: ${connectionError}`, - displayText: - 'Diagnose why my Compass connection is failing and help me debug it.', + metadata: { + displayText: + 'Diagnose why my Compass connection is failing and help me debug it.', + }, }; };