feat(vscode): support multilingual notebook autocomplete

.changeset/bright-notebooks-complete.md

@@ -2,4 +2,4 @@
 "kilo-code": minor
 ---
 
-Enable autocomplete in Jupyter notebooks.
+Enable autocomplete across supported languages in Jupyter notebooks.

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.test.ts

@@ -1,6 +1,7 @@
 import { describe, it, expect } from "vitest"
 import {
   languageForFilepath,
+  languageForId,
   LANGUAGES,
   Typescript,
   JavaScript,
@@ -31,6 +32,24 @@ import {
 } from "./AutocompleteLanguageInfo"
 
 describe("AutocompleteLanguageInfo", () => {
+  describe("languageForId", () => {
+    it("resolves VS Code language identifiers", () => {
+      expect(languageForId("typescript")).toBe(Typescript)
+      expect(languageForId("typescriptreact")).toBe(Typescript)
+      expect(languageForId("javascript")).toBe(JavaScript)
+      expect(languageForId("javascriptreact")).toBe(JavaScript)
+      expect(languageForId("jsonc")).toBe(Json)
+      expect(languageForId("python")).toBe(Python)
+      expect(languageForId("r")).toBe(R)
+      expect(languageForId("julia")).toBe(Julia)
+      expect(languageForId("luau")).toBe(Lua)
+    })
+
+    it("rejects unknown language identifiers", () => {
+      expect(languageForId("custom-language")).toBeUndefined()
+    })
+  })
+
   describe("languageForFilepath", () => {
     describe("TypeScript/JavaScript files", () => {
       it("should return TypeScript for .ts files", () => {

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.ts

@@ -367,6 +367,42 @@ export const LANGUAGES: { [extension: string]: AutocompleteLanguageInfo } = {
   luau: Lua,
 }
 
+const IDS: Record<string, AutocompleteLanguageInfo> = {
+  typescript: Typescript,
+  typescriptreact: Typescript,
+  javascript: JavaScript,
+  javascriptreact: JavaScript,
+  json: Json,
+  jsonc: Json,
+  python: Python,
+  java: Java,
+  cpp: Cpp,
+  c: C,
+  csharp: CSharp,
+  scala: Scala,
+  go: Go,
+  rust: Rust,
+  haskell: Haskell,
+  php: PHP,
+  ruby: Ruby,
+  swift: Swift,
+  kotlin: Kotlin,
+  clojure: Clojure,
+  julia: Julia,
+  fsharp: FSharp,
+  r: R,
+  dart: Dart,
+  solidity: Solidity,
+  yaml: YAML,
+  markdown: Markdown,
+  lua: Lua,
+  luau: Lua,
+}
+
+export function languageForId(id: string): AutocompleteLanguageInfo | undefined {
+  return IDS[id]
+}
+
 export function languageForFilepath(fileUri: string): AutocompleteLanguageInfo {
   const extension = getUriFileExtension(fileUri)
   return LANGUAGES[extension] || Typescript

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/notebook.ts

@@ -2,9 +2,10 @@
 // https://github.com/continuedev/continue/blob/d0a3c0b626b5bebc3bef4742eec05a0242be0bab/extensions/vscode/src/autocomplete/completionProvider.ts#L226-L263
 // Copyright 2023 Continue
 // Licensed under the Apache License, Version 2.0.
-// Modified by Kilo Code for notebook paths, cursor positions, and cache scoping.
+// Modified by Kilo Code for notebook paths, cursor positions, multilingual context, and cache scoping.
 
 import * as vscode from "vscode"
+import { languageForId } from "./constants/AutocompleteLanguageInfo"
 
 export interface NotebookContext {
   contents: string
@@ -53,7 +54,7 @@ export function notebookUri(uri: vscode.Uri): vscode.Uri | undefined {
 export function supportsNotebook(document: vscode.TextDocument): boolean {
   if (document.uri.scheme !== "vscode-notebook-cell") return true
   const resolved = resolveNotebook(document.uri)
-  return resolved?.cell.kind === vscode.NotebookCellKind.Code && document.languageId === "python"
+  return resolved?.cell.kind === vscode.NotebookCellKind.Code && !!languageForId(document.languageId)
 }
 
 export function autocompleteScope(document: vscode.TextDocument): string {
@@ -64,7 +65,7 @@ export function autocompleteScope(document: vscode.TextDocument): string {
   const siblings = resolved.cells
     .filter((_, index) => index !== resolved.index)
     .map((cell) => [cell.document.uri.toString(), cell.kind, cell.document.languageId, cell.document.version])
-  return JSON.stringify([id, resolved.notebook.uri.toString(), resolved.index, siblings])
+  return JSON.stringify([id, document.languageId, resolved.notebook.uri.toString(), resolved.index, siblings])
 }
 
 export function getNotebookContext(
@@ -77,13 +78,24 @@ export function getNotebookContext(
   if (!resolved) return
 
   const cells = resolved.cells
+  const lang = languageForId(document.languageId)
+  if (!lang) return
+
+  const marker = document.languageId === "json" ? undefined : lang.singleLineComment
+  const comment = (text: string, label: string) =>
+    text
+      .split("\n")
+      .map((line, index) => (marker ? `${marker} ${index === 0 ? `[${label}] ` : ""}${line}` : ""))
+      .join("\n")
+
   const contents = cells
     .map((cell) => {
       const text = cell.document.getText()
-      if (cell.kind === vscode.NotebookCellKind.Markup) {
-        return `"""${text}"""`
-      }
-      return text
+      if (cell.kind === vscode.NotebookCellKind.Markup) return comment(text, "markdown")
+      const sibling = languageForId(cell.document.languageId)
+      const strict = document.languageId === "json" || cell.document.languageId === "json"
+      if (sibling === lang && (!strict || cell.document.languageId === document.languageId)) return text
+      return comment(text, cell.document.languageId)
     })
     .join("\n\n")
 

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/templating/constructPrefixSuffix.ts

@@ -1,6 +1,6 @@
 import { IDE } from "../.."
 import { getRangeInString } from "../../util/ranges"
-import { languageForFilepath } from "../constants/AutocompleteLanguageInfo"
+import { languageForFilepath, languageForId } from "../constants/AutocompleteLanguageInfo"
 import { AutocompleteInput } from "../util/types"
 
 /**
@@ -14,7 +14,7 @@ export async function constructInitialPrefixSuffix(
   prefix: string
   suffix: string
 }> {
-  const lang = languageForFilepath(input.filepath)
+  const lang = (input.languageId && languageForId(input.languageId)) || languageForFilepath(input.filepath)
 
   const fileContents = input.manuallyPassFileContents ?? (await ide.readFile(input.filepath))
   const fileLines = fileContents.split("\n")

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/HelperVars.ts

@@ -1,6 +1,6 @@
 import { IDE, TabAutocompleteOptions } from "../.."
 import { countTokens, pruneLinesFromBottom, pruneLinesFromTop } from "../../llm/countTokens"
-import { AutocompleteLanguageInfo, languageForFilepath } from "../constants/AutocompleteLanguageInfo"
+import { AutocompleteLanguageInfo, languageForFilepath, languageForId } from "../constants/AutocompleteLanguageInfo"
 import { constructInitialPrefixSuffix } from "../templating/constructPrefixSuffix"
 
 import { AstPath, getAst, getTreePathAtCursor } from "./ast"
@@ -35,7 +35,7 @@ export const HelperVars = {
     modelName: string,
     ide: IDE,
   ): Promise<HelperVars> => {
-    const lang = languageForFilepath(input.filepath)
+    const lang = (input.languageId && languageForId(input.languageId)) || languageForFilepath(input.filepath)
     const workspaceUris = await ide.getWorkspaceDirs()
     const fileContents = input.manuallyPassFileContents ?? (await ide.readFile(input.filepath))
     const fileLines = fileContents.split("\n")

packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/types.ts

@@ -11,6 +11,7 @@ export interface AutocompleteInput {
   isUntitledFile: boolean
   completionId: string
   filepath: string
+  languageId?: string
   pos: Position
   recentlyVisitedRanges: AutocompleteCodeSnippet[]
   recentlyEditedRanges: RecentlyEditedRange[]

packages/kilo-vscode/src/services/autocomplete/types.ts

@@ -32,6 +32,7 @@ export interface AutocompleteInput {
   isUntitledFile: boolean
   completionId: string
   filepath: string
+  languageId?: string
   pos: Position
   recentlyVisitedRanges: AutocompleteCodeSnippet[]
   recentlyEditedRanges: RecentlyEditedRange[]
@@ -201,6 +202,7 @@ export function contextToAutocompleteInput(context: AutocompleteSuggestionContex
     isUntitledFile: context.document.isUntitled,
     completionId: crypto.randomUUID(),
     filepath: context.document.uri.fsPath,
+    languageId: context.document.languageId,
     pos: { line: position.line, character: position.character },
     recentlyVisitedRanges,
     recentlyEditedRanges,

packages/kilo-vscode/tests/unit/notebook-context.test.ts

@@ -7,7 +7,9 @@ import {
   supportsNotebook,
 } from "../../src/services/autocomplete/continuedev/core/autocomplete/notebook"
 import { accessible } from "../../src/services/autocomplete/classic-auto-complete/AutocompleteInlineCompletionProvider"
+import { constructInitialPrefixSuffix } from "../../src/services/autocomplete/continuedev/core/autocomplete/templating/constructPrefixSuffix"
 import type { FileIgnoreController } from "../../src/services/autocomplete/shims/FileIgnoreController"
+import type { AutocompleteInput } from "../../src/services/autocomplete/types"
 
 function uri(scheme: string, path: string, fragment = ""): vscode.Uri {
   const value = `${scheme}:${path}${fragment ? `#${fragment}` : ""}`
@@ -38,8 +40,8 @@ function notebooks(value: vscode.NotebookDocument[]): void {
 describe("notebook context", () => {
   beforeEach(() => notebooks([]))
 
-  it("flattens notebook cells and translates the cursor", () => {
-    const markdown = document("markdown", "# Title\nNotes")
+  it("projects mixed-language context for the active Python cell", () => {
+    const markdown = document("markdown", "# Title\nNotes", "markdown")
     const code = document("code", "const value = 1\nvalue += 1", "javascript")
     const current = document("current", "print(value)\nprint('done')")
     const notebook = {
@@ -55,34 +57,102 @@ describe("notebook context", () => {
     const context = getNotebookContext(current, new vscode.Position(1, 5))
 
     expect(context).toEqual({
-      contents: `"""# Title\nNotes"""\n\nconst value = 1\nvalue += 1\n\nprint(value)\nprint('done')`,
+      contents: `# [markdown] # Title\n# Notes\n\n# [javascript] const value = 1\n# value += 1\n\nprint(value)\nprint('done')`,
       filepath: "/workspace/example.ipynb",
       position: new vscode.Position(7, 5),
     })
   })
 
-  it("limits notebook completion to Python code cells", () => {
-    const python = document("python", "value = 1")
-    const javascript = document("javascript", "const value = 1", "javascript")
-    const markdown = document("markdown", "# Heading", "markdown")
+  it("projects mixed-language context for the active JavaScript cell", () => {
+    const markdown = document("markdown", "Setup\nvalues", "markdown")
+    const python = document("python", "value = 1\nprint(value)")
+    const current = document("current", "const value = 1", "javascript")
     const notebook = {
       uri: uri("file", "/workspace/example.ipynb"),
       getCells: () => [
-        { kind: vscode.NotebookCellKind.Code, document: python },
-        { kind: vscode.NotebookCellKind.Code, document: javascript },
         { kind: vscode.NotebookCellKind.Markup, document: markdown },
+        { kind: vscode.NotebookCellKind.Code, document: python },
+        { kind: vscode.NotebookCellKind.Code, document: current },
       ],
     } as vscode.NotebookDocument
     notebooks([notebook])
 
-    expect(supportsNotebook(python)).toBe(true)
-    expect(supportsNotebook(javascript)).toBe(false)
-    expect(supportsNotebook(markdown)).toBe(false)
-    expect(getNotebookContext(javascript, new vscode.Position(0, 0))).toBeUndefined()
-    expect(getNotebookContext(markdown, new vscode.Position(0, 0))).toBeUndefined()
+    expect(getNotebookContext(current, new vscode.Position(0, 6))).toEqual({
+      contents: `// [markdown] Setup\n// values\n\n// [python] value = 1\n// print(value)\n\nconst value = 1`,
+      filepath: "/workspace/example.ipynb",
+      position: new vscode.Position(6, 6),
+    })
+  })
+
+  it("supports known code languages and rejects non-code or unknown cells", () => {
+    const cells = [
+      document("python", "value = 1"),
+      document("javascript", "const value = 1", "javascript"),
+      document("typescript", "const value: number = 1", "typescript"),
+      document("r", "value <- 1", "r"),
+      document("julia", "value = 1", "julia"),
+      document("jsonc", "{ // comment\n}", "jsonc"),
+      document("luau", "local value = 1", "luau"),
+      document("unknown", "value = 1", "custom-language"),
+      document("markdown", "# Heading", "markdown"),
+    ]
+    const notebook = {
+      uri: uri("file", "/workspace/example.ipynb"),
+      getCells: () =>
+        cells.map((document, index) => ({
+          kind: index === cells.length - 1 ? vscode.NotebookCellKind.Markup : vscode.NotebookCellKind.Code,
+          document,
+        })),
+    } as vscode.NotebookDocument
+    notebooks([notebook])
+
+    expect(cells.slice(0, 7).every(supportsNotebook)).toBe(true)
+    expect(supportsNotebook(cells[7]!)).toBe(false)
+    expect(supportsNotebook(cells[8]!)).toBe(false)
+    expect(getNotebookContext(cells[7]!, new vscode.Position(0, 0))).toBeUndefined()
     expect(supportsNotebook({ uri: uri("file", "/workspace/file.ts") } as vscode.TextDocument)).toBe(true)
   })
 
+  it("omits foreign and markup content from strict JSON context", () => {
+    const markdown = document("markdown", "Describe values", "markdown")
+    const javascript = document("javascript", "const value = 1", "javascript")
+    const current = document("current", '{"value": 1}', "json")
+    const notebook = {
+      uri: uri("file", "/workspace/example.ipynb"),
+      getCells: () => [
+        { kind: vscode.NotebookCellKind.Markup, document: markdown },
+        { kind: vscode.NotebookCellKind.Code, document: javascript },
+        { kind: vscode.NotebookCellKind.Code, document: current },
+      ],
+    } as vscode.NotebookDocument
+    notebooks([notebook])
+
+    expect(getNotebookContext(current, new vscode.Position(0, 3))).toEqual({
+      contents: `\n\n\n\n{"value": 1}`,
+      filepath: "/workspace/example.ipynb",
+      position: new vscode.Position(4, 3),
+    })
+  })
+
+  it("uses the active cell language when constructing notebook prompts", async () => {
+    const input: AutocompleteInput = {
+      isUntitledFile: false,
+      completionId: "completion",
+      filepath: "/workspace/example.ipynb",
+      languageId: "javascript",
+      pos: { line: 0, character: 5 },
+      recentlyVisitedRanges: [],
+      recentlyEditedRanges: [],
+      manuallyPassFileContents: "value = 1",
+      injectDetails: "notebook context",
+    }
+
+    const result = await constructInitialPrefixSuffix(input, {} as never)
+
+    expect(result.prefix).toBe("\n// notebook context\nvalue")
+    expect(result.suffix).toBe(" = 1")
+  })
+
   it("resolves file and notebook cell URIs", () => {
     const file = uri("file", "/workspace/file.ts")
     const cell = document("code", "value = 1")
@@ -120,6 +190,11 @@ describe("notebook context", () => {
     Object.assign(notebook, { version: 3 })
     expect(autocompleteScope(current)).not.toBe(initial)
     expect(autocompleteScope(current)).not.toBe(autocompleteScope(sibling))
+
+    const changed = autocompleteScope(current)
+    Object.assign(current, { languageId: "javascript" })
+    Object.assign(notebook, { version: 4 })
+    expect(autocompleteScope(current)).not.toBe(changed)
   })
 
   it("changes autocomplete scope when sibling order changes", () => {