feat(chat): Include history from messages to docs chatbot VSCODE-632

src/participant/participant.ts

@@ -40,7 +40,7 @@ import {
 import { DocsChatbotAIService } from './docsChatbotAIService';
 import type TelemetryService from '../telemetry/telemetryService';
 import formatError from '../utils/formatError';
-import type { ModelInput } from './prompts/promptBase';
+import { getContent, type ModelInput } from './prompts/promptBase';
 import { processStreamWithIdentifiers } from './streamParsing';
 import type { PromptIntent } from './prompts/intent';
 import type { DataService } from 'mongodb-data-service';
@@ -1415,10 +1415,12 @@ export default class ParticipantController {
     chatId,
     token,
     stream,
+    context,
   }: {
     prompt: string;
     chatId: string;
     token: vscode.CancellationToken;
+    context: vscode.ChatContext;
     stream: vscode.ChatResponseStream;
   }): Promise<{
     responseContent: string;
@@ -1446,8 +1448,22 @@ export default class ParticipantController {
       log.info('Docs chatbot created for chatId', chatId);
     }
 
+    const history = Prompts.docs.getHistoryMessages({
+      connectionNames: this._getConnectionNames(),
+      context: context,
+    });
+
+    const previousMessages =
+      history.length > 0
+        ? `${history
+            .map((message: vscode.LanguageModelChatMessage) =>
+              getContent(message)
+            )
+            .join('\n\n')}\n\n`
+        : '';
+
     const response = await this._docsChatbotAIService.addMessage({
-      message: prompt,
+      message: `${previousMessages}${prompt}`,
       conversationId: docsChatbotConversationId,
       signal: abortController.signal,
     });
@@ -1553,6 +1569,7 @@ export default class ParticipantController {
         chatId,
         token,
         stream,
+        context,
       });
 
       if (docsResult.responseContent) {

src/participant/prompts/docs.ts

@@ -0,0 +1,44 @@
+import type * as vscode from 'vscode';
+import type { PromptArgsBase } from './promptBase';
+import { PromptBase } from './promptBase';
+
+export class DocsPrompt extends PromptBase<PromptArgsBase> {
+  protected getAssistantPrompt(): string {
+    throw new Error('Not used with docs command');
+  }
+
+  getHistoryMessages({
+    connectionNames,
+    context,
+    databaseName,
+    collectionName,
+  }: {
+    connectionNames?: string[];
+    context?: vscode.ChatContext;
+    databaseName?: string;
+    collectionName?: string;
+  }): vscode.LanguageModelChatMessage[] {
+    if (!context) {
+      return [];
+    }
+    const historySinceLastDocs: (
+      | vscode.ChatRequestTurn
+      | vscode.ChatResponseTurn
+    )[] = [];
+
+    for (let i = context.history.length - 1; i >= 0; i--) {
+      const message = context.history[i];
+
+      if (message.command === 'docs') {
+        break;
+      }
+      historySinceLastDocs.push(context.history[i]);
+    }
+    return this.getFilteredHistoryMessages({
+      connectionNames,
+      history: historySinceLastDocs.reverse(),
+      databaseName,
+      collectionName,
+    });
+  }
+}

src/participant/prompts/index.ts

@@ -5,6 +5,7 @@ import { IntentPrompt } from './intent';
 import { NamespacePrompt } from './namespace';
 import { QueryPrompt } from './query';
 import { SchemaPrompt } from './schema';
+import { DocsPrompt } from './docs';
 import { ExportToPlaygroundPrompt } from './exportToPlayground';
 import { isContentEmpty } from './promptBase';
 
@@ -16,6 +17,7 @@ export class Prompts {
   public static namespace = new NamespacePrompt();
   public static query = new QueryPrompt();
   public static schema = new SchemaPrompt();
+  public static docs = new DocsPrompt();
   public static exportToPlayground = new ExportToPlaygroundPrompt();
 
   public static isPromptEmpty(request: vscode.ChatRequest): boolean {

src/participant/prompts/promptBase.ts

@@ -53,6 +53,26 @@ export function getContentLength(
   return 0;
 }
 
+export function getContent(message: vscode.LanguageModelChatMessage): string {
+  const content = message.content as any;
+  if (typeof content === 'string') {
+    return content;
+  }
+
+  if (Array.isArray(content)) {
+    return content.reduce((agg: string, element) => {
+      const value = element?.value ?? element?.content?.value;
+      if (typeof value === 'string') {
+        return agg + value;
+      }
+
+      return agg;
+    }, '');
+  }
+
+  return '';
+}
+
 export function isContentEmpty(
   message: vscode.LanguageModelChatMessage
 ): boolean {
@@ -88,7 +108,10 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
   }
 
   async buildMessages(args: TArgs): Promise<ModelInput> {
-    let historyMessages = this.getHistoryMessages(args);
+    let historyMessages = this.getFilteredHistoryMessages({
+      history: args.context?.history,
+      ...args,
+    });
     // If the current user's prompt is a connection name, and the last
     // message was to connect. We want to use the last
     // message they sent before the connection name as their prompt.
@@ -158,24 +181,116 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
     };
   }
 
+  private _handleChatResponseTurn({
+    messages,
+    historyItem,
+    namespaceIsKnown,
+  }: {
+    historyItem: vscode.ChatResponseTurn;
+    messages: vscode.LanguageModelChatMessage[];
+    namespaceIsKnown: boolean;
+  }): void {
+    if (
+      historyItem.result.errorDetails?.message ===
+      ParticipantErrorTypes.FILTERED
+    ) {
+      // If the response led to a filtered error, we do not want the
+      // error-causing message to be sent again so we remove it.
+      messages.pop();
+      return;
+    }
+
+    let message = '';
+
+    // Skip a response to an empty user prompt message or connect message.
+    const responseTypesToSkip: ParticipantResponseType[] = [
+      'emptyRequest',
+      'askToConnect',
+    ];
+
+    const responseType = (historyItem.result as ChatResult)?.metadata?.intent;
+    if (responseTypesToSkip.includes(responseType)) {
+      return;
+    }
+
+    // If the namespace is already known, skip including prompts asking for it.
+    if (responseType === 'askForNamespace' && namespaceIsKnown) {
+      return;
+    }
+
+    for (const fragment of historyItem.response) {
+      if (fragment instanceof vscode.ChatResponseMarkdownPart) {
+        message += fragment.value.value;
+
+        if (
+          (historyItem.result as ChatResult)?.metadata?.intent ===
+          'askForNamespace'
+        ) {
+          // When the message is the assistant asking for part of a namespace,
+          // we only want to include the question asked, not the user's
+          // database and collection names in the history item.
+          break;
+        }
+      }
+    }
+
+    // eslint-disable-next-line new-cap
+    messages.push(vscode.LanguageModelChatMessage.Assistant(message));
+  }
+
+  private _handleChatRequestTurn({
+    messages,
+    historyItem,
+    previousItem,
+    connectionNames,
+    namespaceIsKnown,
+  }: {
+    historyItem: vscode.ChatRequestTurn;
+    previousItem: vscode.ChatRequestTurn | vscode.ChatResponseTurn | undefined;
+    messages: vscode.LanguageModelChatMessage[];
+    connectionNames: string[] | undefined;
+    namespaceIsKnown: boolean;
+  }): void {
+    if (
+      historyItem.prompt?.trim().length === 0 ||
+      connectionNames?.includes(historyItem.prompt)
+    ) {
+      // When the message is empty or a connection name then we skip it.
+      // It's probably going to be the response to the connect step.
+      return;
+    }
+
+    if (previousItem instanceof vscode.ChatResponseTurn) {
+      const responseIntent = (previousItem.result as ChatResult).metadata
+        ?.intent;
+
+      if (responseIntent === 'askForNamespace' && namespaceIsKnown) {
+        // If the namespace is already known, skip responses to prompts asking for it.
+        return;
+      }
+    }
+
+    // eslint-disable-next-line new-cap
+    messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt));
+  }
+
   // When passing the history to the model we only want contextual messages
   // to be passed. This function parses through the history and returns
   // the messages that are valuable to keep.
-  // eslint-disable-next-line complexity
-  protected getHistoryMessages({
+  protected getFilteredHistoryMessages({
     connectionNames,
-    context,
+    history,
     databaseName,
     collectionName,
   }: {
     connectionNames?: string[]; // Used to scrape the connecting messages from the history.
-    context?: vscode.ChatContext;
+    history?: vscode.ChatContext['history'];
     databaseName?: string;
     collectionName?: string;
   }): vscode.LanguageModelChatMessage[] {
     const messages: vscode.LanguageModelChatMessage[] = [];
 
-    if (!context) {
+    if (!history) {
       return [];
     }
 
@@ -186,82 +301,21 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
 
     const namespaceIsKnown =
       databaseName !== undefined && collectionName !== undefined;
-    for (const historyItem of context.history) {
+    for (const historyItem of history) {
       if (historyItem instanceof vscode.ChatRequestTurn) {
-        if (
-          historyItem.prompt?.trim().length === 0 ||
-          connectionNames?.includes(historyItem.prompt)
-        ) {
-          // When the message is empty or a connection name then we skip it.
-          // It's probably going to be the response to the connect step.
-          previousItem = historyItem;
-          continue;
-        }
-
-        if (previousItem instanceof vscode.ChatResponseTurn) {
-          const responseIntent = (previousItem.result as ChatResult).metadata
-            ?.intent;
-
-          // If the namespace is already known, skip responses to prompts asking for it.
-          if (responseIntent === 'askForNamespace' && namespaceIsKnown) {
-            previousItem = historyItem;
-            continue;
-          }
-        }
-
-        // eslint-disable-next-line new-cap
-        messages.push(vscode.LanguageModelChatMessage.User(historyItem.prompt));
-      }
-
-      if (historyItem instanceof vscode.ChatResponseTurn) {
-        if (
-          historyItem.result.errorDetails?.message ===
-          ParticipantErrorTypes.FILTERED
-        ) {
-          // If the response led to a filtered error, we do not want the
-          // error-causing message to be sent again so we remove it.
-          messages.pop();
-          continue;
-        }
-
-        let message = '';
-
-        // Skip a response to an empty user prompt message or connect message.
-        const responseTypesToSkip: ParticipantResponseType[] = [
-          'emptyRequest',
-          'askToConnect',
-        ];
-
-        const responseType = (historyItem.result as ChatResult)?.metadata
-          ?.intent;
-        if (responseTypesToSkip.includes(responseType)) {
-          previousItem = historyItem;
-          continue;
-        }
-
-        // If the namespace is already known, skip including prompts asking for it.
-        if (responseType === 'askForNamespace' && namespaceIsKnown) {
-          previousItem = historyItem;
-          continue;
-        }
-
-        for (const fragment of historyItem.response) {
-          if (fragment instanceof vscode.ChatResponseMarkdownPart) {
-            message += fragment.value.value;
-
-            if (
-              (historyItem.result as ChatResult)?.metadata?.intent ===
-              'askForNamespace'
-            ) {
-              // When the message is the assistant asking for part of a namespace,
-              // we only want to include the question asked, not the user's
-              // database and collection names in the history item.
-              break;
-            }
-          }
-        }
-        // eslint-disable-next-line new-cap
-        messages.push(vscode.LanguageModelChatMessage.Assistant(message));
+        this._handleChatRequestTurn({
+          messages,
+          historyItem,
+          previousItem,
+          connectionNames,
+          namespaceIsKnown,
+        });
+      } else if (historyItem instanceof vscode.ChatResponseTurn) {
+        this._handleChatResponseTurn({
+          messages,
+          historyItem,
+          namespaceIsKnown,
+        });
       }
       previousItem = historyItem;
     }