Skip to content

Commit dff3d43

Browse files
gagikgagik
committed
chore(compass-assistant): split drawer and provider usage for connections context
1 parent 6212cb1 commit dff3d43

File tree

8 files changed

+213
-78
lines changed

8 files changed

+213
-78
lines changed

‎packages/compass-assistant/README.md‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
# Compass Assistant Architecture
2+
3+
## Overview
4+
5+
The Compass Assistant package provides a modular architecture for integrating AI-powered assistance into MongoDB Compass. The components are designed to be flexible and composable, allowing the assistant functionality to be placed at different levels in the component tree.
6+
7+
## Architecture
8+
9+
### Components
10+
11+
1. **AssistantProvider** - Context provider that manages chat state and provides assistant functionality
12+
2. **AssistantChat** - Chat interface component that consumes the assistant context
13+
3. **AssistantDrawer** - Pre-configured drawer section that wraps AssistantChat
14+
4. **useAssistant** - Hook for consuming assistant context in components
15+
16+
### Usage Patterns
17+
18+
#### Basic Setup
19+
20+
```tsx
21+
import {
22+
AssistantProvider,
23+
AssistantDrawer,
24+
} from '@mongodb-js/compass-assistant';
25+
26+
// At the top level, wrap your app with the provider
27+
<AssistantProvider chat={chatInstance}>
28+
<YourApp />
29+
30+
{/* Place the drawer anywhere in your component tree */}
31+
<AssistantDrawer />
32+
</AssistantProvider>;
33+
```
34+
35+
#### Custom Implementation
36+
37+
```tsx
38+
import {
39+
AssistantProvider,
40+
AssistantChat,
41+
useAssistant,
42+
} from '@mongodb-js/compass-assistant';
43+
44+
// Custom component that uses the assistant context
45+
function CustomAssistantUI() {
46+
const { messages, sendMessage } = useAssistant();
47+
48+
return (
49+
<div>
50+
<h2>Custom Assistant UI</h2>
51+
<AssistantChat />
52+
<div>Message count: {messages.length}</div>
53+
</div>
54+
);
55+
}
56+
57+
// Usage
58+
<AssistantProvider chat={chatInstance}>
59+
<YourApp />
60+
<CustomAssistantUI />
61+
</AssistantProvider>;
62+
```
63+
64+
## Benefits
65+
66+
- **Separation of Concerns**: Chat UI is separate from state management
67+
- **Flexibility**: Components can be placed at any level in the tree
68+
- **Reusability**: Multiple components can consume the same assistant context
69+
- **Testability**: Easy to mock context for testing individual components
70+
71+
## Migration Guide
72+
73+
If upgrading from the old prop-based approach:
74+
75+
**Before:**
76+
77+
```tsx
78+
<AssistantChat messages={messages} onSendMessage={handleSend} />
79+
```
80+
81+
**After:**
82+
83+
```tsx
84+
<AssistantProvider chat={chatInstance}>
85+
<AssistantChat />
86+
</AssistantProvider>
87+
```

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

Lines changed: 45 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ import { render, screen, userEvent } from '@mongodb-js/testing-library-compass';
33
import { AssistantChat } from './assistant-chat';
44
import { expect } from 'chai';
55
import type { UIMessage } from './@ai-sdk/react/use-chat';
6+
import { Chat } from './@ai-sdk/react/chat-react';
7+
import sinon from 'sinon';
68

79
describe('AssistantChat', function () {
810
const mockMessages: UIMessage[] = [
@@ -23,8 +25,30 @@ describe('AssistantChat', function () {
2325
},
2426
];
2527

28+
let renderWithChat: (messages: UIMessage[]) => {
29+
result: ReturnType<typeof render>;
30+
chat: Chat<UIMessage> & {
31+
sendMessage: sinon.SinonStub;
32+
};
33+
};
34+
35+
beforeEach(() => {
36+
renderWithChat = (messages: UIMessage[]) => {
37+
const newChat = new Chat<UIMessage>({
38+
messages,
39+
});
40+
sinon.replace(newChat, 'sendMessage', sinon.stub());
41+
return {
42+
result: render(<AssistantChat chat={newChat} />),
43+
chat: newChat as unknown as Chat<UIMessage> & {
44+
sendMessage: sinon.SinonStub;
45+
},
46+
};
47+
};
48+
});
49+
2650
it('renders input field and send button', function () {
27-
render(<AssistantChat messages={[]} />);
51+
renderWithChat([]);
2852

2953
const inputField = screen.getByTestId('assistant-chat-input');
3054
const sendButton = screen.getByTestId('assistant-chat-send-button');
@@ -34,7 +58,7 @@ describe('AssistantChat', function () {
3458
});
3559

3660
it('input field accepts text input', function () {
37-
render(<AssistantChat messages={[]} />);
61+
renderWithChat([]);
3862

3963
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
4064
const inputField = screen.getByTestId(
@@ -47,7 +71,7 @@ describe('AssistantChat', function () {
4771
});
4872

4973
it('send button is disabled when input is empty', function () {
50-
render(<AssistantChat messages={[]} />);
74+
renderWithChat([]);
5175

5276
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
5377
const sendButton = screen.getByTestId(
@@ -58,7 +82,7 @@ describe('AssistantChat', function () {
5882
});
5983

6084
it('send button is enabled when input has text', function () {
61-
render(<AssistantChat messages={[]} />);
85+
renderWithChat([]);
6286

6387
const inputField = screen.getByTestId('assistant-chat-input');
6488
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
@@ -72,7 +96,7 @@ describe('AssistantChat', function () {
7296
});
7397

7498
it('send button is disabled for whitespace-only input', function () {
75-
render(<AssistantChat messages={[]} />);
99+
renderWithChat([]);
76100

77101
const inputField = screen.getByTestId('assistant-chat-input');
78102
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
@@ -86,7 +110,7 @@ describe('AssistantChat', function () {
86110
});
87111

88112
it('displays messages in the chat feed', function () {
89-
render(<AssistantChat messages={mockMessages} />);
113+
renderWithChat(mockMessages);
90114

91115
expect(screen.getByTestId('assistant-message-user')).to.exist;
92116
expect(screen.getByTestId('assistant-message-assistant')).to.exist;
@@ -95,27 +119,20 @@ describe('AssistantChat', function () {
95119
.exist;
96120
});
97121

98-
it('calls onSendMessage when form is submitted', function () {
99-
let sentMessage = '';
100-
const handleSendMessage = (message: string) => {
101-
sentMessage = message;
102-
};
103-
104-
render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
105-
122+
it('calls sendMessage when form is submitted', function () {
123+
const { chat } = renderWithChat([]);
106124
const inputField = screen.getByTestId('assistant-chat-input');
107125
const sendButton = screen.getByTestId('assistant-chat-send-button');
108126

109127
userEvent.type(inputField, 'What is aggregation?');
110128
userEvent.click(sendButton);
111129

112-
expect(sentMessage).to.equal('What is aggregation?');
130+
expect(chat.sendMessage.calledWith({ text: 'What is aggregation?' })).to.be
131+
.true;
113132
});
114133

115134
it('clears input field after successful submission', function () {
116-
const handleSendMessage = () => {};
117-
118-
render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
135+
renderWithChat([]);
119136

120137
// eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion
121138
const inputField = screen.getByTestId(
@@ -130,44 +147,35 @@ describe('AssistantChat', function () {
130147
});
131148

132149
it('trims whitespace from input before sending', function () {
133-
let sentMessage = '';
134-
const handleSendMessage = (message: string) => {
135-
sentMessage = message;
136-
};
137-
138-
render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
150+
const { chat } = renderWithChat([]);
139151

140152
const inputField = screen.getByTestId('assistant-chat-input');
141153

142154
userEvent.type(inputField, ' What is sharding? ');
143155
userEvent.click(screen.getByTestId('assistant-chat-send-button'));
144156

145-
expect(sentMessage).to.equal('What is sharding?');
157+
expect(chat.sendMessage.calledWith({ text: 'What is sharding?' })).to.be
158+
.true;
146159
});
147160

148-
it('does not call onSendMessage when input is empty or whitespace-only', function () {
149-
let messageSent = false;
150-
const handleSendMessage = () => {
151-
messageSent = true;
152-
};
153-
154-
render(<AssistantChat messages={[]} onSendMessage={handleSendMessage} />);
161+
it('does not call sendMessage when input is empty or whitespace-only', function () {
162+
const { chat } = renderWithChat([]);
155163

156164
const inputField = screen.getByTestId('assistant-chat-input');
157165
const chatForm = screen.getByTestId('assistant-chat-form');
158166

159167
// Test empty input
160168
userEvent.click(chatForm);
161-
expect(messageSent).to.be.false;
169+
expect(chat.sendMessage.notCalled).to.be.true;
162170

163171
// Test whitespace-only input
164172
userEvent.type(inputField, ' ');
165173
userEvent.click(chatForm);
166-
expect(messageSent).to.be.false;
174+
expect(chat.sendMessage.notCalled).to.be.true;
167175
});
168176

169177
it('displays user and assistant messages with different styling', function () {
170-
render(<AssistantChat messages={mockMessages} />);
178+
renderWithChat(mockMessages);
171179

172180
const userMessage = screen.getByTestId('assistant-message-user');
173181
const assistantMessage = screen.getByTestId('assistant-message-assistant');
@@ -196,7 +204,7 @@ describe('AssistantChat', function () {
196204
},
197205
];
198206

199-
render(<AssistantChat messages={messagesWithMultipleParts} />);
207+
renderWithChat(messagesWithMultipleParts);
200208

201209
expect(screen.getByText('Here is part 1. And here is part 2.')).to.exist;
202210
});
@@ -215,7 +223,7 @@ describe('AssistantChat', function () {
215223
},
216224
];
217225

218-
render(<AssistantChat messages={messagesWithMixedParts} />);
226+
renderWithChat(messagesWithMixedParts);
219227

220228
expect(screen.getByText('This is text content. More text content.')).to
221229
.exist;

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

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,33 @@
11
import React, { useCallback, useState } from 'react';
22
import type { UIMessage } from './@ai-sdk/react/use-chat';
3+
import type { Chat } from './@ai-sdk/react/chat-react';
4+
import { useChat } from './@ai-sdk/react/use-chat';
35

46
interface AssistantChatProps {
5-
messages: UIMessage[];
6-
onSendMessage?: (message: string) => void;
7+
chat: Chat<UIMessage>;
78
}
89

910
/**
1011
* This component is currently using placeholders as Leafygreen UI updates are not available yet.
1112
* Before release, we will replace this with the actual Leafygreen chat components.
1213
*/
1314
export const AssistantChat: React.FunctionComponent<AssistantChatProps> = ({
14-
messages,
15-
onSendMessage,
15+
chat,
1616
}) => {
1717
const [inputValue, setInputValue] = useState('');
18+
const { messages, sendMessage } = useChat({
19+
chat,
20+
});
1821

1922
const handleInputSubmit = useCallback(
2023
(e: React.FormEvent) => {
2124
e.preventDefault();
22-
if (inputValue.trim() && onSendMessage) {
23-
onSendMessage(inputValue.trim());
25+
if (inputValue.trim()) {
26+
void sendMessage({ text: inputValue.trim() });
2427
setInputValue('');
2528
}
2629
},
27-
[inputValue, onSendMessage]
30+
[inputValue, sendMessage]
2831
);
2932

3033
return (

‎packages/compass-assistant/src/assistant-drawer.tsx‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import React, { useContext } from 'react';
2+
import { DrawerSection } from '@mongodb-js/compass-components';
3+
import { AssistantChat } from './assistant-chat';
4+
import { ASSISTANT_DRAWER_ID, AssistantContext } from './assistant-provider';
5+
6+
/**
7+
* AssistantDrawer component that wraps AssistantChat in a DrawerSection.
8+
* This component can be placed at any level in the component tree as long as
9+
* it's within an AssistantProvider.
10+
*/
11+
export const AssistantDrawer: React.FunctionComponent = () => {
12+
const context = useContext(AssistantContext);
13+
14+
if (!context) {
15+
throw new Error(
16+
'CompassAssistantDrawer must be used within an CompassAssistantProvider'
17+
);
18+
}
19+
20+
if (!context.isEnabled) {
21+
return null;
22+
}
23+
24+
const { chat } = context;
25+
26+
return (
27+
<DrawerSection
28+
id={ASSISTANT_DRAWER_ID}
29+
title="MongoDB Assistant"
30+
label="MongoDB Assistant"
31+
glyph="Sparkle"
32+
>
33+
<AssistantChat chat={chat} />
34+
</DrawerSection>
35+
);
36+
};

0 commit comments

Comments
 (0)