✨ Tab Autocomplete (#758)

* ✨ tab autocomplete attempt 1

* 🚧 WIP autocomplete

* 🚧 wip

* 🔧 config for tab autocomplete

* ⚡️ caching for autocomplete

* 🚧 external snippet formatting

* 🎨 stuff

* stuff

* ⚡️ debouncing and other improvements

* ⚡️ use LSP to get definition during function call

* 💄 status bar for Continue tab-autocomplete

* 🐛 small status bar fix

* 📝 basic docs for tab-autocomplete

* 📝 h2 in docs

* 🚸 Explain how to download Ollama and Deepseek Coder

* 📝 autocomplete troubleshooting

core/autocomplete/cache.ts

@@ -0,0 +1,82 @@
+import { open } from "sqlite";
+import sqlite3 from "sqlite3";
+import { DatabaseConnection } from "../indexing/refreshIndex";
+import { getTabAutocompleteCacheSqlitePath } from "../util/paths";
+
+export class AutocompleteLruCache {
+  private static capacity: number = 1000;
+
+  db: DatabaseConnection;
+
+  constructor(db: DatabaseConnection) {
+    this.db = db;
+  }
+
+  static async get(): Promise<AutocompleteLruCache> {
+    const db = await open({
+      filename: getTabAutocompleteCacheSqlitePath(),
+      driver: sqlite3.Database,
+    });
+
+    await db.run(`
+      CREATE TABLE IF NOT EXISTS cache (
+        key TEXT PRIMARY KEY,
+        value TEXT NOT NULL,
+        timestamp INTEGER NOT NULL
+      )
+    `);
+
+    return new AutocompleteLruCache(db);
+  }
+
+  async get(key: string): Promise<string | undefined> {
+    const result = await this.db.get(
+      "SELECT value FROM cache WHERE key = ?",
+      key
+    );
+
+    if (result) {
+      await this.db.run(
+        "UPDATE cache SET timestamp = ? WHERE key = ?",
+        Date.now(),
+        key
+      );
+      return result.value;
+    }
+
+    return undefined;
+  }
+
+  async put(key: string, value: string) {
+    const result = await this.db.get(
+      "SELECT key FROM cache WHERE key = ?",
+      key
+    );
+
+    if (result) {
+      await this.db.run(
+        "UPDATE cache SET value = ?, timestamp = ? WHERE key = ?",
+        value,
+        Date.now(),
+        key
+      );
+    } else {
+      const count = await this.db.get("SELECT COUNT(*) as count FROM cache");
+
+      if (count.count >= AutocompleteLruCache.capacity) {
+        await this.db.run(
+          "DELETE FROM cache WHERE key = (SELECT key FROM cache ORDER BY timestamp ASC LIMIT 1)"
+        );
+      }
+
+      await this.db.run(
+        "INSERT INTO cache (key, value, timestamp) VALUES (?, ?, ?)",
+        key,
+        value,
+        Date.now()
+      );
+    }
+  }
+}
+
+export default AutocompleteLruCache;

core/autocomplete/constructPrompt.ts

@@ -0,0 +1,128 @@
+import Parser from "web-tree-sitter";
+import {
+  countTokens,
+  pruneLinesFromBottom,
+  pruneLinesFromTop,
+} from "../llm/countTokens";
+import { getBasename } from "../util";
+import { getParserForFile } from "../util/treeSitter";
+import { AutocompleteLanguageInfo, Typescript } from "./languages";
+import {
+  MAX_PROMPT_TOKENS,
+  MAX_SUFFIX_PERCENTAGE,
+  PREFIX_PERCENTAGE,
+} from "./parameters";
+
+export function languageForFilepath(
+  filepath: string
+): AutocompleteLanguageInfo {
+  return Typescript;
+}
+
+function formatExternalSnippet(
+  filepath: string,
+  snippet: string,
+  language: AutocompleteLanguageInfo
+) {
+  const comment = language.comment;
+  const lines = [
+    comment + " Path: " + getBasename(filepath),
+    ...snippet.split("\n").map((line) => comment + " " + line),
+    comment,
+  ];
+  return lines.join("\n");
+}
+
+async function getTreePathAtCursor(
+  filepath: string,
+  fileContents: string,
+  cursorIndex: number
+): Promise<Parser.SyntaxNode[]> {
+  const parser = await getParserForFile(filepath);
+  const ast = parser.parse(fileContents);
+  const path = [ast.rootNode];
+  while (path[path.length - 1].childCount > 0) {
+    let foundChild = false;
+    for (let child of path[path.length - 1].children) {
+      if (child.startIndex <= cursorIndex && child.endIndex >= cursorIndex) {
+        path.push(child);
+        foundChild = true;
+        break;
+      }
+    }
+
+    if (!foundChild) {
+      break;
+    }
+  }
+
+  return path;
+}
+
+export interface AutocompleteSnippet {
+  filepath: string;
+  content: string;
+}
+
+export async function constructAutocompletePrompt(
+  filepath: string,
+  fullPrefix: string,
+  fullSuffix: string,
+  clipboardText: string,
+  language: AutocompleteLanguageInfo,
+  getDefinition: (
+    filepath: string,
+    line: number,
+    character: number
+  ) => Promise<AutocompleteSnippet | undefined>
+): Promise<{ prefix: string; suffix: string; useFim: boolean }> {
+  // Find external snippets
+  const snippets: AutocompleteSnippet[] = [];
+
+  const treePath = await getTreePathAtCursor(
+    filepath,
+    fullPrefix + fullSuffix,
+    fullPrefix.length
+  );
+
+  // Get function def when inside call expression
+  let callExpression = undefined;
+  for (let node of treePath.reverse()) {
+    if (node.type === "call_expression") {
+      callExpression = node;
+      break;
+    }
+  }
+  if (callExpression) {
+    const definition = await getDefinition(
+      filepath,
+      callExpression.startPosition.row,
+      callExpression.startPosition.column
+    );
+    if (definition) {
+      snippets.push(definition);
+    }
+  }
+
+  // Construct basic prefix / suffix
+  const formattedSnippets = snippets
+    .map((snippet) =>
+      formatExternalSnippet(snippet.filepath, snippet.content, language)
+    )
+    .join("\n");
+  const maxPrefixTokens =
+    MAX_PROMPT_TOKENS * PREFIX_PERCENTAGE -
+    countTokens(formattedSnippets, "gpt-4");
+  let prefix = pruneLinesFromTop(fullPrefix, maxPrefixTokens);
+  if (formattedSnippets.length > 0) {
+    prefix = formattedSnippets + "\n" + prefix;
+  }
+
+  const maxSuffixTokens = Math.min(
+    MAX_PROMPT_TOKENS - countTokens(prefix, "gpt-4"),
+    MAX_SUFFIX_PERCENTAGE * MAX_PROMPT_TOKENS
+  );
+  let suffix = pruneLinesFromBottom(fullSuffix, maxSuffixTokens);
+
+  return { prefix, suffix, useFim: true };
+}

core/autocomplete/languages.ts

@@ -0,0 +1,11 @@
+export interface AutocompleteLanguageInfo {
+  stopWords: string[];
+  comment: string;
+  endOfLine: string[];
+}
+
+export const Typescript = {
+  stopWords: ["function", "class", "module", "export "],
+  comment: "//",
+  endOfLine: [";"],
+};

core/autocomplete/parameters.ts

@@ -0,0 +1,4 @@
+export const MAX_PROMPT_TOKENS = 1_000;
+export const PREFIX_PERCENTAGE = 0.85;
+export const MAX_SUFFIX_PERCENTAGE = 0.25;
+export const DEBOUNCE_DELAY = 350;

core/autocomplete/templates.ts

@@ -0,0 +1,55 @@
+// Fill in the middle prompts
+
+import { CompletionOptions } from "..";
+
+interface AutocompleteTemplate {
+  template: string | ((prefix: string, suffix: string) => string);
+  completionOptions?: Partial<CompletionOptions>;
+}
+
+// https://huggingface.co/stabilityai/stable-code-3b
+const stableCodeFimTemplate: AutocompleteTemplate = {
+  template: "<fim_prefix>{{{prefix}}}<fim_suffix>{{{suffix}}}<fim_middle>",
+  completionOptions: {
+    stop: ["<fim_prefix>", "<fim_suffix>", "<fim_middle>", "<|endoftext|>"],
+  },
+};
+
+const codeLlamaFimTemplate: AutocompleteTemplate = {
+  template: "<PRE> {{{prefix}}} <SUF>{{{suffix}}} <MID>",
+  completionOptions: { stop: ["<PRE>", "<SUF>", "<MID>"] },
+};
+
+// https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base
+const deepseekFimTemplate: AutocompleteTemplate = {
+  template:
+    "<｜fim▁begin｜>{{{prefix}}}<｜fim▁hole｜>{{{suffix}}}<｜fim▁end｜>",
+  completionOptions: {
+    stop: ["<｜fim▁begin｜>", "<｜fim▁hole｜>", "<｜fim▁end｜>", "//"],
+  },
+};
+
+const deepseekFimTemplateWrongPipeChar: AutocompleteTemplate = {
+  template: "<|fim▁begin|>{{{prefix}}}<|fim▁hole|>{{{suffix}}}<|fim▁end|>",
+  completionOptions: { stop: ["<|fim▁begin|>", "<|fim▁hole|>", "<|fim▁end|>"] },
+};
+
+export function getTemplateForModel(model: string): AutocompleteTemplate {
+  if (
+    model.includes("starcoder") ||
+    model.includes("star-coder") ||
+    model.includes("stable")
+  ) {
+    return stableCodeFimTemplate;
+  }
+
+  if (model.includes("codellama")) {
+    return codeLlamaFimTemplate;
+  }
+
+  if (model.includes("deepseek")) {
+    return deepseekFimTemplate;
+  }
+
+  return stableCodeFimTemplate;
+}

core/config/load.ts

@@ -80,6 +80,20 @@ async function intermediateToFinalConfig(
     models.push(llm);
   }
 
+  let autocompleteLlm: BaseLLM | undefined = undefined;
+  if (config.tabAutocompleteModel) {
+    if (isModelDescription(config.tabAutocompleteModel)) {
+      autocompleteLlm = await llmFromDescription(
+        config.tabAutocompleteModel,
+        readFile,
+        config.completionOptions,
+        config.systemMessage
+      );
+    } else {
+      autocompleteLlm = new CustomLLMClass(config.tabAutocompleteModel);
+    }
+  }
+
   const contextProviders: IContextProvider[] = [new FileContextProvider({})];
   for (const provider of config.contextProviders || []) {
     if (isContextProviderWithParams(provider)) {
@@ -112,6 +126,7 @@ async function intermediateToFinalConfig(
     contextProviders,
     models,
     embeddingsProvider: config.embeddingsProvider as any,
+    tabAutocompleteModel: autocompleteLlm,
   };
 }
 

core/index.d.ts

@@ -452,7 +452,6 @@ export type ModelName =
   | "wizardcoder-34b"
   | "zephyr-7b"
   | "codeup-13b"
-  | "deepseek-1b"
   | "deepseek-7b"
   | "deepseek-33b"
   | "neural-chat-7b"
@@ -465,7 +464,12 @@ export type ModelName =
   // Mistral
   | "mistral-tiny"
   | "mistral-small"
-  | "mistral-medium";
+  | "mistral-medium"
+  // Tab autocomplete
+  | "deepseek-1b"
+  | "starcoder-1b"
+  | "starcoder-3b"
+  | "stable-code-3b";
 
 export interface RequestOptions {
   timeout?: number;
@@ -540,6 +544,11 @@ export interface EmbeddingsProvider {
   embed(chunks: string[]): Promise<number[][]>;
 }
 
+export interface TabAutocompleteOptions {
+  useCopyBuffer?: boolean;
+  useSuffix?: boolean;
+}
+
 export interface SerializedContinueConfig {
   disallowedSteps?: string[];
   allowAnonymousTelemetry?: boolean;
@@ -553,6 +562,8 @@ export interface SerializedContinueConfig {
   disableSessionTitles?: boolean;
   userToken?: string;
   embeddingsProvider?: EmbeddingsProviderDescription;
+  tabAutocompleteModel?: ModelDescription;
+  tabAutocompleteOptions?: TabAutocompleteOptions;
 }
 
 export interface Config {
@@ -581,6 +592,8 @@ export interface Config {
   userToken?: string;
   /** The provider used to calculate embeddings. If left empty, Continue will use transformers.js to calculate the embeddings with all-MiniLM-L6-v2 */
   embeddingsProvider?: EmbeddingsProviderDescription | EmbeddingsProvider;
+  /** The model that Continue will use for tab autocompletions. */
+  tabAutocompleteModel?: CustomLLM | ModelDescription;
 }
 
 export interface ContinueConfig {
@@ -594,4 +607,5 @@ export interface ContinueConfig {
   disableIndexing?: boolean;
   userToken?: string;
   embeddingsProvider: EmbeddingsProvider;
+  tabAutocompleteModel?: ILLM;
 }

core/indexing/chunk/chunk.ts

@@ -1,8 +1,9 @@
 import { Chunk, ChunkWithoutID } from "../..";
 import { MAX_CHUNK_SIZE } from "../../llm/constants";
 import { countTokens } from "../../llm/countTokens";
+import { supportedLanguages } from "../../util/treeSitter";
 import { basicChunker } from "./basic";
-import { codeChunker, supportedLanguages } from "./code";
+import { codeChunker } from "./code";
 
 async function* chunkDocumentWithoutId(
   filepath: string,