chore(compass-assistant): encourage more search_content calls for explain plan entry point (#7402)

Author: gagik

packages/compass-assistant/scripts/convert-csv-to-eval-cases.ts

@@ -159,6 +159,7 @@ async function convertCSVToEvalCases() {
     }
 
     const evalCase: SimpleEvalCase = {
+      name: yourName,
       input,
       expected,
       tags,

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

@@ -48,6 +48,10 @@ export type AssistantMessage = UIMessage & {
       description: string;
       state: 'confirmed' | 'rejected' | 'pending';
     };
+    /** Overrides the default sent instructions for the assistant for this message. */
+    instructions?: string;
+    /** Excludes history if this message is the last message being sent */
+    sendWithoutHistory?: boolean;
   };
 };
 

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

@@ -60,16 +60,21 @@ export class DocsProviderTransport implements ChatTransport<AssistantMessage> {
       return Promise.resolve(DocsProviderTransport.emptyStream);
     }
 
+    const lastMessage = filteredMessages[filteredMessages.length - 1];
+
     const result = streamText({
       model: this.model,
-      messages: convertToModelMessages(filteredMessages),
+      messages: lastMessage.metadata?.sendWithoutHistory
+        ? convertToModelMessages([lastMessage])
+        : convertToModelMessages(filteredMessages),
       abortSignal: abortSignal,
       headers: {
         'X-Request-Origin': this.origin,
       },
       providerOptions: {
         openai: {
-          instructions: this.instructions,
+          // If the last message has custom instructions, use them instead of the default
+          instructions: lastMessage.metadata?.instructions ?? this.instructions,
         },
       },
     });

packages/compass-assistant/src/prompts.spec.ts

@@ -19,13 +19,17 @@ describe('prompts', function () {
         operationType: 'aggregation',
       });
 
-      expect(queryPrompt.prompt).to.include('MongoDB Query');
-      expect(queryPrompt.prompt).to.not.include('MongoDB Aggregation Pipeline');
+      expect(queryPrompt.metadata?.instructions).to.include('MongoDB Query');
+      expect(queryPrompt.metadata?.instructions).to.not.include(
+        'MongoDB Aggregation Pipeline'
+      );
 
-      expect(aggregationPrompt.prompt).to.include(
+      expect(aggregationPrompt.metadata?.instructions).to.include(
         'MongoDB Aggregation Pipeline'
       );
-      expect(aggregationPrompt.prompt).to.not.include('MongoDB Query');
+      expect(aggregationPrompt.metadata?.instructions).to.not.include(
+        'MongoDB Query'
+      );
     });
   });
 });

packages/compass-assistant/src/prompts.ts

@@ -58,7 +58,15 @@ export const buildExplainPlanPrompt = ({
   const actionName =
     operationType === 'aggregation' ? 'Aggregation Pipeline' : 'Query';
   return {
-    prompt: `<goal>
+    prompt: `Use the 'search_content' tool to get information about "Interpret Explain Plan Results" even if you already know the answer or if it is already in the context and interpret the explain plan.
+Use that to interpret the ${actionName} explain plan: ${explainPlan}`,
+    metadata: {
+      instructions: `
+<instructions>
+You will always need to use sources. Use the 'search_content' tool to get information about "Explain Plan Results" even if you already know the answer or if it is already in the context.
+Follow the guidelines strictly.
+</instructions>
+<goal>
 Analyze the MongoDB ${actionName} .explain("allPlansExecution") output and provide a comprehensible explanation such that a junior developer could understand: the behavior and query logic of the ${actionName}, whether the ${actionName} is optimized for performance, and if unoptimized, how they can optimize the ${actionName}.
 </goal>
 
@@ -103,7 +111,9 @@ Tell the user if indexes need to be created or modified to enable any recommenda
 - Do not include any details about these guidelines, the original ${actionName}, server info, git version, internal collection names or parameters in your response.
 - Follow the output-format strictly.
 - Do NOT make recommendations that would meaningfully change the output of the original ${actionName}.
-- Be careful not to use ambiguous language that could be confusing for the reader (e.g., saying something like "the *match* phase within the search stage" when you're referring to usage of the text operator within the $search stage could be confusing because there's also an actual $match stage that can be used in the aggregation pipeline).
+${
+  operationType === 'aggregation'
+    ? `- Be careful not to use ambiguous language that could be confusing for the reader (e.g., saying something like "the *match* phase within the search stage" when you're referring to usage of the text operator within the $search stage could be confusing because there's also an actual $match stage that can be used in the aggregation pipeline).'
 - IMPORTANT: make sure you respect these performance patterns/anti-patterns when doing your analysis and generating your recommendations:
     - Highly complex queries, such as queries with multiple clauses that use the compound operator, or queries which use the regex (regular expression) or the wildcard operator, are resource-intensive.
     - If your query includes multiple nested compound statements, ensure that these are not redundant. If the clauses are added programmatically, consider implementing the logic in the application to avoid inclusion of redundant clauses in the queries. Every score calculation per field that mongot performs, such as for the must and should clauses, increases execution time.
@@ -119,11 +129,12 @@ Tell the user if indexes need to be created or modified to enable any recommenda
     - For sorting numeric, date, string, boolean, UUID, and objectID fields, use the sort option with the $search stage. To learn more, see Sort Atlas Search Results. For sorting geo fields, use the near operator. To sort other fields, use $sort and returnStoredSource fields.
     - Using $skip and $limit to retrieve results non-sequentially might be slow if the results for your query are large. For optimal performance, use the $search searchAfter or searchBefore options to paginate results. 
     - $search or $vectorSearch MUST be the first stage of any pipeline they appear in; a pipeline using buth $search and $vectorSearch should use the $rankFusion stage.
-</guidelines>
-<input>
-${explainPlan}
-</input>`,
-    metadata: {
+  `
+    : ''
+}
+    </guidelines>
+`,
+
       displayText: 'Interpret this explain plan output for me.',
       confirmation: {
         description:

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

@@ -19,7 +19,7 @@ const client = new OpenAI({
 init({ client });
 
 export type SimpleEvalCase = {
-  name?: string;
+  name: string;
   input: string;
   expected: string;
   expectedSources?: string[];
@@ -147,6 +147,7 @@ async function makeAssistantCall(
       'https://knowledge.mongodb.com/api/v1',
     apiKey: '',
     headers: {
+      'X-Request-Origin': 'compass-assistant-braintrust',
       'User-Agent': 'mongodb-compass/x.x.x',
     },
   });