Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 57 additions & 26 deletions packages/compass-assistant/src/compass-assistant-drawer.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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;
}
Expand All @@ -92,16 +80,7 @@ export const CompassAssistantDrawer: React.FunctionComponent<{
<span className={assistantTitleTextStyles}>MongoDB Assistant</span>
<Badge variant="blue">Preview</Badge>
</div>
<IconButton
aria-label="Clear chat"
onClick={() => {
void handleClearChat();
}}
title="Clear chat"
data-testid="assistant-clear-chat"
>
<Icon glyph="Eraser" />
</IconButton>
<ClearChatButton chat={chat} />
</div>
}
label="MongoDB Assistant"
Expand All @@ -123,3 +102,55 @@ export const CompassAssistantDrawer: React.FunctionComponent<{
</DrawerSection>
);
};

export const ClearChatButton: React.FunctionComponent<{
chat: Chat<AssistantMessage>;
}> = ({ 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 (
<Tooltip
trigger={
<IconButton
onClick={() => {
void handleClearChat();
}}
title="Clear chat"
aria-label="Clear chat"
aria-hidden={true}
data-testid="assistant-clear-chat"
>
<Icon glyph="Eraser" />
</IconButton>
}
>
Clear chat
</Tooltip>
);
};
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,11 @@ const TestComponent: React.FunctionComponent<{

return (
<DrawerContentProvider>
<MockedProvider appNameForPrompt="MongoDB Compass" chat={chat}>
<MockedProvider
originForPrompt="mongodb-compass"
appNameForPrompt="MongoDB Compass"
chat={chat}
>
<DrawerAnchor>
<div data-testid="provider-children">Provider children</div>
<CompassAssistantDrawer
Expand All @@ -106,7 +110,11 @@ describe('useAssistantActions', function () {

return (
<DrawerContentProvider>
<MockedProvider appNameForPrompt="MongoDB Compass" chat={chat}>
<MockedProvider
originForPrompt="mongodb-compass"
appNameForPrompt="MongoDB Compass"
chat={chat}
>
{children}
</MockedProvider>
</DrawerContentProvider>
Expand Down Expand Up @@ -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<AssistantMessage>({
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 });

Expand Down Expand Up @@ -609,7 +667,10 @@ describe('CompassAssistantProvider', function () {
render(
<DrawerContentProvider>
<DrawerAnchor />
<MockedProvider appNameForPrompt="MongoDB Compass" />
<MockedProvider
originForPrompt="mongodb-compass"
appNameForPrompt="MongoDB Compass"
/>
</DrawerContentProvider>,
{
preferences: {
Expand Down
13 changes: 3 additions & 10 deletions packages/compass-assistant/src/compass-assistant-provider.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ type AssistantActionsContextType = {
connectionInfo: ConnectionInfo;
error: Error;
}) => void;
clearChat?: () => Promise<void>;
tellMoreAboutInsight?: (context: ProactiveInsightsContext) => void;
ensureOptInAndSend?: (
message: SendMessage,
Expand All @@ -88,7 +87,7 @@ type AssistantActionsContextType = {

type AssistantActionsType = Omit<
AssistantActionsContextType,
'ensureOptInAndSend' | 'clearChat'
'ensureOptInAndSend'
> & {
getIsAssistantEnabled: () => boolean;
};
Expand All @@ -98,7 +97,6 @@ export const AssistantActionsContext =
interpretExplainPlan: () => {},
interpretConnectionError: () => {},
tellMoreAboutInsight: () => {},
clearChat: async () => {},
ensureOptInAndSend: async () => {},
});

Expand Down Expand Up @@ -215,13 +213,6 @@ export const AssistantProvider: React.FunctionComponent<
'performance insights',
buildProactiveInsightsPrompt
),
clearChat: async () => {
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

figured there's no point in keeping this around here

await chat.stop();
chat.clearError();
chat.messages = chat.messages.filter(
(message) => message.metadata?.isPermanent
);
},
ensureOptInAndSend: async (
message: SendMessage,
options: SendOptions,
Expand Down Expand Up @@ -265,6 +256,7 @@ export const CompassAssistantProvider = registerCompassPlugin(
children,
}: PropsWithChildren<{
appNameForPrompt: string;
originForPrompt: string;
chat?: Chat<AssistantMessage>;
atlasAiService?: AtlasAiService;
}>) => {
Expand All @@ -289,6 +281,7 @@ export const CompassAssistantProvider = registerCompassPlugin(
initialProps.chat ??
new Chat({
transport: new DocsProviderTransport({
origin: initialProps.originForPrompt,
instructions: buildConversationInstructionsPrompt({
target: initialProps.appNameForPrompt,
}),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ describe('DocsProviderTransport', function () {
});
abortController = new AbortController();
transport = new DocsProviderTransport({
origin: 'mongodb-compass',
instructions: 'Test instructions for MongoDB assistance',
model: mockModel,
});
Expand Down
7 changes: 7 additions & 0 deletions packages/compass-assistant/src/docs-provider-transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,21 @@ export function shouldExcludeMessage({ metadata }: AssistantMessage) {

export class DocsProviderTransport implements ChatTransport<AssistantMessage> {
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<UIMessageChunk>({
Expand Down Expand Up @@ -60,6 +64,9 @@ export class DocsProviderTransport implements ChatTransport<AssistantMessage> {
model: this.model,
messages: convertToModelMessages(filteredMessages),
abortSignal: abortSignal,
headers: {
'X-Request-Origin': this.origin,
},
providerOptions: {
openai: {
instructions: this.instructions,
Expand Down
1 change: 1 addition & 0 deletions packages/compass-web/src/entrypoint.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -512,6 +512,7 @@ const CompassWeb = ({
>
<CompassInstanceStorePlugin>
<CompassAssistantProvider
originForPrompt="atlas-data-explorer"
appNameForPrompt={
APP_NAMES_FOR_PROMPT.DataExplorer
}
Expand Down
1 change: 1 addition & 0 deletions packages/compass/src/app/components/home.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ function HomeWithConnections({
<ConnectionStorageProvider value={connectionStorage}>
<FileInputBackendProvider createFileInputBackend={createFileInputBackend}>
<CompassAssistantProvider
originForPrompt="mongodb-compass"
appNameForPrompt={APP_NAMES_FOR_PROMPT.Compass}
>
<CompassConnections
Expand Down
Loading