feat(ai-assistant): Assistant Eval improvements

packages/compass-assistant/package.json

@@ -49,7 +49,7 @@
     "reformat": "npm run eslint . -- --fix && npm run prettier -- --write ."
   },
   "dependencies": {
-    "@ai-sdk/openai": "^2.0.4",
+    "@ai-sdk/openai": "^2.0.22",
     "@mongodb-js/atlas-service": "^0.56.0",
     "@mongodb-js/compass-app-registry": "^9.4.20",
     "@mongodb-js/compass-components": "^1.49.0",

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

@@ -9,7 +9,7 @@ import {
 import { atlasServiceLocator } from '@mongodb-js/atlas-service/provider';
 import { DocsProviderTransport } from './docs-provider-transport';
 import { useDrawerActions } from '@mongodb-js/compass-components';
-import { buildConnectionErrorPrompt, buildExplainPlanPrompt } from './prompts';
+import { buildPrompts } from './prompts';
 import { usePreference } from 'compass-preferences-model/provider';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import type { ConnectionInfo } from '@mongodb-js/connection-info';
@@ -21,6 +21,7 @@ export type AssistantMessage = UIMessage & {
   metadata?: {
     /** The text to display instead of the message text. */
     displayText?: string;
+    type?: keyof typeof buildPrompts;
   };
 };
 
@@ -91,14 +92,15 @@ export const AssistantProvider: React.FunctionComponent<
   const assistantActionsContext = useRef<AssistantActionsContextType>({
     interpretExplainPlan: ({ explainPlan }) => {
       openDrawer(ASSISTANT_DRAWER_ID);
-      const { prompt, displayText } = buildExplainPlanPrompt({
+      const { prompt, displayText } = buildPrompts['explain-plan'].user({
         explainPlan,
       });
       void chat.sendMessage(
         {
           text: prompt,
           metadata: {
             displayText,
+            type: 'explain-plan',
           },
         },
         {}
@@ -112,13 +114,16 @@ export const AssistantProvider: React.FunctionComponent<
       );
       const connectionError = error.toString();
 
-      const { prompt } = buildConnectionErrorPrompt({
+      const { prompt } = buildPrompts['connection-error'].user({
         connectionString,
         connectionError,
       });
       void chat.sendMessage(
         {
           text: prompt,
+          metadata: {
+            type: 'connection-error',
+          },
         },
         {}
       );

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

@@ -1,30 +1,39 @@
+import { createOpenAI } from '@ai-sdk/openai';
+import { buildPrompts } from './prompts';
 import {
   type ChatTransport,
+  convertToModelMessages,
   type UIMessage,
   type UIMessageChunk,
-  convertToModelMessages,
+  type LanguageModel,
+  type ModelMessage,
   streamText,
 } from 'ai';
-import { createOpenAI } from '@ai-sdk/openai';
 
 export class DocsProviderTransport implements ChatTransport<UIMessage> {
-  private openai: ReturnType<typeof createOpenAI>;
+  private createResponse: ReturnType<typeof makeCreateResponse>;
 
   constructor({ baseUrl }: { baseUrl: string }) {
-    this.openai = createOpenAI({
+    const openai = createOpenAI({
       baseURL: baseUrl,
       apiKey: '',
     });
+    this.createResponse = makeCreateResponse({
+      languageModel: openai.responses('mongodb-chat-latest'),
+    });
   }
 
   sendMessages({
     messages,
+    // TODO: pass the metadata correctly in a strongly typed manner.
+    // I'm not sure how to do this.
+    metadata,
     abortSignal,
   }: Parameters<ChatTransport<UIMessage>['sendMessages']>[0]) {
-    const result = streamText({
-      model: this.openai.responses('mongodb-chat-latest'),
+    const result = this.createResponse({
       messages: convertToModelMessages(messages),
-      abortSignal: abortSignal,
+      abortSignal,
+      systemPromptType: metadata?.type as keyof typeof buildPrompts,
     });
 
     return Promise.resolve(result.toUIMessageStream());
@@ -35,3 +44,40 @@ export class DocsProviderTransport implements ChatTransport<UIMessage> {
     return Promise.resolve(null);
   }
 }
+
+export function makeCreateResponse({
+  languageModel,
+}: {
+  languageModel: LanguageModel;
+}) {
+  return function createResponse({
+    messages,
+    abortSignal,
+    systemPromptType,
+    headers,
+  }: {
+    messages: ModelMessage[];
+    abortSignal?: AbortSignal;
+    systemPromptType?: keyof typeof buildPrompts;
+    headers?: Record<string, string>;
+  }) {
+    const instructions =
+      systemPromptType && buildPrompts[systemPromptType]?.system
+        ? buildPrompts[systemPromptType].system
+        : undefined;
+
+    const result = streamText({
+      model: languageModel,
+      messages: messages,
+      abortSignal: abortSignal,
+      headers,
+      providerOptions: {
+        openai: {
+          // Have to pass like this because it only accepts JSONValue (not instructions: undefined).
+          ...(instructions ? { instructions } : {}),
+        },
+      },
+    });
+    return result;
+  };
+}

packages/compass-assistant/src/index.tsx

@@ -5,3 +5,7 @@ export {
   compassAssistantServiceLocator,
 } from './compass-assistant-provider';
 export type { CompassAssistantService } from './compass-assistant-provider';
+export type {
+  DocsProviderTransport,
+  makeCreateResponse,
+} from './docs-provider-transport';

packages/compass-assistant/src/prompts.ts

@@ -1,4 +1,6 @@
-export const buildExplainPlanPrompt = ({
+const explainPlanSystemPrompt = `TODO:....add custom prompt stuff here`;
+
+const buildExplainPlanUserPrompt = ({
   explainPlan,
 }: {
   explainPlan: string;
@@ -12,7 +14,9 @@ ${explainPlan}`,
   };
 };
 
-export const buildConnectionErrorPrompt = ({
+const connectionErrorSystemPrompt = `TODO:....add custom prompt stuff here`;
+
+const buildConnectionUserErrorPrompt = ({
   connectionString,
   connectionError,
 }: {
@@ -29,3 +33,14 @@ Error message:
 ${connectionError}`,
   };
 };
+
+export const buildPrompts = {
+  'explain-plan': {
+    user: buildExplainPlanUserPrompt,
+    system: explainPlanSystemPrompt,
+  },
+  'connection-error': {
+    user: buildConnectionUserErrorPrompt,
+    system: connectionErrorSystemPrompt,
+  },
+} as const;