Autocomplete improvements

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 @@
   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 @@
   }
 
   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/config/load.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/types.ts

@@ -628,6 +628,11 @@ declare global {
     | "mistral-large-latest"
     | "mistral-7b"
     | "mistral-8x7b"
+    | "mistral-8x22b"
+    | "mistral-tiny"
+    | "mistral-small"
+    | "mistral-medium"
+    | "mistral-nemo"
     // Llama 2
     | "llama2-7b"
     | "llama2-13b"
@@ -641,6 +646,8 @@ declare global {
     | "llama3-70b"
     // Other Open-source
     | "phi2"
+    | "phi-3-mini"
+    | "phi-3-medium"
     | "phind-codellama-34b"
     | "wizardcoder-7b"
     | "wizardcoder-13b"
@@ -649,9 +656,12 @@ declare global {
     | "codeup-13b"
     | "deepseek-7b"
     | "deepseek-33b"
+    | "deepseek-2-lite"
     | "neural-chat-7b"
     | "gemma-7b-it"
+    | "gemma2-2b-it"
     | "gemma2-9b-it"
+    | "olmo-7b"
     // Anthropic
     | "claude-3-5-sonnet-latest"
     | "claude-3-5-sonnet-20240620"
@@ -669,10 +679,6 @@ declare global {
     | "gemini-1.5-pro"
     | "gemini-1.5-flash-latest"
     | "gemini-1.5-flash"
-    // Mistral
-    | "mistral-tiny"
-    | "mistral-small"
-    | "mistral-medium"
     // Tab autocomplete
     | "deepseek-1b"
     | "starcoder-1b"