Add limit and refactor

src/participant/participant.ts

@@ -45,6 +45,7 @@ import { processStreamWithIdentifiers } from './streamParsing';
 import type { PromptIntent } from './prompts/intent';
 import type { DataService } from 'mongodb-data-service';
 import { ParticipantErrorTypes } from './participantErrorTypes';
+import { PromptHistory } from './prompts/promptHistory';
 
 const log = createLogger('participant');
 
@@ -1448,7 +1449,7 @@ export default class ParticipantController {
       log.info('Docs chatbot created for chatId', chatId);
     }
 
-    const history = Prompts.docs.getHistoryMessages({
+    const history = PromptHistory.getFilteredHistoryForDocs({
       connectionNames: this._getConnectionNames(),
       context: context,
     });

src/participant/prompts/index.ts

@@ -5,7 +5,6 @@ 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';
 
@@ -17,7 +16,6 @@ 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

@@ -1,10 +1,10 @@
 import * as vscode from 'vscode';
-import type { ChatResult, ParticipantResponseType } from '../constants';
+import type { ChatResult } from '../constants';
 import type {
   InternalPromptPurpose,
   ParticipantPromptProperties,
 } from '../../telemetry/telemetryService';
-import { ParticipantErrorTypes } from '../participantErrorTypes';
+import { PromptHistory } from './promptHistory';
 
 export interface PromptArgsBase {
   request: {
@@ -108,7 +108,7 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
   }
 
   async buildMessages(args: TArgs): Promise<ModelInput> {
-    let historyMessages = this.getFilteredHistoryMessages({
+    let historyMessages = PromptHistory.getFilteredHistory({
       history: args.context?.history,
       ...args,
     });
@@ -180,146 +180,4 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
       internal_purpose: this.internalPurposeForTelemetry,
     };
   }
-
-  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.
-  protected getFilteredHistoryMessages({
-    connectionNames,
-    history,
-    databaseName,
-    collectionName,
-  }: {
-    connectionNames?: string[]; // Used to scrape the connecting messages from the history.
-    history?: vscode.ChatContext['history'];
-    databaseName?: string;
-    collectionName?: string;
-  }): vscode.LanguageModelChatMessage[] {
-    const messages: vscode.LanguageModelChatMessage[] = [];
-
-    if (!history) {
-      return [];
-    }
-
-    let previousItem:
-      | vscode.ChatRequestTurn
-      | vscode.ChatResponseTurn
-      | undefined = undefined;
-
-    const namespaceIsKnown =
-      databaseName !== undefined && collectionName !== undefined;
-    for (const historyItem of history) {
-      if (historyItem instanceof vscode.ChatRequestTurn) {
-        this._handleChatRequestTurn({
-          messages,
-          historyItem,
-          previousItem,
-          connectionNames,
-          namespaceIsKnown,
-        });
-      } else if (historyItem instanceof vscode.ChatResponseTurn) {
-        this._handleChatResponseTurn({
-          messages,
-          historyItem,
-          namespaceIsKnown,
-        });
-      }
-      previousItem = historyItem;
-    }
-
-    return messages;
-  }
 }