Merge pull request #2878 from continuedev/nate/autocomplete-cleaning

Nate/autocomplete cleaning

Author: sestinj

core/autocomplete/context/root-path-context/RootPathContextService.ts

@@ -4,7 +4,7 @@ import { LRUCache } from "lru-cache";
 import Parser from "web-tree-sitter";
 
 import { IDE } from "../../..";
-import { getQueryForFile, TSQueryType } from "../../../util/treeSitter";
+import { getQueryForFile } from "../../../util/treeSitter";
 import { AstPath } from "../../util/ast";
 import { ImportDefinitionsService } from "../ImportDefinitionsService";
 import { AutocompleteSnippet } from "../ranking";
@@ -27,6 +27,7 @@ export class RootPathContextService {
     "program",
     "function_declaration",
     "method_definition",
+    "class_declaration",
   ]);
 
   /**
@@ -54,23 +55,12 @@ export class RootPathContextService {
       case "program":
         this.importDefinitionsService.get(filepath);
         break;
-      case "function_declaration":
+      default:
         query = await getQueryForFile(
           filepath,
-          TSQueryType.FunctionDeclaration,
+          `root-path-context-queries/${node.type}`,
         );
         break;
-      case "method_definition":
-        query = await getQueryForFile(filepath, TSQueryType.MethodDefinition);
-        break;
-      case "function_definition":
-        query = await getQueryForFile(filepath, TSQueryType.FunctionDefinition);
-        break;
-      case "method_declaration":
-        query = await getQueryForFile(filepath, TSQueryType.MethodDeclaration);
-        break;
-      default:
-        break;
     }
 
     if (!query) {
@@ -79,28 +69,37 @@ export class RootPathContextService {
 
     await Promise.all(
       query.matches(node).map(async (match) => {
-        const startPosition = match.captures[0].node.startPosition;
-        const endPosition = match.captures[0].node.endPosition;
-        const definitions = await this.ide.gotoDefinition({
-          filepath,
-          position: {
-            line: endPosition.row,
-            character: endPosition.column,
-          },
-        });
-        const newSnippets = await Promise.all(
-          definitions.map(async (def) => ({
-            ...def,
-            contents: await this.ide.readRangeInFile(def.filepath, def.range),
-          })),
-        );
-        snippets.push(...newSnippets);
+        for (const item of match.captures) {
+          const endPosition = item.node.endPosition;
+          const newSnippets = await this.getSnippets(filepath, endPosition);
+          snippets.push(...newSnippets);
+        }
       }),
     );
 
     return snippets;
   }
 
+  private async getSnippets(
+    filepath: string,
+    endPosition: Parser.Point,
+  ): Promise<AutocompleteSnippet[]> {
+    const definitions = await this.ide.gotoDefinition({
+      filepath,
+      position: {
+        line: endPosition.row,
+        character: endPosition.column,
+      },
+    });
+    const newSnippets = await Promise.all(
+      definitions.map(async (def) => ({
+        ...def,
+        contents: await this.ide.readRangeInFile(def.filepath, def.range),
+      })),
+    );
+    return newSnippets;
+  }
+
   async getContextForPath(
     filepath: string,
     astPath: AstPath,

core/autocomplete/context/root-path-context/test/RootPathContextService.test.ts

@@ -0,0 +1,42 @@
+import { testRootPathContext } from "./testUtils";
+
+const TEST_CASES = [
+  {
+    description: "function",
+    fileName: "file1.ts",
+    range: {
+      start: { line: 10, character: 2 },
+      end: { line: 10, character: 24 },
+    },
+    positions: [
+      { row: 9, column: 34 }, // Person
+      { row: 9, column: 44 }, // Address
+    ],
+  },
+  {
+    description: "class method",
+    fileName: "file1.ts",
+    range: {
+      start: { line: 22, character: 4 },
+      end: { line: 22, character: 30 },
+    },
+    positions: [
+      { row: 13, column: 29 }, // BaseClass
+      { row: 13, column: 55 }, // FirstInterface
+      { row: 13, column: 72 }, // SecondInterface
+      { row: 21, column: 33 }, // Person
+      { row: 21, column: 43 }, // Address
+    ],
+  },
+];
+
+describe("RootPathContextService", () => {
+  describe("TypeScript should return expected snippets when editing inside a:", () => {
+    test.each(TEST_CASES)(
+      "should look for correct type definitions when editing inside a $description",
+      async ({ fileName, range, positions }) => {
+        await testRootPathContext("typescript", fileName, range, positions);
+      },
+    );
+  });
+});

core/autocomplete/context/root-path-context/RootPathContextService.test.ts core/autocomplete/context/root-path-context/test/testUtils.ts

@@ -1,12 +1,13 @@
+import { jest } from "@jest/globals";
 import fs from "fs";
 import path from "path";
 
-import { Range } from "../../..";
-import { testIde } from "../../../test/util/fixtures";
-import { getAst, getTreePathAtCursor } from "../../util/ast";
-import { ImportDefinitionsService } from "../ImportDefinitionsService";
-
-import { RootPathContextService } from "./RootPathContextService";
+import Parser from "web-tree-sitter";
+import { Range } from "../../../..";
+import { testIde } from "../../../../test/util/fixtures";
+import { getAst, getTreePathAtCursor } from "../../../util/ast";
+import { ImportDefinitionsService } from "../../ImportDefinitionsService";
+import { RootPathContextService } from "../RootPathContextService";
 
 function splitTextAtRange(fileContent: string, range: Range): [string, string] {
   const lines = fileContent.split("\n");
@@ -42,12 +43,21 @@ export async function testRootPathContext(
   folderName: string,
   relativeFilepath: string,
   rangeToFill: Range,
-  expectedSnippets: string[],
+  expectedDefinitionPositions: Parser.Point[],
 ) {
+  // Create a mocked instance of RootPathContextService
   const ide = testIde;
   const importDefinitionsService = new ImportDefinitionsService(ide);
   const service = new RootPathContextService(importDefinitionsService, ide);
 
+  const getSnippetsMock = jest
+    // @ts-ignore
+    .spyOn(service, "getSnippets")
+    // @ts-ignore
+    .mockImplementation(async (_filepath, _endPosition) => {
+      return [];
+    });
+
   // Copy the folder to the test directory
   const folderPath = path.join(
     __dirname,
@@ -77,23 +87,17 @@ export async function testRootPathContext(
   }
 
   const treePath = await getTreePathAtCursor(ast, prefix.length);
-  const snippets = await service.getContextForPath(startPath, treePath);
+  await service.getContextForPath(startPath, treePath);
 
-  expectedSnippets.forEach((expectedSnippet) => {
-    const found = snippets.find((snippet) =>
-      snippet.contents.includes(expectedSnippet),
-    );
-    expect(found).toBeDefined();
-  });
-}
+  expect(getSnippetsMock).toHaveBeenCalledTimes(
+    expectedDefinitionPositions.length,
+  );
 
-describe("RootPathContextService", () => {
-  it.skip("should be true", async () => {
-    await testRootPathContext(
-      "typescript",
-      "file1.ts",
-      { start: { line: 3, character: 2 }, end: { line: 3, character: 24 } },
-      ["export interface Person", "export interface Address"],
+  expectedDefinitionPositions.forEach((position, index) => {
+    expect(getSnippetsMock).toHaveBeenNthCalledWith(
+      index + 1,
+      expect.any(String), // filepath argument
+      position,
     );
   });
-});
+}

core/autocomplete/context/root-path-context/test/typescript/file1.ts

@@ -1,13 +1,25 @@
-import { Address, Person } from "./types";
+import {
+  Address,
+  Person,
+  BaseClass,
+  FirstInterface,
+  SecondInterface,
+  // @ts-ignore
+} from "./types";
 
 function getAddress(person: Person): Address {
   return person.address;
 }
 
-class Group {
+class Group extends BaseClass implements FirstInterface, SecondInterface {
   people: Person[];
 
   constructor(people: Person[]) {
+    super();
     this.people = people;
   }
+
+  getPersonAddress(person: Person): Address {
+    return getAddress(person);
+  }
 }

core/autocomplete/context/root-path-context/test/typescript/types.ts

@@ -1,11 +1,10 @@
 import { execSync } from "child_process";
+import * as JSONC from "comment-json";
 import * as fs from "fs";
 import os from "os";
 import path from "path";
 
-import * as JSONC from "comment-json";
 import * as tar from "tar";
-
 import {
   BrowserSerializedContinueConfig,
   Config,
@@ -72,7 +71,6 @@ import {
 } from "./promptFile.js";
 import { ConfigValidationError, validateConfig } from "./validation.js";
 
-
 export interface ConfigResult<T> {
   config: T | undefined;
   errors: ConfigValidationError[] | undefined;
@@ -138,6 +136,13 @@ function loadSerializedConfig(
     config.allowAnonymousTelemetry = true;
   }
 
+  if (config.ui?.getChatTitles === undefined) {
+    config.ui = {
+      ...config.ui,
+      getChatTitles: true,
+    };
+  }
+
   if (ideSettings.remoteConfigServerUrl) {
     try {
       const remoteConfigJson = resolveSerializedConfig(

core/config/load.ts

@@ -944,6 +944,7 @@ export interface ContinueUIConfig {
   fontSize?: number;
   displayRawMarkdown?: boolean;
   showChatScrollbar?: boolean;
+  getChatTitles?: boolean;
 }
 
 interface ContextMenuConfig {
@@ -1006,11 +1007,6 @@ interface ExperimentalConfig {
    */
   readResponseTTS?: boolean;
 
-  /**
-   * Prompt the user's LLM for a title given the current chat content
-   */
-  getChatTitles?: boolean;
-
   /**
    * If set to true, we will attempt to pull down and install an instance of Chromium
    * that is compatible with the current version of Puppeteer.

core/index.d.ts

@@ -3,8 +3,8 @@ import { stripImages } from "../llm/images";
 
 import { removeQuotesAndEscapes } from ".";
 
-import type { IMessenger } from "./messenger";
 import type { FromCoreProtocol, ToCoreProtocol } from "../protocol";
+import type { IMessenger } from "./messenger";
 
 /**
  * Removes code blocks from a message.
@@ -42,7 +42,7 @@ export class ChatDescriber {
       return;
     }
 
-    completionOptions.maxTokens = 24;
+    completionOptions.maxTokens = 6;
 
     // Prompt the user's current LLM for the title
     const titleResponse = await model.chat(

core/util/chatDescriber.ts

@@ -140,7 +140,7 @@ export enum TSQueryType {
 
 export async function getQueryForFile(
   filepath: string,
-  queryType: TSQueryType,
+  queryPath: string,
 ): Promise<Parser.Query | undefined> {
   const language = await getLanguageForFile(filepath);
   if (!language) {
@@ -154,7 +154,7 @@ export async function getQueryForFile(
     ...(process.env.NODE_ENV === "test"
       ? ["extensions", "vscode", "tree-sitter"]
       : ["tree-sitter"]),
-    queryType,
+    queryPath,
     `${fullLangName}.scm`,
   );
   if (!fs.existsSync(sourcePath)) {

core/util/treeSitter.ts

@@ -0,0 +1,19 @@
+(
+    (class_declaration
+        (class_heritage
+            (extends_clause
+                (identifier) @base-class
+            )
+        )
+    )
+)
+
+(
+    (class_declaration
+        (class_heritage
+            (implements_clause
+                (type_identifier) @interface
+            )
+        )
+    )
+)

extensions/vscode/tree-sitter/root-path-context-queries/class_declaration/typescript.scm

@@ -2,8 +2,9 @@
     (function_declaration 
         (formal_parameters 
             (_ 
-                (type_annotation) @type
+                (type_annotation) @param_type
             )
         )
+        (type_annotation) @return_type
     )
 )

extensions/vscode/tree-sitter/root-path-context-queries/function_declaration/typescript.scm

@@ -2,8 +2,9 @@
     (method_definition
         (formal_parameters 
             (_ 
-                (type_annotation) @type
+                (type_annotation) @param_type
             )
         )
+        (type_annotation) @return_type
     )
-)
+)