diff --git a/.changeset/multilingual-notebook-autocomplete.md b/.changeset/multilingual-notebook-autocomplete.md new file mode 100644 index 00000000000..1da4072fb4f --- /dev/null +++ b/.changeset/multilingual-notebook-autocomplete.md @@ -0,0 +1,5 @@ +--- +"kilo-code": minor +--- + +Enable autocomplete across supported languages in Jupyter notebooks. diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.test.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.test.ts index a6c35399519..e22b07b40d0 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.test.ts +++ b/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", () => { diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.ts index aee32df9230..012e37828aa 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/constants/AutocompleteLanguageInfo.ts +++ b/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 = { + 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 diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/notebook.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/notebook.ts index 2dca033b18e..15d3180d277 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/notebook.ts +++ b/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 json = document.languageId === "json" || document.languageId === "jsonc" + const marker = 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) => { + .map((cell, index) => { const text = cell.document.getText() - if (cell.kind === vscode.NotebookCellKind.Markup) { - return `"""${text}"""` - } - return text + if (index === resolved.index) return text + if (cell.kind === vscode.NotebookCellKind.Markup) return comment(text, "markdown") + if (!json && languageForId(cell.document.languageId) === lang) return text + return comment(text, cell.document.languageId) }) .join("\n\n") diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/templating/constructPrefixSuffix.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/templating/constructPrefixSuffix.ts index 0e72eafcf79..5f6aa302da2 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/templating/constructPrefixSuffix.ts +++ b/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") diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/HelperVars.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/HelperVars.ts index 000fc7ee905..ed2a8613e16 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/HelperVars.ts +++ b/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 => { - 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") diff --git a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/types.ts b/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/types.ts index 1b11559f8a0..a366bf48668 100644 --- a/packages/kilo-vscode/src/services/autocomplete/continuedev/core/autocomplete/util/types.ts +++ b/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[] diff --git a/packages/kilo-vscode/src/services/autocomplete/types.ts b/packages/kilo-vscode/src/services/autocomplete/types.ts index b7a88433828..cd400d47865 100644 --- a/packages/kilo-vscode/src/services/autocomplete/types.ts +++ b/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, diff --git a/packages/kilo-vscode/tests/unit/notebook-context.test.ts b/packages/kilo-vscode/tests/unit/notebook-context.test.ts index c4f2742936c..99002dd05a4 100644 --- a/packages/kilo-vscode/tests/unit/notebook-context.test.ts +++ b/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,104 @@ 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 sibling = document("sibling", '{"other": 2}', "json") + 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: sibling }, + { 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\n\n{"value": 1}`, + filepath: "/workspace/example.ipynb", + position: new vscode.Position(6, 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 +192,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", () => {