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 = ({
>