diff --git a/packages/compass-assistant/src/compass-assistant-drawer.tsx b/packages/compass-assistant/src/compass-assistant-drawer.tsx index fe0e0f0e107..a895ec5e00a 100644 --- a/packages/compass-assistant/src/compass-assistant-drawer.tsx +++ b/packages/compass-assistant/src/compass-assistant-drawer.tsx @@ -7,17 +7,20 @@ import { IconButton, showConfirmation, spacing, + Tooltip, } from '@mongodb-js/compass-components'; import { AssistantChat } from './components/assistant-chat'; import { ASSISTANT_DRAWER_ID, - AssistantActionsContext, AssistantContext, + type AssistantMessage, } from './compass-assistant-provider'; import { useIsAIFeatureEnabled, usePreference, } from 'compass-preferences-model/provider'; +import { useChat } from './@ai-sdk/react/use-chat'; +import type { Chat } from './@ai-sdk/react/chat-react'; const assistantTitleStyles = css({ display: 'flex', @@ -54,25 +57,10 @@ export const CompassAssistantDrawer: React.FunctionComponent<{ hasNonGenuineConnections?: boolean; }> = ({ appName, autoOpen, hasNonGenuineConnections = false }) => { const chat = useContext(AssistantContext); - const { clearChat } = useContext(AssistantActionsContext); const enableAIAssistant = usePreference('enableAIAssistant'); const isAiFeatureEnabled = useIsAIFeatureEnabled(); - const handleClearChat = useCallback(async () => { - const confirmed = await showConfirmation({ - title: 'Clear this chat?', - description: - 'The current chat will be cleared, and chat history will not be retrievable.', - buttonText: 'Clear chat', - variant: 'danger', - 'data-testid': 'assistant-confirm-clear-chat-modal', - }); - if (confirmed) { - await clearChat?.(); - } - }, [clearChat]); - if (!enableAIAssistant || !isAiFeatureEnabled) { return null; } @@ -92,16 +80,7 @@ export const CompassAssistantDrawer: React.FunctionComponent<{ MongoDB Assistant Preview - { - void handleClearChat(); - }} - title="Clear chat" - data-testid="assistant-clear-chat" - > - - + } label="MongoDB Assistant" @@ -123,3 +102,55 @@ export const CompassAssistantDrawer: React.FunctionComponent<{ ); }; + +export const ClearChatButton: React.FunctionComponent<{ + chat: Chat; +}> = ({ chat }) => { + const { clearError, stop } = useChat({ chat }); + + const handleClearChat = useCallback(async () => { + const confirmed = await showConfirmation({ + title: 'Clear this chat?', + description: + 'The current chat will be cleared, and chat history will not be retrievable.', + buttonText: 'Clear chat', + variant: 'danger', + 'data-testid': 'assistant-confirm-clear-chat-modal', + }); + if (confirmed) { + await stop(); + clearError(); + chat.messages = chat.messages.filter( + (message) => message.metadata?.isPermanent + ); + } + }, [stop, clearError, chat]); + + const isChatEmpty = + chat.messages.filter((message) => !message.metadata?.isPermanent).length === + 0; + + if (isChatEmpty) { + return null; + } + + return ( + { + void handleClearChat(); + }} + title="Clear chat" + aria-label="Clear chat" + aria-hidden={true} + data-testid="assistant-clear-chat" + > + + + } + > + Clear chat + + ); +}; diff --git a/packages/compass-assistant/src/compass-assistant-provider.spec.tsx b/packages/compass-assistant/src/compass-assistant-provider.spec.tsx index 0e1230d88cd..9f12549c284 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.spec.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.spec.tsx @@ -85,7 +85,11 @@ const TestComponent: React.FunctionComponent<{ return ( - +
Provider children
- + {children}
@@ -485,6 +493,56 @@ describe('CompassAssistantProvider', function () { }); describe('clear chat button', function () { + it('is hidden when the chat is empty', async function () { + const mockChat = createMockChat({ messages: [] }); + await renderOpenAssistantDrawer({ chat: mockChat }); + expect(screen.queryByTestId('assistant-clear-chat')).to.not.exist; + }); + + it('is hidden when the chat has only permanent messages', async function () { + const mockChat = createMockChat({ + messages: mockMessages.map((message) => ({ + ...message, + metadata: { isPermanent: true }, + })), + }); + await renderOpenAssistantDrawer({ chat: mockChat }); + expect(screen.queryByTestId('assistant-clear-chat')).to.not.exist; + }); + + it('is visible when the chat has messages', async function () { + const mockChat = createMockChat({ messages: mockMessages }); + await renderOpenAssistantDrawer({ chat: mockChat }); + expect(screen.getByTestId('assistant-clear-chat')).to.exist; + }); + + it('appears after a message is sent', async function () { + const mockChat = new Chat({ + messages: [], + transport: { + sendMessages: sinon.stub().returns( + new Promise(() => { + return new ReadableStream({}); + }) + ), + reconnectToStream: sinon.stub(), + }, + }); + await renderOpenAssistantDrawer({ chat: mockChat }); + + expect(screen.queryByTestId('assistant-clear-chat')).to.not.exist; + + userEvent.type( + screen.getByPlaceholderText('Ask a question'), + 'Hello assistant' + ); + userEvent.click(screen.getByLabelText('Send message')); + + await waitFor(() => { + expect(screen.getByTestId('assistant-clear-chat')).to.exist; + }); + }); + it('clears the chat when the user clicks and confirms', async function () { const mockChat = createMockChat({ messages: mockMessages }); @@ -609,7 +667,10 @@ describe('CompassAssistantProvider', function () { render( - + , { preferences: { diff --git a/packages/compass-assistant/src/compass-assistant-provider.tsx b/packages/compass-assistant/src/compass-assistant-provider.tsx index 2f17904e6e0..a702e023718 100644 --- a/packages/compass-assistant/src/compass-assistant-provider.tsx +++ b/packages/compass-assistant/src/compass-assistant-provider.tsx @@ -77,7 +77,6 @@ type AssistantActionsContextType = { connectionInfo: ConnectionInfo; error: Error; }) => void; - clearChat?: () => Promise; tellMoreAboutInsight?: (context: ProactiveInsightsContext) => void; ensureOptInAndSend?: ( message: SendMessage, @@ -88,7 +87,7 @@ type AssistantActionsContextType = { type AssistantActionsType = Omit< AssistantActionsContextType, - 'ensureOptInAndSend' | 'clearChat' + 'ensureOptInAndSend' > & { getIsAssistantEnabled: () => boolean; }; @@ -98,7 +97,6 @@ export const AssistantActionsContext = interpretExplainPlan: () => {}, interpretConnectionError: () => {}, tellMoreAboutInsight: () => {}, - clearChat: async () => {}, ensureOptInAndSend: async () => {}, }); @@ -215,13 +213,6 @@ export const AssistantProvider: React.FunctionComponent< 'performance insights', buildProactiveInsightsPrompt ), - clearChat: async () => { - await chat.stop(); - chat.clearError(); - chat.messages = chat.messages.filter( - (message) => message.metadata?.isPermanent - ); - }, ensureOptInAndSend: async ( message: SendMessage, options: SendOptions, @@ -265,6 +256,7 @@ export const CompassAssistantProvider = registerCompassPlugin( children, }: PropsWithChildren<{ appNameForPrompt: string; + originForPrompt: string; chat?: Chat; atlasAiService?: AtlasAiService; }>) => { @@ -289,6 +281,7 @@ export const CompassAssistantProvider = registerCompassPlugin( initialProps.chat ?? new Chat({ transport: new DocsProviderTransport({ + origin: initialProps.originForPrompt, instructions: buildConversationInstructionsPrompt({ target: initialProps.appNameForPrompt, }), diff --git a/packages/compass-assistant/src/docs-provider-transport.spec.ts b/packages/compass-assistant/src/docs-provider-transport.spec.ts index 2143d429971..ed570d6d3be 100644 --- a/packages/compass-assistant/src/docs-provider-transport.spec.ts +++ b/packages/compass-assistant/src/docs-provider-transport.spec.ts @@ -62,6 +62,7 @@ describe('DocsProviderTransport', function () { }); abortController = new AbortController(); transport = new DocsProviderTransport({ + origin: 'mongodb-compass', instructions: 'Test instructions for MongoDB assistance', model: mockModel, }); diff --git a/packages/compass-assistant/src/docs-provider-transport.ts b/packages/compass-assistant/src/docs-provider-transport.ts index 9a7ae78200e..36e541c4b66 100644 --- a/packages/compass-assistant/src/docs-provider-transport.ts +++ b/packages/compass-assistant/src/docs-provider-transport.ts @@ -17,17 +17,21 @@ export function shouldExcludeMessage({ metadata }: AssistantMessage) { export class DocsProviderTransport implements ChatTransport { private model: LanguageModel; + private origin: string; private instructions: string; constructor({ instructions, model, + origin, }: { instructions: string; model: LanguageModel; + origin: string; }) { this.instructions = instructions; this.model = model; + this.origin = origin; } static emptyStream = new ReadableStream({ @@ -60,6 +64,9 @@ export class DocsProviderTransport implements ChatTransport { model: this.model, messages: convertToModelMessages(filteredMessages), abortSignal: abortSignal, + headers: { + 'X-Request-Origin': this.origin, + }, providerOptions: { openai: { instructions: this.instructions, diff --git a/packages/compass-web/src/entrypoint.tsx b/packages/compass-web/src/entrypoint.tsx index e43725cc09b..0d19c4ee756 100644 --- a/packages/compass-web/src/entrypoint.tsx +++ b/packages/compass-web/src/entrypoint.tsx @@ -512,6 +512,7 @@ const CompassWeb = ({ >