Skip to content

Commit 8722ef7

Browse files
gagikgagik
committed
chore: add tests, refactor docs provider
1 parent f4aec0c commit 8722ef7

File tree

5 files changed

+245
-34
lines changed

5 files changed

+245
-34
lines changed

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

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
2929
import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
3030
import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
3131
import { buildConversationInstructionsPrompt } from './prompts';
32+
import { createOpenAI } from '@ai-sdk/openai';
3233

3334
export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
3435

@@ -177,18 +178,11 @@ export const AssistantProvider: React.FunctionComponent<
177178
return;
178179
}
179180

180-
const { prompt, displayText } = builder(props);
181+
const { prompt, metadata } = builder(props);
181182
void assistantActionsContext.current.ensureOptInAndSend(
182183
{
183184
text: prompt,
184-
metadata: {
185-
displayText,
186-
confirmation: {
187-
description:
188-
'Explain plan metadata, including the original query, may be used to process your request',
189-
state: 'pending',
190-
},
191-
},
185+
metadata,
192186
},
193187
{},
194188
() => {
@@ -286,10 +280,13 @@ export const CompassAssistantProvider = registerCompassPlugin(
286280
initialProps.chat ??
287281
new Chat({
288282
transport: new DocsProviderTransport({
289-
baseUrl: atlasService.assistantApiEndpoint(),
290283
instructions: buildConversationInstructionsPrompt({
291284
target: initialProps.appNameForPrompt,
292285
}),
286+
model: createOpenAI({
287+
baseURL: atlasService.assistantApiEndpoint(),
288+
apiKey: '',
289+
}).responses('mongodb-chat-latest'),
293290
}),
294291
onError: (err: Error) => {
295292
logger.log.error(

‎packages/compass-assistant/src/components/confirmation-message.spec.tsx‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ describe('ConfirmationMessage', function () {
9797
render(<ConfirmationMessage {...defaultProps} state="confirmed" />);
9898

9999
expect(screen.getByText('Request confirmed')).to.exist;
100-
// This is sic from the icon library
100+
// sic from the icon library
101101
expect(screen.getByLabelText('Checkmark With Circle Icon')).to.exist;
102102

103103
expect(screen.queryByText('Confirm')).to.not.exist;
@@ -108,7 +108,8 @@ describe('ConfirmationMessage', function () {
108108
render(<ConfirmationMessage {...defaultProps} state="rejected" />);
109109

110110
expect(screen.getByText('Request cancelled')).to.exist;
111-
expect(screen.getByLabelText('X With Circle Icon')).to.exist;
111+
// sic from the icon library
112+
expect(screen.getByLabelText('XWith Circle Icon')).to.exist;
112113

113114
expect(screen.queryByText('Confirm')).to.not.exist;
114115
expect(screen.queryByText('Cancel')).to.not.exist;

‎packages/compass-assistant/src/docs-provider-transport.spec.ts‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
import { expect } from 'chai';
2+
import sinon from 'sinon';
3+
import {
4+
DocsProviderTransport,
5+
shouldExcludeMessage,
6+
} from './docs-provider-transport';
7+
import type { AssistantMessage } from './compass-assistant-provider';
8+
import { MockLanguageModelV2 } from 'ai/test';
9+
import type { UIMessageChunk } from 'ai';
10+
import { waitFor } from '@mongodb-js/testing-library-compass';
11+
12+
describe.only('DocsProviderTransport', function () {
13+
describe('shouldExcludeMessage', function () {
14+
it('returns false for messages without confirmation metadata', function () {
15+
const message: AssistantMessage = {
16+
id: 'test-1',
17+
role: 'user',
18+
parts: [{ type: 'text', text: 'Hello' }],
19+
};
20+
21+
expect(shouldExcludeMessage(message)).to.be.false;
22+
});
23+
24+
it('returns true for confirmation messages', function () {
25+
const message: AssistantMessage = {
26+
id: 'test-5',
27+
role: 'assistant',
28+
parts: [{ type: 'text', text: 'Response' }],
29+
metadata: {
30+
confirmation: {
31+
state: 'pending',
32+
description: 'Confirm this action',
33+
},
34+
},
35+
};
36+
37+
expect(shouldExcludeMessage(message)).to.be.true;
38+
});
39+
});
40+
41+
describe('sending messages', function () {
42+
let mockModel: MockLanguageModelV2;
43+
let doStream: sinon.SinonStub;
44+
let transport: DocsProviderTransport;
45+
let abortController: AbortController;
46+
let sendMessages: (
47+
params: Partial<Parameters<typeof transport.sendMessages>[0]>
48+
) => Promise<ReadableStream<UIMessageChunk>>;
49+
50+
beforeEach(function () {
51+
// Mock the OpenAI client
52+
doStream = sinon.stub().returns({
53+
stream: DocsProviderTransport.emptyStream,
54+
request: {
55+
body: {
56+
messages: [],
57+
},
58+
},
59+
});
60+
mockModel = new MockLanguageModelV2({
61+
doStream,
62+
});
63+
abortController = new AbortController();
64+
transport = new DocsProviderTransport({
65+
instructions: 'Test instructions for MongoDB assistance',
66+
model: mockModel,
67+
});
68+
sendMessages = (params) =>
69+
transport.sendMessages({
70+
trigger: 'submit-message',
71+
chatId: 'test-chat',
72+
messageId: undefined,
73+
abortSignal: abortController.signal,
74+
messages: [],
75+
...params,
76+
});
77+
});
78+
79+
afterEach(function () {
80+
sinon.restore();
81+
});
82+
83+
describe('sendMessages', function () {
84+
const userMessage: AssistantMessage = {
85+
id: 'included1',
86+
role: 'user',
87+
parts: [{ type: 'text', text: 'User message' }],
88+
};
89+
const confirmationPendingMessage: AssistantMessage = {
90+
id: 'test',
91+
role: 'assistant',
92+
parts: [{ type: 'text', text: 'Response' }],
93+
metadata: {
94+
confirmation: {
95+
state: 'pending',
96+
description: 'Confirm this action',
97+
},
98+
},
99+
};
100+
const confirmationConfirmedMessage: AssistantMessage = {
101+
id: 'test',
102+
role: 'assistant',
103+
parts: [{ type: 'text', text: 'Response' }],
104+
metadata: {
105+
confirmation: {
106+
state: 'confirmed',
107+
description: 'Confirmed action',
108+
},
109+
},
110+
};
111+
it('returns empty stream when last message should be excluded', async function () {
112+
const messages: AssistantMessage[] = [
113+
userMessage,
114+
confirmationConfirmedMessage,
115+
];
116+
117+
const result = await sendMessages({
118+
messages,
119+
});
120+
121+
expect(result).to.equal(DocsProviderTransport.emptyStream);
122+
expect(mockModel.doStreamCalls).to.be.empty;
123+
});
124+
125+
it('returns empty stream when all messages are filtered out', async function () {
126+
const messages: AssistantMessage[] = [
127+
confirmationPendingMessage,
128+
confirmationConfirmedMessage,
129+
];
130+
131+
const result = await sendMessages({
132+
messages,
133+
});
134+
135+
expect(result).to.equal(DocsProviderTransport.emptyStream);
136+
expect(mockModel.doStreamCalls).to.be.empty;
137+
});
138+
139+
it('sends filtered messages to AI when valid messages exist', async function () {
140+
await sendMessages({
141+
messages: [userMessage],
142+
});
143+
144+
await waitFor(() => {
145+
expect(doStream).to.have.been.calledOnce;
146+
expect(doStream.firstCall.args[0]).to.deep.include({
147+
prompt: [
148+
{
149+
role: 'user',
150+
providerOptions: undefined,
151+
content: [
152+
{
153+
type: 'text',
154+
text: 'User message',
155+
providerOptions: undefined,
156+
},
157+
],
158+
},
159+
],
160+
});
161+
});
162+
});
163+
164+
it('sends only valid messages when confirmation required messages exist', async function () {
165+
await sendMessages({
166+
messages: [
167+
confirmationConfirmedMessage,
168+
confirmationPendingMessage,
169+
userMessage,
170+
],
171+
});
172+
173+
await waitFor(() => {
174+
expect(doStream).to.have.been.calledOnce;
175+
expect(doStream.firstCall.args[0]).to.deep.include({
176+
prompt: [
177+
{
178+
role: 'user',
179+
providerOptions: undefined,
180+
content: [
181+
{
182+
type: 'text',
183+
text: 'User message',
184+
providerOptions: undefined,
185+
},
186+
],
187+
},
188+
],
189+
});
190+
});
191+
});
192+
});
193+
194+
// We currently do not support reconnecting to streams but we may want to in the future
195+
describe('reconnectToStream', function () {
196+
it('always returns null', async function () {
197+
const result = await transport.reconnectToStream();
198+
expect(result).to.be.null;
199+
});
200+
});
201+
});
202+
});

‎packages/compass-assistant/src/docs-provider-transport.ts‎

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,36 +1,33 @@
11
import {
22
type ChatTransport,
3+
type LanguageModel,
34
type UIMessageChunk,
45
convertToModelMessages,
56
streamText,
67
} from 'ai';
7-
import { createOpenAI } from '@ai-sdk/openai';
88
import type { AssistantMessage } from './compass-assistant-provider';
99

1010
/** Returns true if the message should be excluded from being sent to the assistant API. */
1111
export function shouldExcludeMessage({ metadata }: AssistantMessage) {
1212
if (metadata?.confirmation) {
13-
return metadata.confirmation.state !== 'confirmed';
13+
return true;
1414
}
1515
return false;
1616
}
1717

1818
export class DocsProviderTransport implements ChatTransport<AssistantMessage> {
19-
private openai: ReturnType<typeof createOpenAI>;
19+
private model: LanguageModel;
2020
private instructions: string;
2121

2222
constructor({
23-
baseUrl,
2423
instructions,
24+
model,
2525
}: {
26-
baseUrl: string;
2726
instructions: string;
27+
model: LanguageModel;
2828
}) {
29-
this.openai = createOpenAI({
30-
baseURL: baseUrl,
31-
apiKey: '',
32-
});
3329
this.instructions = instructions;
30+
this.model = model;
3431
}
3532

3633
static emptyStream = new ReadableStream<UIMessageChunk>({
@@ -60,7 +57,7 @@ export class DocsProviderTransport implements ChatTransport<AssistantMessage> {
6057
}
6158

6259
const result = streamText({
63-
model: this.openai.responses('mongodb-chat-latest'),
60+
model: this.model,
6461
messages: convertToModelMessages(filteredMessages),
6562
abortSignal: abortSignal,
6663
providerOptions: {

0 commit comments

Comments
 (0)