feat(compass-assistant): add global per-conversation instructions COMPASS-9766 (#7279)

* Add global per-conversation instructions

* add Gagik's prompt tweak

* extend stub for test

* rather pass in a prop to the provider

* also use the assistant prompt in the braintrust evals

Author: lerouxb

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

@@ -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} />
@@ -100,7 +100,9 @@ describe('useAssistantActions', function () {
 
       return (
         <DrawerContentProvider>
-          <MockedProvider chat={chat}>{children}</MockedProvider>
+          <MockedProvider appNameForPrompt="MongoDB Compass" chat={chat}>
+            {children}
+          </MockedProvider>
         </DrawerContentProvider>
       );
     }
@@ -494,7 +496,7 @@ describe('CompassAssistantProvider', function () {
       render(
         <DrawerContentProvider>
           <DrawerAnchor />
-          <MockedProvider />
+          <MockedProvider appNameForPrompt="MongoDB Compass" />
         </DrawerContentProvider>,
         {
           preferences: {

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

@@ -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';
 
@@ -141,6 +142,7 @@ export type CompassAssistantService = {
 
 export const AssistantProvider: React.FunctionComponent<
   PropsWithChildren<{
+    appNameForPrompt: string;
     chat: Chat<AssistantMessage>;
     atlasAiService: AtlasAiService;
   }>
@@ -223,10 +225,12 @@ export const CompassAssistantProvider = registerCompassPlugin(
   {
     name: 'CompassAssistant',
     component: ({
+      appNameForPrompt,
       chat,
       atlasAiService,
       children,
     }: PropsWithChildren<{
+      appNameForPrompt: string;
       chat?: Chat<AssistantMessage>;
       atlasAiService?: AtlasAiService;
     }>) => {
@@ -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>
       );
@@ -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(

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

@@ -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({
@@ -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());

packages/compass-assistant/src/prompts.ts

@@ -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;
@@ -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.',
   };
 };

packages/compass-assistant/test/assistant.eval.ts

@@ -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',
@@ -33,6 +34,7 @@ type ExpectedMessage = OutputMessage;
 
 type ConversationEvalCaseInput = {
   messages: InputMessage[];
+  instructions: Message;
 };
 
 type ConversationEvalCaseExpected = {
@@ -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 || [] }],
@@ -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[] = [];