feat(chat): filter namespace messages from history if it exists in th…

…e metadata VSCODE-611 (#866)
mongodb-js · Nov 12, 2024 · dd80613 · dd80613
1 parent d5a9345
commit dd80613
Show file tree

Hide file tree

Showing 3 changed files with 308 additions and 167 deletions.
diff --git a/src/participant/prompts/promptBase.ts b/src/participant/prompts/promptBase.ts
@@ -13,6 +13,8 @@ export interface PromptArgsBase {
   };
   context?: vscode.ChatContext;
   connectionNames?: string[];
+  databaseName?: string;
+  collectionName?: string;
 }
 
 export interface UserPromptResponse {
@@ -163,16 +165,27 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
   protected getHistoryMessages({
     connectionNames,
     context,
+    databaseName,
+    collectionName,
   }: {
     connectionNames?: string[]; // Used to scrape the connecting messages from the history.
     context?: vscode.ChatContext;
+    databaseName?: string;
+    collectionName?: string;
   }): vscode.LanguageModelChatMessage[] {
     const messages: vscode.LanguageModelChatMessage[] = [];
 
     if (!context) {
       return [];
     }
 
+    let previousItem:
+      | vscode.ChatRequestTurn
+      | vscode.ChatResponseTurn
+      | undefined = undefined;
+
+    const namespaceIsKnown =
+      databaseName !== undefined && collectionName !== undefined;
     for (const historyItem of context.history) {
       if (historyItem instanceof vscode.ChatRequestTurn) {
         if (
@@ -181,9 +194,21 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
         ) {
           // 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));
       }
@@ -206,11 +231,17 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
           'emptyRequest',
           'askToConnect',
         ];
-        if (
-          responseTypesToSkip.indexOf(
-            (historyItem.result as ChatResult)?.metadata?.intent
-          ) > -1
-        ) {
+
+        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;
         }
 
@@ -232,6 +263,7 @@ export abstract class PromptBase<TArgs extends PromptArgsBase> {
         // eslint-disable-next-line new-cap
         messages.push(vscode.LanguageModelChatMessage.Assistant(message));
       }
+      previousItem = historyItem;
     }
 
     return messages;