Merge pull request #6055 from continuedev/adarsh/enhancement/deduplication

Simple Context Deduplication for Autocomplete

Author: adarshramiyer

core/autocomplete/templating/filtering.ts

@@ -3,6 +3,7 @@ import { SnippetPayload } from "../snippets";
 import {
   AutocompleteCodeSnippet,
   AutocompleteSnippet,
+  AutocompleteSnippetType,
 } from "../snippets/types";
 import { HelperVars } from "../util/HelperVars";
 import { formatOpenedFilesContext } from "./formatOpenedFilesContext";
@@ -130,21 +131,16 @@ export const getSnippets = (
   const finalSnippets = [];
   let remainingTokenCount = getRemainingTokenCount(helper);
 
+  // tracks already added filepaths for deduplication
+  const addedFilepaths = new Set<string>();
+
   // Process snippets in priority order
   for (const { key } of snippetOrder) {
     // Special handling for recentlyOpenedFiles
     if (key === "recentlyOpenedFiles" && helper.options.useRecentlyOpened) {
-      const recentlyOpenedFilesSnippets =
-        payload.recentlyOpenedFileSnippets.filter(
-          (snippet) =>
-            !(snippet as AutocompleteCodeSnippet).filepath?.startsWith(
-              "output:extension-output-Continue.continue",
-            ),
-        );
-
       // Custom trimming
       const processedSnippets = formatOpenedFilesContext(
-        recentlyOpenedFilesSnippets,
+        payload.recentlyOpenedFileSnippets,
         remainingTokenCount,
         helper,
         finalSnippets,
@@ -160,18 +156,18 @@ export const getSnippets = (
 
         if (remainingTokenCount >= snippetSize) {
           finalSnippets.push(snippet);
+          addedFilepaths.add(snippet.filepath);
           remainingTokenCount -= snippetSize;
         } else {
-          break; // Out of tokens
+          continue; // Not enough tokens, try again with next snippet
         }
       }
     } else {
       // Normal processing for other snippet types
       const snippetsToProcess = snippets[key].filter(
         (snippet) =>
-          !(snippet as AutocompleteCodeSnippet).filepath?.startsWith(
-            "output:extension-output-Continue.continue",
-          ),
+          snippet.type !== AutocompleteSnippetType.Code ||
+          !addedFilepaths.has(snippet.filepath),
       );
 
       for (const snippet of snippetsToProcess) {
@@ -182,9 +178,14 @@ export const getSnippets = (
 
         if (remainingTokenCount >= snippetSize) {
           finalSnippets.push(snippet);
+
+          if ((snippet as AutocompleteCodeSnippet).filepath) {
+            addedFilepaths.add((snippet as AutocompleteCodeSnippet).filepath);
+          }
+
           remainingTokenCount -= snippetSize;
         } else {
-          break; // Out of tokens
+          continue; // Not enough tokens, try again with next snippet
         }
       }
     }

core/autocomplete/templating/formatOpenedFilesContext.ts

@@ -2,6 +2,7 @@ import { countTokens, pruneStringFromBottom } from "../../llm/countTokens";
 import {
   AutocompleteCodeSnippet,
   AutocompleteSnippet,
+  AutocompleteSnippetType,
 } from "../snippets/types";
 import { HelperVars } from "../util/HelperVars";
 
@@ -20,13 +21,23 @@ export function formatOpenedFilesContext(
   recentlyOpenedFilesSnippets: AutocompleteCodeSnippet[],
   remainingTokenCount: number,
   helper: HelperVars,
-  alreadyAddedSnippets: AutocompleteSnippet[], // TODO use this to deduplicate context
+  alreadyAddedSnippets: AutocompleteSnippet[],
   TOKEN_BUFFER: number,
 ): AutocompleteCodeSnippet[] {
   if (recentlyOpenedFilesSnippets.length === 0) {
     return [];
   }
 
+  // deduplication; if a snippet is already added, don't include it here
+  for (const snippet of alreadyAddedSnippets) {
+    if (snippet.type !== AutocompleteSnippetType.Code) {
+      continue;
+    }
+    recentlyOpenedFilesSnippets = recentlyOpenedFilesSnippets.filter(
+      (s) => s.filepath !== snippet.filepath,
+    );
+  }
+
   // Calculate how many full snippets would fit within the remaining token count
   let numSnippetsThatFit = 0;
   let totalTokens = 0;

core/autocomplete/templating/validation.ts

@@ -1,5 +1,6 @@
 import {
   AutocompleteClipboardSnippet,
+  AutocompleteCodeSnippet,
   AutocompleteSnippet,
   AutocompleteSnippetType,
 } from "../snippets/types";
@@ -25,5 +26,13 @@ export const isValidSnippet = (snippet: AutocompleteSnippet): boolean => {
     return isValidClipboardSnippet(snippet);
   }
 
+  if (
+    (snippet as AutocompleteCodeSnippet).filepath?.startsWith(
+      "output:extension-output-Continue.continue",
+    )
+  ) {
+    return false;
+  }
+
   return true;
 };