fix(chat): update response handling to stream and inline code block p…

…arsing VSCODE-620 (#835)
mongodb-js · Sep 26, 2024 · 7eef0e7 · 7eef0e7
1 parent 118aee0
commit 7eef0e7
Show file tree

Hide file tree

Showing 10 changed files with 508 additions and 112 deletions.
diff --git a/src/participant/constants.ts b/src/participant/constants.ts
@@ -14,6 +14,11 @@ export type ParticipantResponseType =
   | 'askToConnect'
   | 'askForNamespace';
 
+export const codeBlockIdentifier = {
+  start: '```javascript',
+  end: '```',
+};
+
 interface Metadata {
   intent: Exclude<ParticipantResponseType, 'askForNamespace' | 'docs'>;
   chatId: string;

diff --git a/src/participant/participant.ts b/src/participant/participant.ts
@@ -21,6 +21,7 @@ import {
   docsRequestChatResult,
   schemaRequestChatResult,
   createCancelledRequestChatResult,
+  codeBlockIdentifier,
 } from './constants';
 import { SchemaFormatter } from './schema';
 import { getSimplifiedSampleDocuments } from './sampleDocuments';
@@ -38,7 +39,8 @@ import {
 } from '../telemetry/telemetryService';
 import { DocsChatbotAIService } from './docsChatbotAIService';
 import type TelemetryService from '../telemetry/telemetryService';
-import { IntentPrompt, type PromptIntent } from './prompts/intent';
+import { processStreamWithIdentifiers } from './streamParsing';
+import type { PromptIntent } from './prompts/intent';
 
 const log = createLogger('participant');
 
@@ -59,16 +61,6 @@ export type ParticipantCommand = '/query' | '/schema' | '/docs';
 
 const MAX_MARKDOWN_LIST_LENGTH = 10;
 
-export function getRunnableContentFromString(text: string): string {
-  const matchedJSresponseContent = text.match(/```javascript((.|\n)*)```/);
-
-  const code =
-    matchedJSresponseContent && matchedJSresponseContent.length > 1
-      ? matchedJSresponseContent[1]
-      : '';
-  return code.trim();
-}
-
 export default class ParticipantController {
   _participant?: vscode.ChatParticipant;
   _connectionController: ConnectionController;
@@ -171,48 +163,113 @@ export default class ParticipantController {
     });
   }
 
-  async getChatResponseContent({
+  async _getChatResponse({
     messages,
     token,
   }: {
     messages: vscode.LanguageModelChatMessage[];
     token: vscode.CancellationToken;
-  }): Promise<string> {
+  }): Promise<vscode.LanguageModelChatResponse> {
     const model = await getCopilotModel();
-    let responseContent = '';
-    if (model) {
-      const chatResponse = await model.sendRequest(messages, {}, token);
-      for await (const fragment of chatResponse.text) {
-        responseContent += fragment;
-      }
+
+    if (!model) {
+      throw new Error('Copilot model not found');
     }
 
-    return responseContent;
+    return await model.sendRequest(messages, {}, token);
   }
 
-  _streamRunnableContentActions({
-    responseContent,
+  async streamChatResponse({
+    messages,
     stream,
+    token,
   }: {
-    responseContent: string;
+    messages: vscode.LanguageModelChatMessage[];
+    stream: vscode.ChatResponseStream;
+    token: vscode.CancellationToken;
+  }): Promise<void> {
+    const chatResponse = await this._getChatResponse({
+      messages,
+      token,
+    });
+    for await (const fragment of chatResponse.text) {
+      stream.markdown(fragment);
+    }
+  }
+
+  _streamCodeBlockActions({
+    runnableContent,
+    stream,
+  }: {
+    runnableContent: string;
     stream: vscode.ChatResponseStream;
   }): void {
-    const runnableContent = getRunnableContentFromString(responseContent);
-    if (runnableContent) {
-      const commandArgs: RunParticipantQueryCommandArgs = {
-        runnableContent,
-      };
-      stream.button({
-        command: EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY,
-        title: vscode.l10n.t('▶️ Run'),
-        arguments: [commandArgs],
-      });
-      stream.button({
-        command: EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND,
-        title: vscode.l10n.t('Open in playground'),
-        arguments: [commandArgs],
-      });
+    runnableContent = runnableContent.trim();
+
+    if (!runnableContent) {
+      return;
     }
+
+    const commandArgs: RunParticipantQueryCommandArgs = {
+      runnableContent,
+    };
+    stream.button({
+      command: EXTENSION_COMMANDS.RUN_PARTICIPANT_QUERY,
+      title: vscode.l10n.t('▶️ Run'),
+      arguments: [commandArgs],
+    });
+    stream.button({
+      command: EXTENSION_COMMANDS.OPEN_PARTICIPANT_QUERY_IN_PLAYGROUND,
+      title: vscode.l10n.t('Open in playground'),
+      arguments: [commandArgs],
+    });
+  }
+
+  async streamChatResponseContentWithCodeActions({
+    messages,
+    stream,
+    token,
+  }: {
+    messages: vscode.LanguageModelChatMessage[];
+    stream: vscode.ChatResponseStream;
+    token: vscode.CancellationToken;
+  }): Promise<void> {
+    const chatResponse = await this._getChatResponse({
+      messages,
+      token,
+    });
+
+    await processStreamWithIdentifiers({
+      processStreamFragment: (fragment: string) => {
+        stream.markdown(fragment);
+      },
+      onStreamIdentifier: (content: string) => {
+        this._streamCodeBlockActions({ runnableContent: content, stream });
+      },
+      inputIterable: chatResponse.text,
+      identifier: codeBlockIdentifier,
+    });
+  }
+
+  // This will stream all of the response content and create a string from it.
+  // It should only be used when the entire response is needed at one time.
+  async getChatResponseContent({
+    messages,
+    token,
+  }: {
+    messages: vscode.LanguageModelChatMessage[];
+    token: vscode.CancellationToken;
+  }): Promise<string> {
+    let responseContent = '';
+    const chatResponse = await this._getChatResponse({
+      messages,
+      token,
+    });
+    for await (const fragment of chatResponse.text) {
+      responseContent += fragment;
+    }
+
+    return responseContent;
   }
 
   async _handleRoutedGenericRequest(
@@ -227,14 +284,9 @@ export default class ParticipantController {
       connectionNames: this._getConnectionNames(),
     });
 
-    const responseContent = await this.getChatResponseContent({
+    await this.streamChatResponseContentWithCodeActions({
       messages,
       token,
-    });
-    stream.markdown(responseContent);
-
-    this._streamRunnableContentActions({
-      responseContent,
       stream,
     });
 
@@ -293,7 +345,7 @@ export default class ParticipantController {
       token,
     });
 
-    return IntentPrompt.getIntentFromModelResponse(responseContent);
+    return Prompts.intent.getIntentFromModelResponse(responseContent);
   }
 
   async handleGenericRequest(
@@ -1001,11 +1053,11 @@ export default class ParticipantController {
       connectionNames: this._getConnectionNames(),
       ...(sampleDocuments ? { sampleDocuments } : {}),
     });
-    const responseContent = await this.getChatResponseContent({
+    await this.streamChatResponse({
       messages,
+      stream,
       token,
     });
-    stream.markdown(responseContent);
 
     stream.button({
       command: EXTENSION_COMMANDS.PARTICIPANT_OPEN_RAW_SCHEMA_OUTPUT,
@@ -1104,16 +1156,11 @@ export default class ParticipantController {
       connectionNames: this._getConnectionNames(),
       ...(sampleDocuments ? { sampleDocuments } : {}),
     });
-    const responseContent = await this.getChatResponseContent({
-      messages,
-      token,
-    });
-
-    stream.markdown(responseContent);
 
-    this._streamRunnableContentActions({
-      responseContent,
+    await this.streamChatResponseContentWithCodeActions({
+      messages,
       stream,
+      token,
     });
 
     return queryRequestChatResult(context.history);
@@ -1181,32 +1228,41 @@ export default class ParticipantController {
       vscode.ChatResponseStream,
       vscode.CancellationToken
     ]
-  ): Promise<{
-    responseContent: string;
-    responseReferences?: Reference[];
-  }> {
-    const [request, context, , token] = args;
+  ): Promise<void> {
+    const [request, context, stream, token] = args;
     const messages = await Prompts.generic.buildMessages({
       request,
       context,
       connectionNames: this._getConnectionNames(),
     });
 
-    const responseContent = await this.getChatResponseContent({
+    await this.streamChatResponseContentWithCodeActions({
       messages,
+      stream,
       token,
     });
-    const responseReferences = [
-      {
+
+    this._streamResponseReference({
+      reference: {
         url: MONGODB_DOCS_LINK,
         title: 'View MongoDB documentation',
       },
-    ];
+      stream,
+    });
+  }
 
-    return {
-      responseContent,
-      responseReferences,
-    };
+  _streamResponseReference({
+    reference,
+    stream,
+  }: {
+    reference: Reference;
+    stream: vscode.ChatResponseStream;
+  }): void {
+    const link = new vscode.MarkdownString(
+      `- [${reference.title}](${reference.url})\n`
+    );
+    link.supportHtml = true;
+    stream.markdown(link);
   }
 
   async handleDocsRequest(
@@ -1235,6 +1291,19 @@ export default class ParticipantController {
         token,
         stream,
       });
+
+      if (docsResult.responseReferences) {
+        for (const reference of docsResult.responseReferences) {
+          this._streamResponseReference({
+            reference,
+            stream,
+          });
+        }
+      }
+
+      if (docsResult.responseContent) {
+        stream.markdown(docsResult.responseContent);
+      }
     } catch (error) {
       // If the docs chatbot API is not available, fall back to Copilot’s LLM and include
       // the MongoDB documentation link for users to go to our documentation site directly.
@@ -1255,25 +1324,7 @@ export default class ParticipantController {
         }
       );
 
-      docsResult = await this._handleDocsRequestWithCopilot(...args);
-    }
-
-    if (docsResult.responseContent) {
-      stream.markdown(docsResult.responseContent);
-      this._streamRunnableContentActions({
-        responseContent: docsResult.responseContent,
-        stream,
-      });
-    }
-
-    if (docsResult.responseReferences) {
-      for (const ref of docsResult.responseReferences) {
-        const link = new vscode.MarkdownString(
-          `- [${ref.title}](${ref.url})\n`
-        );
-        link.supportHtml = true;
-        stream.markdown(link);
-      }
+      await this._handleDocsRequestWithCopilot(...args);
     }
 
     return docsRequestChatResult({

diff --git a/src/participant/prompts/generic.ts b/src/participant/prompts/generic.ts
@@ -3,6 +3,8 @@ import * as vscode from 'vscode';
 import type { PromptArgsBase } from './promptBase';
 import { PromptBase } from './promptBase';
 
+import { codeBlockIdentifier } from '../constants';
+
 export class GenericPrompt extends PromptBase<PromptArgsBase> {
   protected getAssistantPrompt(): string {
     return `You are a MongoDB expert.
@@ -12,7 +14,7 @@ Rules:
 1. Keep your response concise.
 2. You should suggest code that is performant and correct.
 3. Respond with markdown.
-4. When relevant, provide code in a Markdown code block that begins with \`\`\`javascript and ends with \`\`\`.
+4. When relevant, provide code in a Markdown code block that begins with ${codeBlockIdentifier.start} and ends with ${codeBlockIdentifier.end}
 5. Use MongoDB shell syntax for code unless the user requests a specific language.
 6. If you require additional information to provide a response, ask the user for it.
 7. When specifying a database, use the MongoDB syntax use('databaseName').`;

diff --git a/src/participant/prompts/intent.ts b/src/participant/prompts/intent.ts
@@ -34,7 +34,7 @@ Response:
 Docs`;
   }
 
-  static getIntentFromModelResponse(response: string): PromptIntent {
+  getIntentFromModelResponse(response: string): PromptIntent {
     response = response.trim();
     switch (response) {
       case 'Query':