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
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ const TestComponent: React.FunctionComponent<{

return (
<DrawerContentProvider>
<MockedProvider chat={chat}>
<MockedProvider appNameForPrompt="MongoDB Compass" chat={chat}>
<DrawerAnchor>
<div data-testid="provider-children">Provider children</div>
<CompassAssistantDrawer autoOpen={autoOpen} />
Expand All @@ -100,7 +100,9 @@ describe('useAssistantActions', function () {

return (
<DrawerContentProvider>
<MockedProvider chat={chat}>{children}</MockedProvider>
<MockedProvider appNameForPrompt="MongoDB Compass" chat={chat}>
{children}
</MockedProvider>
</DrawerContentProvider>
);
}
Expand Down Expand Up @@ -494,7 +496,7 @@ describe('CompassAssistantProvider', function () {
render(
<DrawerContentProvider>
<DrawerAnchor />
<MockedProvider />
<MockedProvider appNameForPrompt="MongoDB Compass" />
</DrawerContentProvider>,
{
preferences: {
Expand Down
13 changes: 12 additions & 1 deletion packages/compass-assistant/src/compass-assistant-provider.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import type { ConnectionInfo } from '@mongodb-js/connection-info';
import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
import { buildConversationInstructionsPrompt } from './prompts';

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

Expand Down Expand Up @@ -141,6 +142,7 @@ export type CompassAssistantService = {

export const AssistantProvider: React.FunctionComponent<
PropsWithChildren<{
appNameForPrompt: string;
chat: Chat<AssistantMessage>;
atlasAiService: AtlasAiService;
}>
Expand Down Expand Up @@ -223,10 +225,12 @@ export const CompassAssistantProvider = registerCompassPlugin(
{
name: 'CompassAssistant',
component: ({
appNameForPrompt,
chat,
atlasAiService,
children,
}: PropsWithChildren<{
appNameForPrompt: string;
chat?: Chat<AssistantMessage>;
atlasAiService?: AtlasAiService;
}>) => {
Expand All @@ -237,7 +241,11 @@ export const CompassAssistantProvider = registerCompassPlugin(
throw new Error('atlasAiService was not provided by the state');
}
return (
<AssistantProvider chat={chat} atlasAiService={atlasAiService}>
<AssistantProvider
appNameForPrompt={appNameForPrompt}
chat={chat}
atlasAiService={atlasAiService}
>
{children}
</AssistantProvider>
);
Expand All @@ -248,6 +256,9 @@ export const CompassAssistantProvider = registerCompassPlugin(
new Chat({
transport: new DocsProviderTransport({
baseUrl: atlasService.assistantApiEndpoint(),
instructions: buildConversationInstructionsPrompt({
target: initialProps.appNameForPrompt,
}),
}),
onError: (err: Error) => {
logger.log.error(
Expand Down
15 changes: 14 additions & 1 deletion packages/compass-assistant/src/docs-provider-transport.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,20 @@ import { createOpenAI } from '@ai-sdk/openai';

export class DocsProviderTransport implements ChatTransport<UIMessage> {
private openai: ReturnType<typeof createOpenAI>;
private instructions: string;

constructor({ baseUrl }: { baseUrl: string }) {
constructor({
baseUrl,
instructions,
}: {
baseUrl: string;
instructions: string;
}) {
this.openai = createOpenAI({
baseURL: baseUrl,
apiKey: '',
});
this.instructions = instructions;
}

sendMessages({
Expand All @@ -25,6 +33,11 @@ export class DocsProviderTransport implements ChatTransport<UIMessage> {
model: this.openai.responses('mongodb-chat-latest'),
messages: convertToModelMessages(messages),
abortSignal: abortSignal,
providerOptions: {
openai: {
instructions: this.instructions,
},
},
});

return Promise.resolve(result.toUIMessageStream());
Expand Down
12 changes: 11 additions & 1 deletion packages/compass-assistant/src/prompts.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,16 @@
import type { ConnectionInfo } from '@mongodb-js/connection-info';
import { redactConnectionString } from 'mongodb-connection-string-url';

export const buildConversationInstructionsPrompt = ({
target,
}: {
target: string;
}) => {
// TODO: we'll want to greatly expand on this, but at minimum this is where we
// make the distinction between running inside Data Explorer vs Compass.
return `You are an assistant running in a side-panel inside ${target}. Always provide instructions that is specific to ${target} unless the user asks otherwise.`;
};

export type EntryPointMessage = {
prompt: string;
displayText?: string;
Expand Down Expand Up @@ -78,7 +88,7 @@ Tell the user if indexes need to be created or modified to enable any recommenda
</guidelines>
<input>
${explainPlan}
</explain-plan>`,
</input>`,
displayText: 'Interpret this explain plan output for me.',
};
};
Expand Down
28 changes: 27 additions & 1 deletion packages/compass-assistant/test/assistant.eval.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import { evalCases } from './eval-cases';
import { fuzzyLinkMatch } from './fuzzylinkmatch';
import { binaryNdcgAtK } from './binaryndcgatk';
import { makeEntrypointCases } from './entrypoints';
import { buildConversationInstructionsPrompt } from '../src/prompts';

const client = new OpenAI({
baseURL: 'https://api.braintrust.dev/v1/proxy',
Expand All @@ -33,6 +34,7 @@ type ExpectedMessage = OutputMessage;

type ConversationEvalCaseInput = {
messages: InputMessage[];
instructions: Message;
};

type ConversationEvalCaseExpected = {
Expand Down Expand Up @@ -79,13 +81,32 @@ function getScorerTemperature(): number | undefined {
}

function makeEvalCases(): ConversationEvalCase[] {
const entrypointCases: ConversationEvalCase[] = makeEntrypointCases();
const instructions = buildConversationInstructionsPrompt({
target: 'MongoDB Compass',
});

const entrypointCases: ConversationEvalCase[] = makeEntrypointCases().map(
(c) => {
return {
name: c.name ?? c.input,
input: {
messages: [{ text: c.input }],
instructions: { text: instructions },
},
expected: {
messages: [{ text: c.expected, sources: c.expectedSources || [] }],
},
metadata: {},
};
}
);

const userCases: ConversationEvalCase[] = evalCases.map((c) => {
return {
name: c.name ?? c.input,
input: {
messages: [{ text: c.input }],
instructions: { text: instructions },
},
expected: {
messages: [{ text: c.expected, sources: c.expectedSources || [] }],
Expand Down Expand Up @@ -115,6 +136,11 @@ async function makeAssistantCall(
model: openai.responses('mongodb-chat-latest'),
temperature: getChatTemperature(),
prompt,
providerOptions: {
openai: {
instructions: input.instructions.text,
},
},
});

const chunks: string[] = [];
Expand Down
Loading
Loading