diff --git a/.gitignore b/.gitignore index 595435d9d637..1162a696a744 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,6 @@ __pycache__ # ignore cargo flamegraph temp files flamegraph.svg cargo-flamegraph.stacks +# for temp NES docs +/docs + diff --git a/clients/tabby-agent/src/codeCompletion/contexts.ts b/clients/tabby-agent/src/codeCompletion/contexts.ts index 6db935a50ed8..10d3752caad2 100644 --- a/clients/tabby-agent/src/codeCompletion/contexts.ts +++ b/clients/tabby-agent/src/codeCompletion/contexts.ts @@ -3,6 +3,7 @@ import type { ConfigData } from "../config/type"; import path from "path"; import hashObject from "object-hash"; import { splitLines, isBlank, regOnlyAutoClosingCloseChars } from "../utils/string"; +import { EditHistory } from "./editHistory"; export type CompletionRequest = { filepath: string; @@ -30,6 +31,8 @@ export type CompletionRequest = { insertSeg?: string; currSeg?: string; }; + mode?: "standard" | "next_edit_suggestion"; + editHistory?: EditHistory; }; export type Declaration = { @@ -309,4 +312,22 @@ export class CompletionContext { clipboard, }; } + + /** + * Build segments object specifically for next edit suggestion mode + * Includes minimal segments with edit_history + * @param editHistory Converted edit history object in snake_case format + * @returns Segments object with edit_history + */ + static buildSegmentsForNextEditSuggestion( + filepath: string | undefined, + editHistory: any, + ): TabbyApiComponents["schemas"]["Segments"] { + return { + prefix: "", + suffix: "", + filepath, + edit_history: editHistory, + }; + } } diff --git a/clients/tabby-agent/src/codeCompletion/editHistory.ts b/clients/tabby-agent/src/codeCompletion/editHistory.ts new file mode 100644 index 000000000000..4ecc61a113d7 --- /dev/null +++ b/clients/tabby-agent/src/codeCompletion/editHistory.ts @@ -0,0 +1,605 @@ +import { TextDocument } from "vscode-languageserver-textdocument"; +import { Position, Range, FoldingRange, Connection } from "vscode-languageserver"; +import { TextDocuments } from "../lsp/textDocuments"; +import type { Configurations } from "../config"; +import { getLogger } from "../logger"; +import * as diff from "diff"; +import type { components as TabbyApiComponents } from "tabby-openapi/compatible"; +/** + * Interface for the cursor position in a document + */ +export interface CursorPosition { + line: number; + character: number; +} + +/** + * Interface for the current version of a document + */ +export interface CurrentVersion { + content: string; + cursorPosition: CursorPosition; +} + +/** + * Interface for the edit history of a document + */ +export interface EditHistory { + originalCode: string; + editsDiff: string; + currentVersion: CurrentVersion; +} + +/** + * Configuration options for window strategy + */ +interface WindowStrategyConfig { + useFoldingRanges: boolean; + defaultLinesAround: number; + minLines: number; + maxLines: number; + requireCompleteBrackets: boolean; +} + +/** + * Uniquely identifies a code block within a file + */ +interface BlockIdentifier { + startLine: number; + endLine: number; +} + +/** + * Represents a tracked code block with its edit history + */ +interface TrackedBlock { + originalContent: string; + edits: string[]; + lastUpdated: number; +} + +/** + * Class responsible for tracking edit history for files + */ +export class EditHistoryTracker { + private readonly logger = getLogger("EditHistoryTracker"); + private fileContentCache: Map = new Map(); + // Map structure: uri -> blockKey -> TrackedBlock + private trackedBlocksCache: Map> = new Map(); + private maxEditsPerBlock: number = 10; + private maxDiffContextSize: number = 5000; + private windowConfig: WindowStrategyConfig; + private connection: Connection | null = null; + private initPromises: Map> = new Map(); + + constructor( + private readonly documents: TextDocuments, + private config: Configurations, + connection?: Connection, + ) { + this.connection = connection || null; + + this.windowConfig = { + useFoldingRanges: true, + defaultLinesAround: 20, + minLines: 5, + maxLines: 100, + requireCompleteBrackets: true, + }; + + for (const document of documents.all()) { + this.logger.info(`Initializing edit history for ${document.uri}`); + this.trackedBlocksCache.set(document.uri, new Map()); + const initPromise = Promise.resolve(); + this.initPromises.set(document.uri, initPromise); + } + + documents.onDidChangeContent((change) => { + this.trackDocumentChange(change.document).catch((err) => { + this.logger.error(`Error tracking document change: ${err}`, err); + }); + }); + + this.updateConfig(config); + } + + /** + * Get a string key for a block identifier + */ + private getBlockKey(block: BlockIdentifier): string { + return `${block.startLine}-${block.endLine}`; + } + + /** + * Track document changes + */ + private async trackDocumentChange(document: TextDocument): Promise { + const uri = document.uri; + const newContent = document.getText(); + const oldContent = this.fileContentCache.get(uri); + + if (this.initPromises.has(uri)) { + try { + await this.initPromises.get(uri); + this.initPromises.delete(uri); + } catch (error) { + this.logger.error(`Error waiting for document initialization: ${error}`, error); + } + } + + if (oldContent === undefined) { + this.logger.info(`First time seeing file ${uri}, initializing cache`); + this.fileContentCache.set(uri, newContent); + if (!this.trackedBlocksCache.has(uri)) { + this.trackedBlocksCache.set(uri, new Map()); + } + return; + } + + if (oldContent === newContent) { + return; + } + + // Update file content cache + this.fileContentCache.set(uri, newContent); + } + + /** + * Get the edit history for a document at a specific position + */ + public async getEditHistory(uri: string, position: Position): Promise { + if (this.initPromises.has(uri)) { + try { + await this.initPromises.get(uri); + this.initPromises.delete(uri); + } catch (error) { + this.logger.error(`Error waiting for document initialization: ${error}`, error); + } + } + + const document = this.documents.get(uri); + if (!document) { + this.logger.info(`No document available for ${uri}`); + return undefined; + } + + try { + // Get the current context window for the position + const contextWindow = await this.getSmartContextWindow(document, position); + + // Identify the block + const blockId = await this.identifyCodeBlock(document, position); + const blockKey = this.getBlockKey(blockId); + + // Get or create tracked block + let trackedBlockMap = this.trackedBlocksCache.get(uri); + if (!trackedBlockMap) { + trackedBlockMap = new Map(); + this.trackedBlocksCache.set(uri, trackedBlockMap); + } + + let trackedBlock = trackedBlockMap.get(blockKey); + if (!trackedBlock) { + // First time seeing this block, initialize it + trackedBlock = { + originalContent: contextWindow, + edits: [], + lastUpdated: Date.now(), + }; + trackedBlockMap.set(blockKey, trackedBlock); + this.logger.info(`Initialized new tracked block at ${blockKey} for ${uri}`); + } else { + // Block exists, calculate diff + const unifiedDiff = this.generateUnifiedDiff(uri, trackedBlock.originalContent, contextWindow); + + if (trackedBlock.edits.length === 0 || trackedBlock.edits[trackedBlock.edits.length - 1] !== unifiedDiff) { + trackedBlock.edits.push(unifiedDiff); + trackedBlock.lastUpdated = Date.now(); + + // Cap the edits array size + if (trackedBlock.edits.length > this.maxEditsPerBlock) { + trackedBlock.edits.shift(); + } + + this.logger.info( + `Updated tracked block at ${blockKey} for ${uri}, edits queue size: ${trackedBlock.edits.length}`, + ); + } + } + + // Get the most recent diff + const latestDiff = trackedBlock.edits.length > 0 ? trackedBlock.edits[trackedBlock.edits.length - 1] : ""; + + const result: EditHistory = { + originalCode: trackedBlock.originalContent, + editsDiff: latestDiff || "", + currentVersion: { + content: contextWindow, + cursorPosition: { + line: position.line, + character: position.character, + }, + }, + }; + + this.logger.info( + `Generated edit history for block ${blockKey} in ${uri} with context window size ${this.countLines(contextWindow)} lines`, + ); + return result; + } catch (error) { + this.logger.error(`Failed to get edit history for ${uri}: ${error}`, error); + return undefined; + } + } + + /** + * Identify the code block containing the position + */ + private async identifyCodeBlock(document: TextDocument, position: Position): Promise { + if (this.windowConfig.useFoldingRanges && this.connection) { + try { + const foldingRanges = (await this.connection.sendRequest("textDocument/foldingRange", { + textDocument: { uri: document.uri }, + })) as FoldingRange[]; + + if (foldingRanges && foldingRanges.length > 0) { + const containingRange = this.findSmallestContainingRange(foldingRanges, position); + + if (containingRange) { + return { + startLine: containingRange.startLine, + endLine: containingRange.endLine, + }; + } + } + } catch (error) { + this.logger.error(`Error getting folding ranges: ${error}`, error); + } + } + + // Fallback to fixed window + return { + startLine: Math.max(0, position.line - this.windowConfig.defaultLinesAround), + endLine: position.line + this.windowConfig.defaultLinesAround, + }; + } + + /** + * Reset the original content to the current version after a NES request + * This ensures that subsequent edits are compared against the accepted state + */ + public async updateOriginalContentToCurrentVersion(uri: string, position: Position): Promise { + if (this.initPromises.has(uri)) { + try { + await this.initPromises.get(uri); + this.initPromises.delete(uri); + } catch (error) { + this.logger.error(`Error waiting for document initialization: ${error}`, error); + return; + } + } + + const document = this.documents.get(uri); + if (!document) { + this.logger.info(`No document available for ${uri}`); + return; + } + + try { + // Get the current context window + const contextWindow = await this.getSmartContextWindow(document, position); + + // Identify the block + const blockId = await this.identifyCodeBlock(document, position); + const blockKey = this.getBlockKey(blockId); + + // Get or create tracked block map + let trackedBlockMap = this.trackedBlocksCache.get(uri); + if (!trackedBlockMap) { + trackedBlockMap = new Map(); + this.trackedBlocksCache.set(uri, trackedBlockMap); + } + + // Update the original content for this block + const trackedBlock = { + originalContent: contextWindow, + edits: [], + lastUpdated: Date.now(), + }; + trackedBlockMap.set(blockKey, trackedBlock); + + this.logger.info(`Reset original content for block ${blockKey} in ${uri} to current version`); + } catch (error) { + this.logger.error(`Failed to update original content to current version for ${uri}: ${error}`, error); + } + } + + /** + * Generate a unified diff between old and new content + */ + private generateUnifiedDiff(uri: string, oldContent: string, newContent: string): string { + const filename = uri.split("/").pop() || uri; + + const patch = diff + .createPatch(filename, oldContent, newContent, undefined, undefined, { + context: 0, + }) + .split("\n") + .slice(2) + .join("\n"); + + if (patch.length > this.maxDiffContextSize) { + this.logger.warn(`Diff for ${uri} exceeds max size (${patch.length} > ${this.maxDiffContextSize}), truncating`); + return patch.substring(0, this.maxDiffContextSize); + } + + return patch; + } + + /** + * Get a smart context window around a position using folding ranges + */ + private async getSmartContextWindow(document: TextDocument, position: Position): Promise { + if (this.windowConfig.useFoldingRanges && this.connection) { + try { + this.logger.info(`Attempting to get folding ranges for ${document.uri}`); + + const foldingRanges = (await this.connection.sendRequest("textDocument/foldingRange", { + textDocument: { uri: document.uri }, + })) as FoldingRange[]; + + if (foldingRanges && foldingRanges.length > 0) { + const containingRange = this.findSmallestContainingRange(foldingRanges, position); + + if (containingRange) { + const lines = document.getText().split("\n"); + const endLine = Math.min(containingRange.endLine + 1, lines.length); + + const range = Range.create(containingRange.startLine, 0, endLine, 0); + + const contextWindow = this.getTextForRange(document, range); + const lineCount = this.countLines(contextWindow); + + this.logger.info(`Found containing folding range with ${lineCount} lines`); + + if (this.isContextSufficient(contextWindow)) { + return contextWindow; + } else { + this.logger.info(`Folding range context window not sufficient, falling back to fixed line window`); + } + } else { + this.logger.info(`No containing folding range found for position line ${position.line}`); + } + } else { + this.logger.info(`No folding ranges returned for ${document.uri}`); + } + } catch (error) { + this.logger.error(`Error getting folding ranges: ${error}`, error); + } + } + + return this.getFixedLineWindow(document, position); + } + + /** + * Find the smallest folding range that contains the position + */ + private findSmallestContainingRange(ranges: FoldingRange[], position: Position): FoldingRange | null { + let smallestRange: FoldingRange | null = null; + let smallestSize = Infinity; + + for (const range of ranges) { + if (range.startLine <= position.line && range.endLine >= position.line) { + const size = range.endLine - range.startLine + 1; + + if (size < smallestSize) { + smallestRange = range; + smallestSize = size; + } + } + } + + return smallestRange; + } + + /** + * Get text for a specific range in a document + */ + private getTextForRange(document: TextDocument, range: Range): string { + const content = document.getText(); + const lines = content.split("\n"); + + const startLine = Math.max(0, Math.min(range.start.line, lines.length - 1)); + const endLine = Math.max(0, Math.min(range.end.line, lines.length - 1)); + + const linesInRange = lines.slice(startLine, endLine + 1); + + return linesInRange.join("\n"); + } + + /** + * Get a fixed number of lines around a position + */ + private getFixedLineWindow(document: TextDocument, position: Position): string { + const content = document.getText(); + const lines = content.split("\n"); + + const startLine = Math.max(0, position.line - this.windowConfig.defaultLinesAround); + const endLine = Math.min(lines.length - 1, position.line + this.windowConfig.defaultLinesAround); + + const linesInRange = lines.slice(startLine, endLine + 1); + + this.logger.info( + `Generated fixed line window with ${linesInRange.length} lines around position line ${position.line}`, + ); + + return linesInRange.join("\n"); + } + + /** + * Check if a context window is sufficient based on configuration + */ + private isContextSufficient(context: string): boolean { + const lineCount = this.countLines(context); + if (lineCount < this.windowConfig.minLines) { + return false; + } + + if (lineCount > this.windowConfig.maxLines) { + return false; + } + + if (this.windowConfig.requireCompleteBrackets && !this.hasBalancedBrackets(context)) { + return false; + } + + return true; + } + + /** + * Count the number of lines in a string + */ + private countLines(text: string): number { + return (text.match(/\n/g) || []).length + 1; + } + + /** + * Check if a string has balanced brackets + * This is a simplified check that accepts more code fragments + */ + private hasBalancedBrackets(text: string): boolean { + const stack: string[] = []; + + for (let i = 0; i < text.length; i++) { + const char = text[i]; + if (char === "{") { + stack.push(char); + } else if (char === "}") { + if (stack.length === 0 || stack.pop() !== "{") { + return false; + } + } + } + + return true; + } + + /** + * Find a position where the text differs between old and new content + */ + private findChangePosition(oldContent: string, newContent: string): Position { + if (oldContent.length > 1000000 || newContent.length > 1000000) { + this.logger.info("Large file detected, using optimized change detection"); + + const lengthDiff = newContent.length - oldContent.length; + const estimatedChangePct = Math.abs(lengthDiff) / Math.max(oldContent.length, newContent.length); + + if (estimatedChangePct < 0.05) { + let line = 0; + let char = 0; + + const minLength = Math.min(oldContent.length, newContent.length); + for (let i = 0; i < minLength; i++) { + if (oldContent[i] !== newContent[i]) { + const textBefore = oldContent.substring(0, i); + const lines = textBefore.split("\n"); + line = lines.length - 1; + char = (lines[lines.length - 1] || "").length; + break; + } + } + + return Position.create(line, char); + } + } + + const oldLines = oldContent.split("\n"); + const newLines = newContent.split("\n"); + + for (let i = 0; i < Math.max(oldLines.length, newLines.length); i++) { + if (oldLines[i] !== newLines[i]) { + const oldLine = oldLines[i] || ""; + const newLine = newLines[i] || ""; + + for (let j = 0; j < Math.max(oldLine.length, newLine.length); j++) { + if (oldLine[j] !== newLine[j]) { + return Position.create(i, j); + } + } + + return Position.create(i, Math.min(oldLine.length, newLine.length)); + } + } + + return Position.create(0, 0); + } + + /** + * Update the tracker's configuration + */ + public updateConfig(config: Configurations): void { + const mergedConfig = config.getMergedConfig(); + + const nextEditConfig = mergedConfig.completion.nextEditSuggestion || { + enabled: true, + maxEditsPerFile: 10, + maxDiffContextSize: 5000, + }; + + this.maxEditsPerBlock = nextEditConfig.maxEditsPerFile || 10; + this.maxDiffContextSize = nextEditConfig.maxDiffContextSize || 5000; + + this.windowConfig = { + useFoldingRanges: true, + defaultLinesAround: 20, + minLines: 5, + maxLines: 100, + requireCompleteBrackets: true, + }; + + this.logger.info( + `Updated config: useFoldingRanges=${this.windowConfig.useFoldingRanges}, ` + + `defaultLinesAround=${this.windowConfig.defaultLinesAround}, ` + + `maxEditsPerBlock=${this.maxEditsPerBlock}, ` + + `maxDiffContextSize=${this.maxDiffContextSize}`, + ); + } + + /** + * Clear the edit history for a specific file or all files + */ + public clearHistory(uri?: string): void { + if (uri) { + this.trackedBlocksCache.delete(uri); + this.logger.info(`Cleared edit history for ${uri}`); + } else { + this.trackedBlocksCache.clear(); + this.logger.info("Cleared all edit history"); + } + } +} + +/** + * Build edit history for next edit suggestion mode + * Converts from camelCase to snake_case format for the API + */ +export function buildEditHistoryForRequest( + editHistory: EditHistory, +): TabbyApiComponents["schemas"]["Segments"]["edit_history"] { + const originalCode = editHistory.originalCode; + const editsDiff = editHistory.editsDiff; + const currentVersionContent = editHistory.currentVersion.content; + const cursorPosition = { + line: editHistory.currentVersion.cursorPosition.line, + character: editHistory.currentVersion.cursorPosition.character, + }; + + return { + original_code: originalCode, + edits_diff: editsDiff, + current_version: { + content: currentVersionContent, + cursor_position: cursorPosition, + }, + }; +} diff --git a/clients/tabby-agent/src/codeCompletion/index.ts b/clients/tabby-agent/src/codeCompletion/index.ts index b39e5f83bab9..a49a4227d9e2 100644 --- a/clients/tabby-agent/src/codeCompletion/index.ts +++ b/clients/tabby-agent/src/codeCompletion/index.ts @@ -10,8 +10,8 @@ import type { NotebookCell, CompletionParams, CompletionOptions, - InlineCompletionParams, TextDocumentPositionParams, + InlineCompletionParams, } from "vscode-languageserver"; import type { TextDocument } from "vscode-languageserver-textdocument"; import type { TextDocuments } from "../lsp/textDocuments"; @@ -33,6 +33,7 @@ import { CompletionList, CompletionItem as LspCompletionItem, InlineCompletionRequest, + InlineNESCompletionRequest, InlineCompletionList, InlineCompletionItem, TelemetryEventNotification, @@ -51,6 +52,7 @@ import { CompletionStats } from "./statistics"; import { CompletionContext, CompletionRequest } from "./contexts"; import { CompletionSolution, CompletionItem } from "./solution"; import { preCacheProcess, postCacheProcess } from "./postprocess"; +import { buildEditHistoryForRequest, EditHistoryTracker } from "./editHistory"; import { getLogger } from "../logger"; import { abortSignalFromAnyOf } from "../utils/signal"; import { splitLines, extractNonReservedWordList } from "../utils/string"; @@ -64,6 +66,7 @@ export class CompletionProvider implements Feature { private readonly completionCache = new CompletionCache(); private readonly completionDebounce = new CompletionDebounce(); private readonly completionStats = new CompletionStats(); + private readonly editHistoryTracker: EditHistoryTracker; private submitStatsTimer: ReturnType | undefined = undefined; @@ -73,8 +76,10 @@ export class CompletionProvider implements Feature { private completionFeatureOptions: CompletionOptions | undefined = undefined; private completionFeatureRegistration: Disposable | undefined = undefined; private inlineCompletionFeatureRegistration: Disposable | undefined = undefined; + private inlineNESCompletionFeatureRegistration: Disposable | undefined = undefined; private mutexAbortController: AbortController | undefined = undefined; + private nesMutexAbortController: AbortController | undefined = undefined; constructor( private readonly configurations: Configurations, @@ -85,7 +90,17 @@ export class CompletionProvider implements Feature { private readonly gitContextProvider: GitContextProvider, private readonly recentlyChangedCodeSearch: RecentlyChangedCodeSearch, private readonly fileTracker: FileTracker, - ) {} + ) { + // Initialize edit history tracker + getLogger("EditHistoryTracker").info("Initializing edit history tracker"); + getLogger("docs").info("docs :" + JSON.stringify(this.documents.all())); + this.editHistoryTracker = new EditHistoryTracker(documents, configurations); + + // Listen for configuration changes to update edit history settings + configurations.on("updated", () => { + this.editHistoryTracker.updateConfig(configurations); + }); + } initialize(connection: Connection, clientCapabilities: ClientCapabilities): ServerCapabilities { this.lspConnection = connection; @@ -113,6 +128,9 @@ export class CompletionProvider implements Feature { connection.onRequest(InlineCompletionRequest.type, async (params, token) => { return this.provideInlineCompletion(params, token); }); + connection.onRequest(InlineNESCompletionRequest.type, async (params, token) => { + return this.provideInlineNESCompletion(params, token); + }); if (!clientCapabilities.textDocument?.inlineCompletion.dynamicRegistration) { serverCapabilities = { ...serverCapabilities, @@ -156,11 +174,16 @@ export class CompletionProvider implements Feature { ) { this.inlineCompletionFeatureRegistration = await connection.client.register(InlineCompletionRequest.type); } + if (!this.inlineNESCompletionFeatureRegistration) { + this.inlineNESCompletionFeatureRegistration = await connection.client.register(InlineNESCompletionRequest.type); + } } else { this.completionFeatureRegistration?.dispose(); this.completionFeatureRegistration = undefined; this.inlineCompletionFeatureRegistration?.dispose(); this.inlineCompletionFeatureRegistration = undefined; + this.inlineNESCompletionFeatureRegistration?.dispose(); + this.inlineNESCompletionFeatureRegistration = undefined; } } @@ -211,9 +234,12 @@ export class CompletionProvider implements Feature { if (token.isCancellationRequested) { return null; } + const abortController = new AbortController(); token.onCancellationRequested(() => abortController.abort()); + try { + // Standard inline completion flow const request = await this.inlineCompletionParamsToCompletionRequest(params, token); if (!request) { return null; @@ -228,6 +254,171 @@ export class CompletionProvider implements Feature { } } + async provideInlineNESCompletion( + params: InlineCompletionParams, + token: CancellationToken, + ): Promise { + if (!this.tabbyApiClient.isCodeCompletionApiAvailable()) { + throw { + name: "CodeCompletionFeatureNotAvailableError", + message: "Code completion feature not available", + }; + } + if (token.isCancellationRequested) { + return null; + } + + const abortController = new AbortController(); + token.onCancellationRequested(() => abortController.abort()); + + try { + this.logger.info("Handling dedicated NES completion request"); + const config = this.configurations.getMergedConfig(); + if (!config.completion.nextEditSuggestion?.enabled) { + this.logger.info("Not a NES request enabled, rejecting"); + return null; + } + + const response = await this.provideNextEditSuggestion(params, token, abortController.signal); + return response; + } catch (error) { + this.logger.error("Error in NES completion", error); + return null; + } + } + + private async provideNextEditSuggestion( + params: InlineCompletionParams, + token: CancellationToken, + signal?: AbortSignal, + ): Promise { + const document = this.documents.get(params.textDocument.uri); + if (!document) { + this.logger.info("No document available for next edit suggestion"); + return null; + } + + this.logger.info(`Providing next edit suggestion for ${document.uri}`); + + // Get edit history from tracker + const editHistory = await this.editHistoryTracker.getEditHistory(document.uri, params.position); + + if (!editHistory) { + this.logger.info("No edit history available for next edit suggestion"); + return null; + } + + this.logger.info( + `Got edit history for ${document.uri}, original code length: ${editHistory.originalCode.length}, current version length: ${editHistory.currentVersion.content.length}`, + ); + + // Create base request from parameters + const request = await this.textDocumentPositionParamsToCompletionRequest(params, token); + if (!request) { + this.logger.info("Failed to create completion request for next edit suggestion"); + return null; + } + + // Add next edit suggestion specific fields + request.request.mode = "next_edit_suggestion"; + + // Convert editHistory from camelCase to snake_case for the API + // Similar to how segments is built in the provideCompletions method + const editHistoryForApi = buildEditHistoryForRequest(editHistory); + if (!editHistoryForApi) { + this.logger.info("Failed to build edit history for API"); + return null; + } + // Log the edit history being sent to the API + this.logger.info( + `Edit history for API - original_code length: ${editHistoryForApi.original_code.length}, current_version.content length: ${editHistoryForApi.current_version.content.length}`, + ); + + // Get the file path from document URI + const filepath = document.uri.split("/").pop(); + this.logger.info(`File path for next edit suggestion: ${filepath}`); + + // Create a request with segments containing edit_history + const modifiedRequest = { + ...request.request, + segments: CompletionContext.buildSegmentsForNextEditSuggestion(filepath, editHistoryForApi), + }; + + request.request = modifiedRequest as any; + + try { + const config = this.configurations.getMergedConfig(); + const temperature = config.completion.solution.temperature * 1.2; + + const requestOptions = { + ...request.request, + temperature, + }; + + getLogger("TabbyApiClient").info("Fetching next edit suggestion..." + JSON.stringify(editHistoryForApi)); + const response = await this.tabbyApiClient.fetchCompletion(requestOptions, signal, this.completionStats); + // TODO(Sma1lboy): add a timeout here, also have to update editHistoryTracker + getLogger("TabbyApiClient").info("Received next edit suggestion response: " + JSON.stringify(response)); + if (!response || !response.choices || response.choices.length === 0) { + return null; + } + + // Create inline completion items from response + const items: InlineCompletionItem[] = response.choices.map((choice) => { + // Use edit_range if provided, otherwise use normal range + if (choice.edit_range) { + const range = { + start: { + line: choice.edit_range.start_line, + character: choice.edit_range.start_character, + }, + end: { + line: choice.edit_range.end_line, + character: choice.edit_range.end_character, + }, + }; + + return { + insertText: choice.text, + range, + data: { + eventId: { + completionId: response.id, + choiceIndex: choice.index, + }, + }, + }; + } else { + // Fallback to standard inline completion item + return { + insertText: choice.text, + range: { + start: document.positionAt(params.position.character), + end: document.positionAt(params.position.character + choice.text.length), + }, + data: { + eventId: { + completionId: response.id, + choiceIndex: choice.index, + }, + }, + }; + } + }); + + this.logger.info(`Successfully generated next edit suggestion, updating original content to current version`); + await this.editHistoryTracker.updateOriginalContentToCurrentVersion(document.uri, params.position); + + return { + isIncomplete: false, + items, + }; + } catch (error) { + this.logger.error("Error providing next edit suggestion", error); + return null; + } + } + async postEvent(params: EventParams): Promise { this.completionStats.addEvent(params.type); const request = { @@ -377,7 +568,7 @@ export class CompletionProvider implements Feature { const config = this.configurations.getMergedConfig(); - // Mutex Control + // Mutex Control for normal completions if (this.mutexAbortController && !this.mutexAbortController.signal.aborted) { this.mutexAbortController.abort(new MutexAbortError()); } diff --git a/clients/tabby-agent/src/config/default.ts b/clients/tabby-agent/src/config/default.ts index a7b3683bfa56..344490392955 100644 --- a/clients/tabby-agent/src/config/default.ts +++ b/clients/tabby-agent/src/config/default.ts @@ -61,6 +61,11 @@ export const defaultConfigData: ConfigData = { maxTries: 6, temperature: 0.8, }, + nextEditSuggestion: { + enabled: true, + maxEditsPerFile: 10, + maxDiffContextSize: 5000, + }, }, postprocess: { limitScope: {}, diff --git a/clients/tabby-agent/src/config/type.d.ts b/clients/tabby-agent/src/config/type.d.ts index 3399abff6f77..eb67504c853a 100644 --- a/clients/tabby-agent/src/config/type.d.ts +++ b/clients/tabby-agent/src/config/type.d.ts @@ -68,6 +68,15 @@ export type ConfigData = { // The temperature for fetching the second and subsequent choices temperature: number; }; + // Next edit suggestion settings + nextEditSuggestion?: { + // Whether next edit suggestion is enabled + enabled: boolean; + // Maximum number of edits to track per file + maxEditsPerFile: number; + // Maximum size of diff content to track (in characters) + maxDiffContextSize: number; + }; }; postprocess: { limitScope: any; diff --git a/clients/tabby-agent/src/protocol.ts b/clients/tabby-agent/src/protocol.ts index c140e7626cbb..b66580f11533 100644 --- a/clients/tabby-agent/src/protocol.ts +++ b/clients/tabby-agent/src/protocol.ts @@ -31,7 +31,6 @@ import { CompletionList as LspCompletionList, CompletionItem as LspCompletionItem, InlineCompletionRequest as LspInlineCompletionRequest, - InlineCompletionParams, InlineCompletionList as LspInlineCompletionList, InlineCompletionItem as LspInlineCompletionItem, DeclarationParams, @@ -41,6 +40,7 @@ import { SemanticTokens, SemanticTokensLegend, WorkspaceEdit, + InlineCompletionParams, } from "vscode-languageserver-protocol"; /** @@ -355,6 +355,23 @@ export type InlineCompletionItem = LspInlineCompletionItem & { }; }; +/** + * [Tabby] Inline NES Completion Request(↩️) + * + * This method is sent from the client to the server to get next edit suggestion. + * - method: `tabby/inlineNESCompletion` + * - params: {@link InlineNESCompletionParams} + * - result: {@link InlineCompletionList} | null + * - error: {@link ChatFeatureNotAvailableError} + */ +export namespace InlineNESCompletionRequest { + export const method = "tabby/inlineNESCompletion"; + export const messageDirection = MessageDirection.clientToServer; + export const type = new ProtocolRequestType( + method, + ); +} + /** * [Tabby] Chat Edit Suggestion Command Request(↩️) * diff --git a/clients/tabby-openapi/lib/tabby.d.ts b/clients/tabby-openapi/lib/tabby.d.ts index 01e4cb46537e..55012cd6f045 100644 --- a/clients/tabby-openapi/lib/tabby.d.ts +++ b/clients/tabby-openapi/lib/tabby.d.ts @@ -95,34 +95,7 @@ export interface components { /** Format: int32 */ index: number; text: string; - }; - CodeSearchDocument: { - body: string; - filepath: string; - git_url: string; - language: string; - start_line: number; - }; - CodeSearchHit: { - scores: components["schemas"]["CodeSearchScores"]; - doc: components["schemas"]["CodeSearchDocument"]; - }; - CodeSearchQuery: { - git_url: string; - filepath?: string | null; - language?: string | null; - content: string; - }; - CodeSearchScores: { - /** - * Format: float - * @description Reciprocal rank fusion score: https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html - */ - rrf: number; - /** Format: float */ - bm25: number; - /** Format: float */ - embedding: number; + edit_range?: components["schemas"]["EditRange"] | null; }; /** @example { * "language": "python", @@ -153,6 +126,9 @@ export interface components { * @description The seed used for randomly selecting tokens */ seed?: number | null; + /** @description The mode for completion. Use 'standard' for normal code completions or 'next_edit_suggestion' + * to predict the next edit the user will make. */ + mode?: string; }; /** @example { * "choices": [ @@ -168,6 +144,25 @@ export interface components { choices: components["schemas"]["Choice"][]; debug_data?: components["schemas"]["DebugData"] | null; }; + /** @description Current version of the code after all edits */ + CurrentVersion: { + /** @description Current content after all edits */ + content: string; + cursor_position: components["schemas"]["CursorPosition"]; + }; + /** @description Cursor position in the current version */ + CursorPosition: { + /** + * Format: int32 + * @description Line number (0-based) + */ + line: number; + /** + * Format: int32 + * @description Character position within the line (0-based) + */ + character: number; + }; DebugData: { snippets?: components["schemas"]["Snippet"][] | null; prompt?: string | null; @@ -195,15 +190,36 @@ export interface components { /** @description Body of the snippet. */ body: string; }; - DocSearchDocument: { - title: string; - link: string; - snippet: string; + /** @description Contains information about edit history for next edit suggestion mode */ + EditHistory: { + /** @description Original code content before edits */ + original_code: string; + /** @description Unified git-style diff of all edits made to the file */ + edits_diff: string; + current_version: components["schemas"]["CurrentVersion"]; }; - DocSearchHit: { - /** Format: float */ - score: number; - doc: components["schemas"]["DocSearchDocument"]; + /** @description Range information for next edit suggestion mode */ + EditRange: { + /** + * Format: int32 + * @description Start line of the edit (0-based) + */ + start_line: number; + /** + * Format: int32 + * @description Start character position within the line (0-based) + */ + start_character: number; + /** + * Format: int32 + * @description End line of the edit (0-based) + */ + end_line: number; + /** + * Format: int32 + * @description End character position within the line (0-based) + */ + end_character: number; }; HealthState: { model?: string | null; @@ -257,10 +273,18 @@ export interface components { * * Sorted in descending order of [Snippet::score]. */ relevant_snippets_from_changed_files?: components["schemas"]["Snippet"][] | null; - /** @description The relevant code snippets extracted from recently opened files. These snippets are selected from candidates found within code chunks based on the last visited location. Current Active file is excluded from the search candidates. When provided with [Segments::relevant_snippets_from_changed_files], the snippets have already been deduplicated to ensure no duplication with entries in [Segments::relevant_snippets_from_changed_files]. */ + /** @description The relevant code snippets extracted from recently opened files. + * These snippets are selected from candidates found within code chunks + * based on the last visited location. + * + * Current Active file is excluded from the search candidates. + * When provided with [Segments::relevant_snippets_from_changed_files], the snippets have + * already been deduplicated to ensure no duplication with entries + * in [Segments::relevant_snippets_from_changed_files]. */ relevant_snippets_from_recently_opened_files?: components["schemas"]["Snippet"][] | null; /** @description Clipboard content when requesting code completion. */ clipboard?: string | null; + edit_history?: components["schemas"]["EditHistory"] | null; }; ServerSetting: { disable_client_side_telemetry: boolean; diff --git a/clients/tabby-openapi/openapi.json b/clients/tabby-openapi/openapi.json index 2382e69deb6e..5c321db6cc8d 100644 --- a/clients/tabby-openapi/openapi.json +++ b/clients/tabby-openapi/openapi.json @@ -1 +1 @@ -{"openapi":"3.0.3","info":{"title":"Tabby Server","description":"\n[![tabby stars](https://img.shields.io/github/stars/TabbyML/tabby)](https://github.com/TabbyML/tabby)\n[![Join Slack](https://shields.io/badge/Join-Tabby%20Slack-red?logo=slack)](https://links.tabbyml.com/join-slack)\n\nInstall following IDE / Editor extensions to get started with [Tabby](https://github.com/TabbyML/tabby).\n* [VSCode Extension](https://github.com/TabbyML/tabby/tree/main/clients/vscode) – Install from the [marketplace](https://marketplace.visualstudio.com/items?itemName=TabbyML.vscode-tabby), or [open-vsx.org](https://open-vsx.org/extension/TabbyML/vscode-tabby)\n* [VIM Extension](https://github.com/TabbyML/tabby/tree/main/clients/vim)\n* [IntelliJ Platform Plugin](https://github.com/TabbyML/tabby/tree/main/clients/intellij) – Install from the [marketplace](https://plugins.jetbrains.com/plugin/22379-tabby)\n","contact":{"name":"TabbyML Team"},"license":{"name":"Apache 2.0","url":"https://github.com/TabbyML/tabby/blob/main/LICENSE"},"version":"0.17.0-dev.0"},"servers":[{"url":"/","description":"Server"}],"paths":{"/v1/chat/completions":{"post":{"tags":["v1"],"operationId":"chat_completions","requestBody":{"description":"","content":{"application/json":{"schema":{}}},"required":true},"responses":{"200":{"description":"Success"},"405":{"description":"When chat model is not specified, the endpoint returns 405 Method Not Allowed"},"422":{"description":"When the prompt is malformed, the endpoint returns 422 Unprocessable Entity"}},"security":[{"token":[]}]}},"/v1/completions":{"post":{"tags":["v1"],"operationId":"completion","requestBody":{"content":{"application/json":{"schema":{"$ref":"#/components/schemas/CompletionRequest"}}},"required":true},"responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/CompletionResponse"}}}},"400":{"description":"Bad Request"}},"security":[{"token":[]}]}},"/v1/events":{"post":{"tags":["v1"],"operationId":"event","requestBody":{"content":{"application/json":{"schema":{"$ref":"#/components/schemas/LogEventRequest"}}},"required":true},"responses":{"200":{"description":"Success"},"400":{"description":"Bad Request"}},"security":[{"token":[]}]}},"/v1/health":{"get":{"tags":["v1"],"operationId":"health","responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/HealthState"}}}}},"security":[{"token":[]}]}},"/v1beta/server_setting":{"get":{"tags":["v1beta"],"operationId":"config","responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/ServerSetting"}}}}},"security":[{"token":[]}]}}},"components":{"schemas":{"Choice":{"type":"object","required":["index","text"],"properties":{"index":{"type":"integer","format":"int32","minimum":0},"text":{"type":"string"}}},"CodeSearchDocument":{"type":"object","required":["body","filepath","git_url","language","start_line"],"properties":{"body":{"type":"string"},"filepath":{"type":"string"},"git_url":{"type":"string"},"language":{"type":"string"},"start_line":{"type":"integer","minimum":0}}},"CodeSearchHit":{"type":"object","required":["scores","doc"],"properties":{"scores":{"$ref":"#/components/schemas/CodeSearchScores"},"doc":{"$ref":"#/components/schemas/CodeSearchDocument"}}},"CodeSearchQuery":{"type":"object","required":["git_url","content"],"properties":{"git_url":{"type":"string"},"filepath":{"type":"string","nullable":true},"language":{"type":"string","nullable":true},"content":{"type":"string"}}},"CodeSearchScores":{"type":"object","required":["rrf","bm25","embedding"],"properties":{"rrf":{"type":"number","format":"float","description":"Reciprocal rank fusion score: https://www.elastic.co/guide/en/elasticsearch/reference/current/rrf.html"},"bm25":{"type":"number","format":"float"},"embedding":{"type":"number","format":"float"}}},"CompletionRequest":{"type":"object","properties":{"language":{"type":"string","description":"Language identifier, full list is maintained at\nhttps://code.visualstudio.com/docs/languages/identifiers","example":"python","nullable":true},"segments":{"allOf":[{"$ref":"#/components/schemas/Segments"}],"nullable":true},"user":{"type":"string","description":"A unique identifier representing your end-user, which can help Tabby to monitor & generating\nreports.","nullable":true},"debug_options":{"allOf":[{"$ref":"#/components/schemas/DebugOptions"}],"nullable":true},"temperature":{"type":"number","format":"float","description":"The temperature parameter for the model, used to tune variance and \"creativity\" of the model output","nullable":true},"seed":{"type":"integer","format":"int64","description":"The seed used for randomly selecting tokens","nullable":true,"minimum":0}},"example":{"language":"python","segments":{"prefix":"def fib(n):\n ","suffix":"\n return fib(n - 1) + fib(n - 2)"}}},"CompletionResponse":{"type":"object","required":["id","choices"],"properties":{"id":{"type":"string"},"choices":{"type":"array","items":{"$ref":"#/components/schemas/Choice"}},"debug_data":{"allOf":[{"$ref":"#/components/schemas/DebugData"}],"nullable":true}},"example":{"choices":[{"index":0,"text":"string"}],"id":"string"}},"DebugData":{"type":"object","properties":{"snippets":{"type":"array","items":{"$ref":"#/components/schemas/Snippet"},"nullable":true},"prompt":{"type":"string","nullable":true}}},"DebugOptions":{"type":"object","properties":{"raw_prompt":{"type":"string","description":"When `raw_prompt` is specified, it will be passed directly to the inference engine for completion. `segments` field in `CompletionRequest` will be ignored.\n\nThis is useful for certain requests that aim to test the tabby's e2e quality.","nullable":true},"return_snippets":{"type":"boolean","description":"When true, returns `snippets` in `debug_data`."},"return_prompt":{"type":"boolean","description":"When true, returns `prompt` in `debug_data`."},"disable_retrieval_augmented_code_completion":{"type":"boolean","description":"When true, disable retrieval augmented code completion."}}},"Declaration":{"type":"object","description":"A snippet of declaration code that is relevant to the current completion request.","required":["filepath","body"],"properties":{"filepath":{"type":"string","description":"Filepath of the file where the snippet is from.\n- When the file belongs to the same workspace as the current file,\nthis is a relative filepath, use the same rule as [Segments::filepath].\n- When the file located outside the workspace, such as in a dependency package,\nthis is a file URI with an absolute filepath."},"body":{"type":"string","description":"Body of the snippet."}}},"DocSearchDocument":{"type":"object","required":["title","link","snippet"],"properties":{"title":{"type":"string"},"link":{"type":"string"},"snippet":{"type":"string"}}},"DocSearchHit":{"type":"object","required":["score","doc"],"properties":{"score":{"type":"number","format":"float"},"doc":{"$ref":"#/components/schemas/DocSearchDocument"}}},"HealthState":{"type":"object","required":["device","arch","cpu_info","cpu_count","cuda_devices","version"],"properties":{"model":{"type":"string","nullable":true},"chat_model":{"type":"string","nullable":true},"chat_device":{"type":"string","nullable":true},"device":{"type":"string"},"arch":{"type":"string"},"cpu_info":{"type":"string"},"cpu_count":{"type":"integer","minimum":0},"cuda_devices":{"type":"array","items":{"type":"string"}},"version":{"$ref":"#/components/schemas/Version"},"webserver":{"type":"boolean","nullable":true}}},"LogEventRequest":{"type":"object","required":["type","completion_id","choice_index"],"properties":{"type":{"type":"string","description":"Event type, should be `view`, `select` or `dismiss`.","example":"view"},"completion_id":{"type":"string"},"choice_index":{"type":"integer","format":"int32","minimum":0},"view_id":{"type":"string","nullable":true},"elapsed":{"type":"integer","format":"int32","nullable":true,"minimum":0}}},"Segments":{"type":"object","required":["prefix"],"properties":{"prefix":{"type":"string","description":"Content that appears before the cursor in the editor window."},"suffix":{"type":"string","description":"Content that appears after the cursor in the editor window.","nullable":true},"filepath":{"type":"string","description":"The relative path of the file that is being edited.\n- When [Segments::git_url] is set, this is the path of the file in the git repository.\n- When [Segments::git_url] is empty, this is the path of the file in the workspace.","nullable":true},"git_url":{"type":"string","description":"The remote URL of the current git repository.\nLeave this empty if the file is not in a git repository,\nor the git repository does not have a remote URL.","nullable":true},"declarations":{"type":"array","items":{"$ref":"#/components/schemas/Declaration"},"description":"The relevant declaration code snippets provided by the editor's LSP,\ncontain declarations of symbols extracted from [Segments::prefix].","nullable":true},"relevant_snippets_from_changed_files":{"type":"array","items":{"$ref":"#/components/schemas/Snippet"},"description":"The relevant code snippets extracted from recently edited files.\nThese snippets are selected from candidates found within code chunks\nbased on the edited location.\nThe current editing file is excluded from the search candidates.\n\nWhen provided alongside [Segments::declarations], the snippets have\nalready been deduplicated to ensure no duplication with entries\nin [Segments::declarations].\n\nSorted in descending order of [Snippet::score].","nullable":true},"relevant_snippets_from_recently_opened_files": {"type": "array","items": {"$ref": "#/components/schemas/Snippet"},"description": "The relevant code snippets extracted from recently opened files. These snippets are selected from candidates found within code chunks based on the last visited location. Current Active file is excluded from the search candidates. When provided with [Segments::relevant_snippets_from_changed_files], the snippets have already been deduplicated to ensure no duplication with entries in [Segments::relevant_snippets_from_changed_files].","nullable": true},"clipboard":{"type":"string","description":"Clipboard content when requesting code completion.","nullable":true}}},"ServerSetting":{"type":"object","required":["disable_client_side_telemetry"],"properties":{"disable_client_side_telemetry":{"type":"boolean"}}},"Snippet":{"type":"object","required":["filepath","body","score"],"properties":{"filepath":{"type":"string"},"body":{"type":"string"},"score":{"type":"number","format":"float"}}},"Version":{"type":"object","required":["build_date","build_timestamp","git_sha","git_describe"],"properties":{"build_date":{"type":"string"},"build_timestamp":{"type":"string"},"git_sha":{"type":"string"},"git_describe":{"type":"string"}}}},"securitySchemes":{"token":{"type":"http","scheme":"bearer","bearerFormat":"token"}}}} \ No newline at end of file +{"openapi":"3.0.3","info":{"title":"Tabby Server","description":"\n[![tabby stars](https://img.shields.io/github/stars/TabbyML/tabby)](https://github.com/TabbyML/tabby)\n[![Join Slack](https://shields.io/badge/Join-Tabby%20Slack-red?logo=slack)](https://links.tabbyml.com/join-slack)\n\nInstall following IDE / Editor extensions to get started with [Tabby](https://github.com/TabbyML/tabby).\n* [VSCode Extension](https://github.com/TabbyML/tabby/tree/main/clients/vscode) – Install from the [marketplace](https://marketplace.visualstudio.com/items?itemName=TabbyML.vscode-tabby), or [open-vsx.org](https://open-vsx.org/extension/TabbyML/vscode-tabby)\n* [VIM Extension](https://github.com/TabbyML/tabby/tree/main/clients/vim)\n* [IntelliJ Platform Plugin](https://github.com/TabbyML/tabby/tree/main/clients/intellij) – Install from the [marketplace](https://plugins.jetbrains.com/plugin/22379-tabby)\n","contact":{"name":"TabbyML Team"},"license":{"name":"Apache 2.0","url":"https://github.com/TabbyML/tabby/blob/main/LICENSE"},"version":"0.27.0-dev.0"},"servers":[{"url":"/","description":"Server"}],"paths":{"/v1/chat/completions":{"post":{"tags":["v1"],"operationId":"chat_completions","requestBody":{"description":"","content":{"application/json":{"schema":{}}},"required":true},"responses":{"200":{"description":"Success"},"405":{"description":"When chat model is not specified, the endpoint returns 405 Method Not Allowed"},"422":{"description":"When the prompt is malformed, the endpoint returns 422 Unprocessable Entity"}},"security":[{"token":[]}]}},"/v1/completions":{"post":{"tags":["v1"],"operationId":"completion","requestBody":{"content":{"application/json":{"schema":{"$ref":"#/components/schemas/CompletionRequest"}}},"required":true},"responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/CompletionResponse"}}}},"400":{"description":"Bad Request"}},"security":[{"token":[]}]}},"/v1/events":{"post":{"tags":["v1"],"operationId":"event","requestBody":{"content":{"application/json":{"schema":{"$ref":"#/components/schemas/LogEventRequest"}}},"required":true},"responses":{"200":{"description":"Success"},"400":{"description":"Bad Request"}},"security":[{"token":[]}]}},"/v1/health":{"get":{"tags":["v1"],"operationId":"health","responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/HealthState"}}}}},"security":[{"token":[]}]}},"/v1beta/server_setting":{"get":{"tags":["v1beta"],"operationId":"config","responses":{"200":{"description":"Success","content":{"application/json":{"schema":{"$ref":"#/components/schemas/ServerSetting"}}}}},"security":[{"token":[]}]}}},"components":{"schemas":{"Choice":{"type":"object","required":["index","text"],"properties":{"index":{"type":"integer","format":"int32","minimum":0},"text":{"type":"string"},"edit_range":{"allOf":[{"$ref":"#/components/schemas/EditRange"}],"nullable":true}}},"CompletionRequest":{"type":"object","properties":{"language":{"type":"string","description":"Language identifier, full list is maintained at\nhttps://code.visualstudio.com/docs/languages/identifiers","example":"python","nullable":true},"segments":{"allOf":[{"$ref":"#/components/schemas/Segments"}],"nullable":true},"user":{"type":"string","description":"A unique identifier representing your end-user, which can help Tabby to monitor & generating\nreports.","nullable":true},"debug_options":{"allOf":[{"$ref":"#/components/schemas/DebugOptions"}],"nullable":true},"temperature":{"type":"number","format":"float","description":"The temperature parameter for the model, used to tune variance and \"creativity\" of the model output","nullable":true},"seed":{"type":"integer","format":"int64","description":"The seed used for randomly selecting tokens","nullable":true,"minimum":0},"mode":{"type":"string","description":"The mode for completion. Use 'standard' for normal code completions or 'next_edit_suggestion'\nto predict the next edit the user will make."}},"example":{"language":"python","segments":{"prefix":"def fib(n):\n ","suffix":"\n return fib(n - 1) + fib(n - 2)"}}},"CompletionResponse":{"type":"object","required":["id","choices"],"properties":{"id":{"type":"string"},"choices":{"type":"array","items":{"$ref":"#/components/schemas/Choice"}},"debug_data":{"allOf":[{"$ref":"#/components/schemas/DebugData"}],"nullable":true}},"example":{"choices":[{"index":0,"text":"string"}],"id":"string"}},"CurrentVersion":{"type":"object","description":"Current version of the code after all edits","required":["content","cursor_position"],"properties":{"content":{"type":"string","description":"Current content after all edits"},"cursor_position":{"$ref":"#/components/schemas/CursorPosition"}}},"CursorPosition":{"type":"object","description":"Cursor position in the current version","required":["line","character"],"properties":{"line":{"type":"integer","format":"int32","description":"Line number (0-based)","minimum":0},"character":{"type":"integer","format":"int32","description":"Character position within the line (0-based)","minimum":0}}},"DebugData":{"type":"object","properties":{"snippets":{"type":"array","items":{"$ref":"#/components/schemas/Snippet"},"nullable":true},"prompt":{"type":"string","nullable":true}}},"DebugOptions":{"type":"object","properties":{"raw_prompt":{"type":"string","description":"When `raw_prompt` is specified, it will be passed directly to the inference engine for completion. `segments` field in `CompletionRequest` will be ignored.\n\nThis is useful for certain requests that aim to test the tabby's e2e quality.","nullable":true},"return_snippets":{"type":"boolean","description":"When true, returns `snippets` in `debug_data`."},"return_prompt":{"type":"boolean","description":"When true, returns `prompt` in `debug_data`."},"disable_retrieval_augmented_code_completion":{"type":"boolean","description":"When true, disable retrieval augmented code completion."}}},"Declaration":{"type":"object","description":"A snippet of declaration code that is relevant to the current completion request.","required":["filepath","body"],"properties":{"filepath":{"type":"string","description":"Filepath of the file where the snippet is from.\n- When the file belongs to the same workspace as the current file,\nthis is a relative filepath, use the same rule as [Segments::filepath].\n- When the file located outside the workspace, such as in a dependency package,\nthis is a file URI with an absolute filepath."},"body":{"type":"string","description":"Body of the snippet."}}},"EditHistory":{"type":"object","description":"Contains information about edit history for next edit suggestion mode","required":["original_code","edits_diff","current_version"],"properties":{"original_code":{"type":"string","description":"Original code content before edits"},"edits_diff":{"type":"string","description":"Unified git-style diff of all edits made to the file"},"current_version":{"$ref":"#/components/schemas/CurrentVersion"}}},"EditRange":{"type":"object","description":"Range information for next edit suggestion mode","required":["start_line","start_character","end_line","end_character"],"properties":{"start_line":{"type":"integer","format":"int32","description":"Start line of the edit (0-based)","minimum":0},"start_character":{"type":"integer","format":"int32","description":"Start character position within the line (0-based)","minimum":0},"end_line":{"type":"integer","format":"int32","description":"End line of the edit (0-based)","minimum":0},"end_character":{"type":"integer","format":"int32","description":"End character position within the line (0-based)","minimum":0}}},"HealthState":{"type":"object","required":["device","arch","cpu_info","cpu_count","cuda_devices","version"],"properties":{"model":{"type":"string","nullable":true},"chat_model":{"type":"string","nullable":true},"chat_device":{"type":"string","nullable":true},"device":{"type":"string"},"arch":{"type":"string"},"cpu_info":{"type":"string"},"cpu_count":{"type":"integer","minimum":0},"cuda_devices":{"type":"array","items":{"type":"string"}},"version":{"$ref":"#/components/schemas/Version"},"webserver":{"type":"boolean","nullable":true}}},"LogEventRequest":{"type":"object","required":["type","completion_id","choice_index"],"properties":{"type":{"type":"string","description":"Event type, should be `view`, `select` or `dismiss`.","example":"view"},"completion_id":{"type":"string"},"choice_index":{"type":"integer","format":"int32","minimum":0},"view_id":{"type":"string","nullable":true},"elapsed":{"type":"integer","format":"int32","nullable":true,"minimum":0}}},"Segments":{"type":"object","required":["prefix"],"properties":{"prefix":{"type":"string","description":"Content that appears before the cursor in the editor window."},"suffix":{"type":"string","description":"Content that appears after the cursor in the editor window.","nullable":true},"filepath":{"type":"string","description":"The relative path of the file that is being edited.\n- When [Segments::git_url] is set, this is the path of the file in the git repository.\n- When [Segments::git_url] is empty, this is the path of the file in the workspace.","nullable":true},"git_url":{"type":"string","description":"The remote URL of the current git repository.\nLeave this empty if the file is not in a git repository,\nor the git repository does not have a remote URL.","nullable":true},"declarations":{"type":"array","items":{"$ref":"#/components/schemas/Declaration"},"description":"The relevant declaration code snippets provided by the editor's LSP,\ncontain declarations of symbols extracted from [Segments::prefix].","nullable":true},"relevant_snippets_from_changed_files":{"type":"array","items":{"$ref":"#/components/schemas/Snippet"},"description":"The relevant code snippets extracted from recently edited files.\nThese snippets are selected from candidates found within code chunks\nbased on the edited location.\nThe current editing file is excluded from the search candidates.\n\nWhen provided alongside [Segments::declarations], the snippets have\nalready been deduplicated to ensure no duplication with entries\nin [Segments::declarations].\n\nSorted in descending order of [Snippet::score].","nullable":true},"relevant_snippets_from_recently_opened_files":{"type":"array","items":{"$ref":"#/components/schemas/Snippet"},"description":"The relevant code snippets extracted from recently opened files.\nThese snippets are selected from candidates found within code chunks\nbased on the last visited location.\n\nCurrent Active file is excluded from the search candidates.\nWhen provided with [Segments::relevant_snippets_from_changed_files], the snippets have\nalready been deduplicated to ensure no duplication with entries\nin [Segments::relevant_snippets_from_changed_files].","nullable":true},"clipboard":{"type":"string","description":"Clipboard content when requesting code completion.","nullable":true},"edit_history":{"allOf":[{"$ref":"#/components/schemas/EditHistory"}],"nullable":true}}},"ServerSetting":{"type":"object","required":["disable_client_side_telemetry"],"properties":{"disable_client_side_telemetry":{"type":"boolean"}}},"Snippet":{"type":"object","required":["filepath","body","score"],"properties":{"filepath":{"type":"string"},"body":{"type":"string"},"score":{"type":"number","format":"float"}}},"Version":{"type":"object","required":["build_date","build_timestamp","git_sha","git_describe"],"properties":{"build_date":{"type":"string"},"build_timestamp":{"type":"string"},"git_sha":{"type":"string"},"git_describe":{"type":"string"}}}},"securitySchemes":{"token":{"type":"http","scheme":"bearer","bearerFormat":"token"}}}} \ No newline at end of file diff --git a/clients/vscode/package.json b/clients/vscode/package.json index abfe6afc6974..a3df734c82e9 100644 --- a/clients/vscode/package.json +++ b/clients/vscode/package.json @@ -44,6 +44,16 @@ "title": "Trigger Code Completion Manually", "category": "Tabby" }, + { + "command": "tabby.nextEditSuggestion.trigger", + "title": "Trigger Next Edit Suggestion", + "category": "Tabby" + }, + { + "command": "tabby.acceptNESCompletion", + "title": "Accept Next Edit Suggestion", + "category": "Tabby" + }, { "command": "tabby.connectToServer", "title": "Connect to Server...", @@ -210,6 +220,9 @@ "command": "tabby.inlineCompletion.trigger", "when": "tabby.inlineCompletionTriggerMode === 'manual' && !editorHasSelection && !inlineSuggestionsVisible" }, + { + "command": "tabby.nextEditSuggestion.trigger" + }, { "command": "tabby.updateToken", "when": "tabby.status === 'unauthorized'" @@ -412,6 +425,11 @@ "command": "tabby.inlineCompletion.trigger", "when": "tabby.inlineCompletionTriggerMode === 'manual' && editorTextFocus && !editorHasSelection && !inlineSuggestionsVisible" }, + { + "command": "tabby.acceptNESCompletion", + "key": "tab", + "when": "editorTextFocus && tabby.nesCompletionVisible && !editorReadonly && !suggestWidgetVisible" + }, { "command": "tabby.inlineCompletion.accept", "key": "tab", @@ -563,4 +581,4 @@ "uuid": "^9.0.0", "vscode-languageclient": "^9.0.1" } -} +} \ No newline at end of file diff --git a/clients/vscode/src/InlineCompletionProvider.ts b/clients/vscode/src/InlineCompletionProvider.ts index 996b30610486..8611a2dd78e9 100644 --- a/clients/vscode/src/InlineCompletionProvider.ts +++ b/clients/vscode/src/InlineCompletionProvider.ts @@ -9,19 +9,107 @@ import { SnippetString, Range, window, + TextEditorDecorationType, + ThemeColor, + MarkdownString, + DecorationOptions, + commands, + Selection, } from "vscode"; -import { InlineCompletionParams } from "vscode-languageclient"; -import { InlineCompletionRequest, InlineCompletionList, EventParams } from "tabby-agent"; +import { InlineCompletionRequest, InlineNESCompletionRequest, InlineCompletionList, EventParams } from "tabby-agent"; import { EventEmitter } from "events"; import { getLogger } from "./logger"; import { Client } from "./lsp/client"; import { Config } from "./Config"; +import { InlineCompletionParams } from "vscode-languageclient"; + +// Create unique decoration type for NES completions +const NES_LABEL = "✨ next edit"; + +// NES decoration types +const nesDecorationTypes: Record = {}; + +// NES content highlight decoration type - slight background color change +const nesContentHighlightType = window.createTextEditorDecorationType({ + backgroundColor: new ThemeColor("editor.symbolHighlightBackground"), + opacity: "0.3", + borderRadius: "3px", + isWholeLine: false, +}); + +// NES content replacement decoration type - strikethrough effect +const nesContentReplaceType = window.createTextEditorDecorationType({ + textDecoration: "line-through", + color: new ThemeColor("editorError.foreground"), + backgroundColor: new ThemeColor("diffEditor.removedTextBackground"), + opacity: "0.3", + isWholeLine: false, +}); + +// NES intelligent suggestion decoration type - deep green with high transparency +const nesIntelligentSuggestionType = window.createTextEditorDecorationType({ + backgroundColor: "#004400", + opacity: "0.15", + borderRadius: "3px", + fontWeight: "normal", + border: "1px dotted #006600", + isWholeLine: false, +}); + +// NES exact match content decoration type +const nesMatchedContentType = window.createTextEditorDecorationType({ + backgroundColor: new ThemeColor("diffEditor.unchangedTextBackground"), + border: "1px dotted " + new ThemeColor("editorHint.border"), + borderRadius: "2px", + opacity: "0.25", + isWholeLine: false, +}); + +// Create different decoration types for different scenarios +// nest edit label color +function getNESDecorationType(id: string): TextEditorDecorationType { + if (!nesDecorationTypes[id]) { + // Create a new decoration type with more prominent markers + nesDecorationTypes[id] = window.createTextEditorDecorationType({ + after: { + contentText: ` ${NES_LABEL}`, + color: new ThemeColor("editorInlayHint.foreground"), + fontStyle: "italic", + margin: "0 0 0 20px", + }, + light: { + after: { + color: "#1b80b2", + }, + }, + dark: { + after: { + color: "#69c0fa", + }, + }, + }); + } + + // Return a valid decoration type, create a new empty one as fallback if not found + return nesDecorationTypes[id] || window.createTextEditorDecorationType({}); +} + +// Clean up unused decoration types +function disposeNESDecorationType(id: string) { + const decorationType = nesDecorationTypes[id]; + if (decorationType) { + decorationType.dispose(); + // eslint-disable-next-line @typescript-eslint/no-dynamic-delete + delete nesDecorationTypes[id]; + } +} interface DisplayedCompletion { id: string; displayedAt: number; completions: InlineCompletionList; index: number; + isNES?: boolean; // Marks if this is an NES completion } export class InlineCompletionProvider extends EventEmitter implements InlineCompletionItemProvider { @@ -29,6 +117,9 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp private displayedCompletion: DisplayedCompletion | null = null; private ongoing: Promise | null = null; private triggerMode: "automatic" | "manual"; + private activeNESDecorations: Record = {}; // Track active NES decoration IDs + private temporaryNormalCompletion = false; + private normalCompletionAbortController: AbortController | null = null; constructor( private readonly client: Client, @@ -39,6 +130,357 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp this.config.on("updated", () => { this.triggerMode = this.config.inlineCompletionTriggerMode; }); + + // Listen for editor changes to update decorations + window.onDidChangeActiveTextEditor((editor) => { + if (editor && this.displayedCompletion?.isNES) { + // When editor changes, update NES decorations + this.updateNESDecorations(editor.document.uri.toString()); + } + }); + } + + /** + * Analyze differences between current text and NES completion + */ + private analyzeNESDifferences( + document: TextDocument, + position: Position, + completionText: string, + ): DecorationOptions[] { + const decorations: DecorationOptions[] = []; + const currentLineText = document.lineAt(position.line).text; + const currentLineSuffix = currentLineText.substring(position.character); + + // Handle multi-line completions + const completionLines = completionText.split("\n"); + + if (completionLines.length === 1 && completionLines[0]) { + // Single-line completion case + const completionLine = completionLines[0]; + + // Use character-level diff calculation for more detailed difference info + const diff = this.calculateStringDiff(currentLineSuffix, completionLine); + + // Use matching content decoration for matched parts + if (diff.matched.length > 0) { + // Simplified handling: create separate decorations for each matching character + for (const pos of diff.matched) { + decorations.push({ + range: new Range( + new Position(position.line, position.character + pos), + new Position(position.line, position.character + pos + 1), + ), + hoverMessage: new MarkdownString("**Preserved Content**\n\nThis text will be kept"), + }); + } + + // Handle newly added content + for (const pos of diff.added) { + decorations.push({ + range: new Range( + new Position(position.line, position.character + pos), + new Position(position.line, position.character + pos + 1), + ), + hoverMessage: new MarkdownString( + "**Suggested New Content**\n\n✨ Tabby intelligent suggestion to add this content", + ), + }); + } + + // If there is content to be deleted, add the entire suffix as deletion + if (diff.removed.length > 0 && currentLineSuffix.length > 0) { + decorations.push({ + range: new Range(position, new Position(position.line, position.character + currentLineSuffix.length)), + hoverMessage: new MarkdownString( + "**Content To Be Replaced**\n\nThis content will be removed when accepting the suggestion", + ), + }); + } + } else { + // Completely different, mark all as replacement and new content + if (completionLine.length > 0) { + decorations.push({ + range: new Range(position, new Position(position.line, position.character + completionLine.length)), + hoverMessage: new MarkdownString("**Suggested New Content**\n\n✨ Tabby intelligent complete suggestion"), + }); + } + + if (currentLineSuffix.length > 0) { + decorations.push({ + range: new Range(position, new Position(position.line, position.character + currentLineSuffix.length)), + hoverMessage: new MarkdownString( + "**Content To Be Replaced**\n\nThis content will be completely replaced when accepting the suggestion", + ), + }); + } + } + } else if (completionLines.length > 1) { + // Multi-line completion case + // First line + if (currentLineSuffix.length > 0) { + decorations.push({ + range: new Range(position, new Position(position.line, currentLineText.length)), + hoverMessage: new MarkdownString("**Content To Be Replaced**\n\nThe remainder of this line will be replaced"), + }); + } + + const firstLine = completionLines[0] || ""; + if (firstLine.length > 0) { + decorations.push({ + range: new Range(position, new Position(position.line, position.character + firstLine.length)), + hoverMessage: new MarkdownString( + "**Suggested New Content (Line 1)**\n\n✨ Tabby intelligent suggestion for first line", + ), + }); + } + + // Middle lines + for (let i = 1; i < completionLines.length; i++) { + const lineContent = completionLines[i] || ""; + if (lineContent.length > 0) { + decorations.push({ + range: new Range(new Position(position.line + i, 0), new Position(position.line + i, lineContent.length)), + hoverMessage: new MarkdownString( + `**Suggested New Content (Line ${i + 1})**\n\n✨ Tabby intelligent multi-line suggestion`, + ), + }); + } + } + } + + return decorations; + } + + /** + * Calculate character-level differences, returning deleted and added character indices + * @param oldText Original text + * @param newText New text + * @returns Returns the changed character positions + */ + private calculateStringDiff( + oldText: string, + newText: string, + ): { matched: number[]; added: number[]; removed: number[] } { + const result = { + matched: [] as number[], // Matching character positions (in the new text) + added: [] as number[], // Added character positions (in the new text) + removed: [] as number[], // Removed character positions (in the original text) + }; + + // Calculate matching prefix from the text beginning + let prefixLength = 0; + const minLength = Math.min(oldText.length, newText.length); + while (prefixLength < minLength && oldText[prefixLength] === newText[prefixLength]) { + result.matched.push(prefixLength); + prefixLength++; + } + + // Calculate matching suffix from the text end + let oldSuffixIndex = oldText.length - 1; + let newSuffixIndex = newText.length - 1; + while ( + oldSuffixIndex >= prefixLength && + newSuffixIndex >= prefixLength && + oldText[oldSuffixIndex] === newText[newSuffixIndex] + ) { + result.matched.unshift(newSuffixIndex); + oldSuffixIndex--; + newSuffixIndex--; + } + + // Calculate added and removed characters + for (let i = prefixLength; i <= newSuffixIndex; i++) { + if (!result.matched.includes(i)) { + result.added.push(i); + } + } + + for (let i = prefixLength; i <= oldSuffixIndex; i++) { + result.removed.push(i); + } + + return result; + } + + /** + * Update NES completion decorations + */ + private updateNESDecorations(documentUri: string) { + if (!this.displayedCompletion?.isNES) return; + + const editor = window.activeTextEditor; + if (!editor || editor.document.uri.toString() !== documentUri) return; + + const item = this.displayedCompletion.completions.items[this.displayedCompletion.index]; + if (!item || !item.range) return; + + // Create decoration position + const range = new Range( + item.range.start.line, + item.range.start.character, + item.range.end.line, + item.range.end.character, + ); + + // Get or create decoration type + const decorationType = getNESDecorationType(this.displayedCompletion.id); + + // Apply label decoration + editor.setDecorations(decorationType, [range]); + + // Apply content highlight decoration - using diff editor colors to highlight completion content + editor.setDecorations(nesContentHighlightType, [range]); + + // Analyze and apply detailed differences + const document = editor.document; + + // Analyze differences and apply specific decorations + const diffDecorations = this.analyzeNESDifferences( + document, + item.range.start instanceof Position + ? item.range.start + : new Position(item.range.start.line, item.range.start.character), + typeof item.insertText === "string" ? item.insertText : item.insertText.value, + ); + + const replaceDecorations = diffDecorations.filter((d) => { + const msg = d.hoverMessage; + return msg instanceof MarkdownString + ? msg.value.includes("Content To Be Replaced") + : String(msg).includes("Content To Be Replaced"); + }); + editor.setDecorations(nesContentReplaceType, replaceDecorations); + + const newContentDecorations = diffDecorations.filter((d) => { + const msg = d.hoverMessage; + return msg instanceof MarkdownString + ? msg.value.includes("Suggested New Content") + : String(msg).includes("Suggested New Content"); + }); + editor.setDecorations(nesIntelligentSuggestionType, newContentDecorations); + + const matchedDecorations = diffDecorations.filter((d) => { + const msg = d.hoverMessage; + return msg instanceof MarkdownString + ? msg.value.includes("Preserved Content") + : String(msg).includes("Preserved Content"); + }); + editor.setDecorations(nesMatchedContentType, matchedDecorations); + + if (!this.activeNESDecorations[documentUri]) { + this.activeNESDecorations[documentUri] = []; + } + + const decorations = this.activeNESDecorations[documentUri]; + if (decorations) { + decorations.push(this.displayedCompletion.id); + } + } + + /** + * Clean up decorations + */ + private clearNESDecorations(documentUri?: string) { + if (documentUri) { + const decorationIds = this.activeNESDecorations[documentUri] || []; + decorationIds.forEach((id) => disposeNESDecorationType(id)); + + const editor = window.activeTextEditor; + if (editor && editor.document.uri.toString() === documentUri) { + // Clear all decoration types + editor.setDecorations(nesContentHighlightType, []); + editor.setDecorations(nesContentReplaceType, []); + editor.setDecorations(nesIntelligentSuggestionType, []); + editor.setDecorations(nesMatchedContentType, []); + } + + if (documentUri in this.activeNESDecorations) { + // eslint-disable-next-line @typescript-eslint/no-dynamic-delete + delete this.activeNESDecorations[documentUri]; + } + } else { + const uris = Object.keys(this.activeNESDecorations); + uris.forEach((uri) => { + const decorationIds = this.activeNESDecorations[uri]; + if (decorationIds) { + decorationIds.forEach((id) => disposeNESDecorationType(id)); + } + + window.visibleTextEditors.forEach((editor) => { + if (editor.document.uri.toString() === uri) { + // Clear all decoration types + editor.setDecorations(nesContentHighlightType, []); + editor.setDecorations(nesContentReplaceType, []); + editor.setDecorations(nesIntelligentSuggestionType, []); + editor.setDecorations(nesMatchedContentType, []); + } + }); + }); + this.activeNESDecorations = {}; + } + } + + /** + * Triggers a next edit suggestion by sending a request to the LSP server + */ + public async provideManuallyTriggerNextEditSuggestionTest(): Promise { + this.logger.debug("Function provideNextEditSuggestion called."); + + const document = window.activeTextEditor?.document; + const position = window.activeTextEditor?.selection.active; + + if (!document || !position) { + this.logger.debug("No active document or position found."); + return; + } + + const params: InlineCompletionParams = { + context: { + triggerKind: InlineCompletionTriggerKind.Invoke, + selectedCompletionInfo: undefined, + }, + textDocument: { + uri: document.uri.toString(), + }, + position: { + line: position.line, + character: position.character, + }, + }; + + try { + const request: Promise = this.client.languageClient.sendRequest( + InlineCompletionRequest.method, + params, + ); + this.ongoing = request; + this.emit("didChangeLoading", true); + + const result = await this.ongoing; + this.ongoing = null; + this.emit("didChangeLoading", false); + + if (!result || result.items.length === 0) { + this.logger.debug("No next edit suggestions received."); + return; + } + + this.logger.debug("Next edit suggestion received:", result); + + window.showInformationMessage(`Next edit suggestion received with ${result.items.length} items`); + + this.logger.debug("Inline completions shown successfully."); + } catch (error) { + if (this.ongoing) { + this.ongoing = null; + this.emit("didChangeLoading", false); + } + this.logger.error("Error requesting next edit suggestion:", error); + + // Show error message to user for testing purposes + window.showErrorMessage(`Next edit suggestion failed: ${error instanceof Error ? error.message : String(error)}`); + } } get isLoading(): boolean { @@ -56,6 +498,20 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp if (this.displayedCompletion) { // auto dismiss by new completion this.handleEvent("dismiss"); + + // Clean up decorations + this.clearNESDecorations(document.uri.toString()); + + // Update NES visibility + this.updateNESVisibilityContext(false); + } + + // Cancel any pending normal completion if we're displaying temporary ones + if (this.temporaryNormalCompletion && this.normalCompletionAbortController) { + this.logger.debug("Canceling temporary normal completion as new request is initiated"); + this.normalCompletionAbortController.abort(); + this.normalCompletionAbortController = null; + this.temporaryNormalCompletion = false; } if (context.triggerKind === InlineCompletionTriggerKind.Automatic && this.triggerMode === "manual") { @@ -85,7 +541,8 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp return null; } - const params: InlineCompletionParams = { + // Base request parameters + const baseParams: InlineCompletionParams = { context, textDocument: { uri: document.uri.toString(), @@ -95,80 +552,303 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp character: position.character, }, }; - let request: Promise | undefined = undefined; + + // NES specific parameters + const nesParams: InlineCompletionParams = { + context, + textDocument: { + uri: document.uri.toString(), + }, + position: { + line: position.line, + character: position.character, + }, + }; + try { this.client.fileTrack.addingChangeEditor(window.activeTextEditor); - request = this.client.languageClient.sendRequest(InlineCompletionRequest.method, params, token); - this.ongoing = request; + + // Create an abort controller for normal completions + this.normalCompletionAbortController = new AbortController(); + + // Create both request promises + const normalRequest: Promise = this.client.languageClient.sendRequest( + InlineCompletionRequest.method, + baseParams, + token, + ); + + const nesRequest: Promise = this.client.languageClient.sendRequest( + InlineNESCompletionRequest.method, + nesParams, + token, + ); + + // Setup loading indicator + this.ongoing = Promise.race([normalRequest, nesRequest]); this.emit("didChangeLoading", true); - const result = await this.ongoing; - this.ongoing = null; - this.emit("didChangeLoading", false); - if (token.isCancellationRequested) { - return null; - } - if (!result || result.items.length === 0) { - return null; + // Create a flag to track if NES has completed + let nesCompleted = false; + + // Use Promise.race to handle NES priority + const firstResult = await Promise.race([ + nesRequest.then((result) => { + nesCompleted = true; + return { type: "nes", result }; + }), + normalRequest.then((result) => { + return { type: "normal", result }; + }), + ]); + + // If NES completed first and has valid results + if ( + firstResult.type === "nes" && + firstResult.result && + firstResult.result.items && + firstResult.result.items.length > 0 + ) { + this.logger.info("NES completed first with valid results, using NES"); + + // Update NES visibility to true + this.updateNESVisibilityContext(true); + + // Cancel normal completion + if (this.normalCompletionAbortController) { + this.normalCompletionAbortController.abort(); + this.normalCompletionAbortController = null; + } + + this.ongoing = null; + this.emit("didChangeLoading", false); + + // Process and return NES completion + return this.processNESCompletions(document, firstResult.result); } - this.handleEvent("show", result); + // If normal completed first, but we need to wait for NES + if (firstResult.type === "normal" && !nesCompleted) { + // Show normal completion temporarily while waiting for NES + if (firstResult.result && firstResult.result.items && firstResult.result.items.length > 0) { + this.logger.info("Normal completed first, showing temporary results while waiting for NES"); + this.temporaryNormalCompletion = true; - return result.items.map((item, index) => { - return new InlineCompletionItem( - typeof item.insertText === "string" ? item.insertText : new SnippetString(item.insertText.value), - item.range - ? new Range( - item.range.start.line, - item.range.start.character, - item.range.end.line, - item.range.end.character, - ) - : undefined, - { - title: "", - command: "tabby.applyCallback", - arguments: [ - () => { - this.handleEvent("accept", result, index); + // Start a new promise to wait for NES result + nesRequest + .then((nesResult) => { + // If NES eventually returns valid results, we should replace normal completions + if (nesResult && nesResult.items && nesResult.items.length > 0 && this.temporaryNormalCompletion) { + this.logger.info("NES completed after normal, replacing with NES results"); + this.temporaryNormalCompletion = false; + + // Force a completion refresh to show NES results + const editor = window.activeTextEditor; + if (editor) { + // Use empty string to trigger a minimal change that forces refresh + editor.edit((editBuilder) => { + editBuilder.insert(editor.selection.active, ""); + }); + } + } + }) + .catch((err) => { + this.logger.error("Error waiting for NES after normal completion:", err); + this.temporaryNormalCompletion = false; + }); + + // Show normal completion temporarily + this.handleEvent("show", firstResult.result); + + // Setup normal completion items + return firstResult.result.items.map((item, index) => { + return new InlineCompletionItem( + typeof item.insertText === "string" ? item.insertText : new SnippetString(item.insertText.value), + item.range + ? new Range( + item.range.start.line, + item.range.start.character, + item.range.end.line, + item.range.end.character, + ) + : undefined, + { + title: "", + command: "tabby.applyCallback", + arguments: [ + () => { + this.handleEvent("accept", firstResult.result, index); + this.temporaryNormalCompletion = false; + }, + ], }, - ], - }, - ); - }); + ); + }); + } + } + + // If both completed, wait for final results + this.ongoing = null; + this.emit("didChangeLoading", false); + + // Wait for both to complete + const [normalResult, nesResult] = await Promise.all([ + nesCompleted + ? Promise.resolve(firstResult.type === "nes" ? firstResult.result : null) + : normalRequest.catch((err) => { + this.logger.error("Error in normal completion request:", err); + return null; + }), + nesCompleted + ? Promise.resolve(firstResult.result) + : nesRequest.catch((err) => { + this.logger.error("Error in NES completion request:", err); + return null; + }), + ]); + + // Always prioritize NES if it has results + if (nesResult && nesResult.items && nesResult.items.length > 0) { + this.logger.info("Using NES result for completion (final decision)"); + return this.processNESCompletions(document, nesResult); + } else if (normalResult && normalResult.items && normalResult.items.length > 0) { + this.handleEvent("show", normalResult); + + return normalResult.items.map((item, index) => { + return new InlineCompletionItem( + typeof item.insertText === "string" ? item.insertText : new SnippetString(item.insertText.value), + item.range + ? new Range( + item.range.start.line, + item.range.start.character, + item.range.end.line, + item.range.end.character, + ) + : undefined, + { + title: "", + command: "tabby.applyCallback", + arguments: [ + () => { + this.handleEvent("accept", normalResult, index); + }, + ], + }, + ); + }); + } + + // No results from either source + return null; } catch (error) { - if (this.ongoing === request) { - // the request was not replaced by a new request + if (this.ongoing) { this.ongoing = null; this.emit("didChangeLoading", false); } + this.logger.error("Error in completion requests:", error); return null; + } finally { + // Clean up any temporary state + this.temporaryNormalCompletion = false; + if (this.normalCompletionAbortController) { + this.normalCompletionAbortController = null; + } } } + /** + * Process NES completions with special styling and markers + */ + private processNESCompletions(_document: TextDocument, nesResult: InlineCompletionList): InlineCompletionItem[] { + // Record as NES completion + this.handleEvent("show", nesResult, 0, true); + + // Update NES visibility + this.updateNESVisibilityContext(true); + + // Special handling for NES + return nesResult.items.map((item, index) => { + // Create a completion item with NES marker + const insertText = + typeof item.insertText === "string" ? item.insertText : new SnippetString(item.insertText.value); + + const range = item.range + ? new Range(item.range.start.line, item.range.start.character, item.range.end.line, item.range.end.character) + : undefined; + + // Create custom documentation + const documentation = new MarkdownString(); + documentation.appendMarkdown(`**${NES_LABEL}**\n\n`); + documentation.appendMarkdown("Smart next edit suggestion"); + documentation.isTrusted = true; + + // Create standard VS Code inline completion item + const completionItem = new InlineCompletionItem(insertText, range); + + completionItem.command = { + title: "", + command: "tabby.applyCallback", + arguments: [ + () => { + this.logger.info(`NES completion accepted through standard mechanism, index ${index}`); + this.handleEvent("accept", nesResult, index); + + if (window.activeTextEditor) { + this.clearNESDecorations(window.activeTextEditor.document.uri.toString()); + } + + this.updateNESVisibilityContext(false); + }, + ], + }; + + return completionItem; + }); + } + // FIXME: We don't listen to the user cycling through the items, // so we don't know the 'index' (except for the 'accept' event). // For now, just use the first item to report other events. async handleEvent( event: "show" | "accept" | "dismiss" | "accept_word" | "accept_line", - completions?: InlineCompletionList, + completions?: InlineCompletionList | null, index = 0, + isNES = false, ) { if (event === "show" && completions) { const item = completions.items[index]; - const cmplId = item?.data?.eventId?.completionId.replace("cmpl-", ""); + const cmplId = item?.data?.eventId?.completionId?.replace("cmpl-", "") || Math.random().toString(36).substring(7); const timestamp = Date.now(); this.displayedCompletion = { id: `view-${cmplId}-at-${timestamp}`, completions, index, displayedAt: timestamp, + isNES, }; + + // If NES completion, add decoration + if (isNES && window.activeTextEditor) { + this.updateNESDecorations(window.activeTextEditor.document.uri.toString()); + // Update NES visibility + this.updateNESVisibilityContext(true); + } + await this.postEvent(event, this.displayedCompletion); } else if (this.displayedCompletion) { this.displayedCompletion.index = index; await this.postEvent(event, this.displayedCompletion); - this.displayedCompletion = null; + + // Cleanup after event handling + if (event === "accept" || event === "dismiss") { + // If NES completion, clean up decorations + if (this.displayedCompletion.isNES && window.activeTextEditor) { + this.clearNESDecorations(window.activeTextEditor.document.uri.toString()); + // Update NES visibility + this.updateNESVisibilityContext(false); + } + this.displayedCompletion = null; + } } } @@ -176,7 +856,7 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp event: "show" | "accept" | "dismiss" | "accept_word" | "accept_line", displayedCompletion: DisplayedCompletion, ) { - const { id, completions, index, displayedAt } = displayedCompletion; + const { id, completions, index, displayedAt, isNES } = displayedCompletion; const eventId = completions.items[index]?.data?.eventId; if (!eventId) { return; @@ -188,7 +868,12 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp eventData = { type: "view" }; break; case "accept": - eventData = { type: "select", elapsed }; + eventData = { + type: "select", + elapsed, + // Add extra marker for NES completions + selectKind: isNES ? "line" : undefined, + }; break; case "dismiss": eventData = { type: "dismiss", elapsed }; @@ -235,4 +920,106 @@ export class InlineCompletionProvider extends EventEmitter implements InlineComp this.logger.debug("Calculate edited range for displayed completion item:", completionRange); return completionRange; } + + // Function to accept the current NES completion + public acceptCurrentNESCompletion() { + if (this.displayedCompletion?.isNES) { + this.logger.info("Manual acceptance of NES completion triggered"); + const item = this.displayedCompletion.completions.items[this.displayedCompletion.index]; + + // Only proceed if there's an active item + if (item) { + const editor = window.activeTextEditor; + if (editor) { + // Get the insertion range + const range = item.range + ? new Range( + item.range.start.line, + item.range.start.character, + item.range.end.line, + item.range.end.character, + ) + : editor.selection; + + // Check if cursor is within the NES range + const cursorPosition = editor.selection.active; + const isInRange = range.contains(cursorPosition); + + if (!isInRange) { + this.logger.info("Cursor not in NES range, skipping acceptance"); + return; + } + + // Get the text to insert + const insertText = typeof item.insertText === "string" ? item.insertText : item.insertText.value; + + this.logger.info( + `Applying NES edit: "${insertText}" at range ${range.start.line}:${range.start.character} to ${range.end.line}:${range.end.character}`, + ); + + // Perform the edit + editor + .edit((editBuilder) => { + editBuilder.replace(range, insertText); + }) + .then((success) => { + if (success) { + this.logger.info("NES edit successfully applied"); + + // Move cursor to the end of inserted text + if (insertText) { + const textLines = insertText.split("\n"); + if (textLines.length > 0) { + const lastLineLength = textLines[textLines.length - 1]?.length || 0; + const newLine = range.start.line + textLines.length - 1; + const newCharacter = textLines.length > 1 ? lastLineLength : range.start.character + lastLineLength; + + const newPosition = new Position(newLine, newCharacter); + editor.selection = new Selection(newPosition, newPosition); + } + } + + // Handle the acceptance event + this.handleEvent("accept", this.displayedCompletion?.completions, this.displayedCompletion?.index || 0); + + // Clean up decorations + if (window.activeTextEditor) { + this.clearNESDecorations(window.activeTextEditor.document.uri.toString()); + } + + // Clear the displayed completion + this.displayedCompletion = null; + + // Update NES visibility + this.updateNESVisibilityContext(false); + } else { + this.logger.error("Failed to apply NES edit"); + } + }); + } else { + this.logger.error("No active editor found when trying to accept NES completion"); + } + } else { + this.logger.warn("No item found in displayedCompletion when trying to accept NES completion"); + } + } else { + this.logger.warn("No NES completion to accept"); + } + } + + // Update NES visibility context variable for keybinding conditions + private updateNESVisibilityContext(visible: boolean) { + commands.executeCommand("setContext", "tabby.nesCompletionVisible", visible); + this.logger.info(`NES completion visibility set to: ${visible}`); + } + + /** + * Handle resource disposal + */ + public dispose() { + // Clean up all decorations + this.clearNESDecorations(); + // Reset context + this.updateNESVisibilityContext(false); + } } diff --git a/clients/vscode/src/commands/index.ts b/clients/vscode/src/commands/index.ts index 7a9f7098c397..3a0f122cd421 100644 --- a/clients/vscode/src/commands/index.ts +++ b/clients/vscode/src/commands/index.ts @@ -55,6 +55,12 @@ export class Commands { applyCallback: (callback: (() => void) | undefined) => { callback?.(); }, + acceptNESCompletion: () => { + console.log("acceptNESCompletion command triggered from Commands class"); + if (this.inlineCompletionProvider) { + this.inlineCompletionProvider.acceptCurrentNESCompletion(); + } + }, toggleInlineCompletionTriggerMode: async (value: "automatic" | "manual" | undefined) => { let target = value; if (!target) { @@ -191,6 +197,11 @@ export class Commands { "inlineCompletion.trigger": () => { commands.executeCommand("editor.action.inlineSuggest.trigger"); }, + "nextEditSuggestion.trigger": () => { + console.log("Triggering next edit suggestion"); + // Call internal next edit suggestion implementation + this.inlineCompletionProvider.provideManuallyTriggerNextEditSuggestionTest(); + }, "inlineCompletion.accept": async () => { const editor = window.activeTextEditor; if (!editor) { diff --git a/clients/vscode/src/lsp/InlineNESCompletionFeature.ts b/clients/vscode/src/lsp/InlineNESCompletionFeature.ts new file mode 100644 index 000000000000..5fcc9d2450b8 --- /dev/null +++ b/clients/vscode/src/lsp/InlineNESCompletionFeature.ts @@ -0,0 +1,52 @@ +import { BaseLanguageClient, DynamicFeature, FeatureState, RegistrationData } from "vscode-languageclient"; +import { InlineNESCompletionRequest } from "tabby-agent"; +import { getLogger } from "../logger"; + +/** + * Implementation of the dynamic feature for Next Edit Suggestion (NES) completions + */ +export class InlineNESCompletionFeature implements DynamicFeature { + private readonly logger = getLogger("InlineNESCompletionFeature"); + private registration: string | undefined = undefined; + + constructor(private readonly client: BaseLanguageClient) {} + + // Required properties + readonly registrationType = InlineNESCompletionRequest.type; + + getState(): FeatureState { + return { kind: "workspace", id: this.registrationType.method, registrations: this.isAvailable }; + } + + get isAvailable(): boolean { + return !!this.registration; + } + + fillClientCapabilities() { + // No special client capabilities needed + } + + initialize() { + // Register request handler + this.client.onRequest(InlineNESCompletionRequest.type, (params, token) => { + this.logger.debug("Handling InlineNESCompletionRequest via feature handler"); + return this.client.sendRequest(InlineNESCompletionRequest.method, params, token); + }); + } + + register(data: RegistrationData): void { + this.registration = data.id; + this.logger.debug(`Registered InlineNESCompletionRequest with id: ${this.registration}`); + } + + unregister(id: string): void { + if (this.registration === id) { + this.registration = undefined; + this.logger.debug(`Unregistered InlineNESCompletionRequest with id: ${id}`); + } + } + + clear(): void { + this.registration = undefined; + } +} diff --git a/clients/vscode/src/lsp/client.ts b/clients/vscode/src/lsp/client.ts index 4b0deefeab43..1db308e882fd 100644 --- a/clients/vscode/src/lsp/client.ts +++ b/clients/vscode/src/lsp/client.ts @@ -14,6 +14,7 @@ import { EditorOptionsFeature } from "./EditorOptionsFeature"; import { GitProviderFeature } from "./GitProviderFeature"; import { InitializationFeature } from "./InitializationFeature"; import { InlineCompletionFeature } from "./InlineCompletionFeature"; +import { InlineNESCompletionFeature } from "./InlineNESCompletionFeature"; import { LanguageSupportFeature } from "./LanguageSupportFeature"; import { TelemetryFeature } from "./TelemetryFeature"; import { WorkspaceFileSystemFeature } from "./WorkspaceFileSystemFeature"; @@ -25,6 +26,13 @@ import { WorkSpaceFeature } from "./WorkspaceFeature"; import { FileTrackerFeature } from "./FileTrackFeature"; import { isBrowser } from "../env"; +let outputChannelTransport: { log: (message: string) => void } | undefined; +export function logToOutputChannel(message: string) { + if (outputChannelTransport) { + outputChannelTransport.log(message); + } +} + export function createClient(context: ExtensionContext, logger: LogOutputChannel): Client { const clientOptions: LanguageClientOptions = { documentSelector: [ @@ -69,7 +77,7 @@ export class Client { constructor( private readonly context: ExtensionContext, - readonly languageClient: BaseLanguageClient, + public readonly languageClient: BaseLanguageClient, ) { this.status = new AgentStatusFeature(this.languageClient); this.agentConfig = new AgentConfigFeature(this.languageClient); @@ -88,9 +96,17 @@ export class Client { this.languageClient.registerFeature(new LanguageSupportFeature(this.languageClient)); this.languageClient.registerFeature(new WorkspaceFileSystemFeature(this.languageClient)); + this.languageClient.registerFeature(new InlineNESCompletionFeature(this.languageClient)); + const codeLensMiddleware = new CodeLensMiddleware(); this.languageClient.middleware.provideCodeLenses = (document, token, next) => codeLensMiddleware.provideCodeLenses(document, token, next); + + outputChannelTransport = { + log: (message: string) => { + this.languageClient.outputChannel.appendLine(message); + }, + }; } async start(): Promise { diff --git a/crates/http-api-bindings/src/completion/llama.rs b/crates/http-api-bindings/src/completion/llama.rs index 45858d9fb240..35524d4fb511 100644 --- a/crates/http-api-bindings/src/completion/llama.rs +++ b/crates/http-api-bindings/src/completion/llama.rs @@ -3,7 +3,7 @@ use async_trait::async_trait; use futures::{stream::BoxStream, StreamExt}; use reqwest_eventsource::{Event, EventSource}; use serde::{Deserialize, Serialize}; -use tabby_inference::{CompletionOptions, CompletionStream}; +use tabby_inference::{code::NEXT_EDIT_SUGGESTION_MODE, CompletionOptions, CompletionStream}; use crate::create_reqwest_client; @@ -45,41 +45,107 @@ struct CompletionResponseChunk { #[async_trait] impl CompletionStream for LlamaCppEngine { async fn generate(&self, prompt: &str, options: CompletionOptions) -> BoxStream { - let request = CompletionRequest { + let use_stream = !(options.mode == NEXT_EDIT_SUGGESTION_MODE); + + let request_body = CompletionRequest { seed: options.seed, prompt: prompt.to_owned(), n_predict: options.max_decoding_tokens, temperature: options.sampling_temperature, - stream: true, + stream: use_stream, penalty_last_n: 0, presence_penalty: options.presence_penalty, }; - let mut request = self.client.post(&self.api_endpoint).json(&request); + tracing::info!("Request URL: {}", &self.api_endpoint); + let request_json = serde_json::to_string_pretty(&request_body) + .unwrap_or_else(|_| "Failed to serialize request".to_string()); + tracing::info!("Request body: \n{}", request_json); + + let mut request = self.client.post(&self.api_endpoint).json(&request_body); if let Some(api_key) = &self.api_key { request = request.bearer_auth(api_key); } - let s = stream! { - let mut es = EventSource::new(request).expect("Failed to create event source"); - while let Some(event) = es.next().await { - match event { - Ok(Event::Open) => {} - Ok(Event::Message(message)) => { - let x: CompletionResponseChunk = serde_json::from_str(&message.data).unwrap(); - yield x.content.clone(); - if x.stop { + if use_stream { + let s = stream! { + let mut es = EventSource::new(request).expect("Failed to create event source"); + while let Some(event) = es.next().await { + match event { + Ok(Event::Open) => {} + Ok(Event::Message(message)) => { + let x: CompletionResponseChunk = serde_json::from_str(&message.data).unwrap(); + yield x.content.clone(); + if x.stop { + break; + } + } + Err(_) => { + // StreamEnd break; } } - Err(_) => { - // StreamEnd - break; - } } - } - }; + }; + + Box::pin(s) + } else { + let s = stream! { + tracing::info!("Using non-streaming mode for next_edit_suggestion"); + + match self.client.post(&self.api_endpoint) + .json(&request_body) + .send() + .await { + Ok(response) => { + if !response.status().is_success() { + tracing::error!("Request failed with status: {}", response.status()); + yield String::new(); + } else { + let response_text = match response.text().await { + Ok(text) => text, + Err(e) => { + tracing::error!("Failed to get response text: {}", e); + yield String::new(); + return; + } + }; - Box::pin(s) + tracing::info!("Received raw response: {}", response_text); + + if let Ok(json) = serde_json::from_str::(&response_text) { + if let Some(choices) = json.get("choices").and_then(|c| c.as_array()) { + if let Some(choice) = choices.first() { + if let Some(text) = choice.get("text").and_then(|t| t.as_str()) { + tracing::info!("Extracted text from JSON choices: {}", text); + yield text.to_string(); + return; + } + } + } + + if let Some(content) = json.get("content").and_then(|c| c.as_str()) { + tracing::info!("Extracted content from JSON: {}", content); + yield content.to_string(); + return; + } + + tracing::error!("No recognized content field in JSON: {:?}", json); + yield String::new(); + } else { + tracing::error!("Response is not valid JSON: {}", response_text); + yield String::new(); + } + } + }, + Err(e) => { + tracing::error!("Request failed: {}", e); + yield String::new(); + } + } + }; + + Box::pin(s) + } } } diff --git a/crates/tabby-inference/src/code.rs b/crates/tabby-inference/src/code.rs index 72b027ea1b92..638ec3115c1e 100644 --- a/crates/tabby-inference/src/code.rs +++ b/crates/tabby-inference/src/code.rs @@ -9,6 +9,9 @@ use crate::{ clip_prompt, decoding::StopConditionFactory, CompletionOptionsBuilder, CompletionStream, }; +pub const DEFAULT_MODE: &str = "standard"; +pub const NEXT_EDIT_SUGGESTION_MODE: &str = "next_edit_suggestion"; + #[derive(Builder, Debug)] pub struct CodeGenerationOptions { #[builder(default = "1024")] @@ -25,6 +28,9 @@ pub struct CodeGenerationOptions { #[builder(default = "None")] pub language: Option<&'static Language>, + + #[builder(default = "String::from(DEFAULT_MODE)")] + pub mode: String, } /// CodeGeneration utilizes the CompletionStream to generate code completions. @@ -71,10 +77,12 @@ impl CodeGeneration { .max_decoding_tokens(options.max_decoding_tokens) .sampling_temperature(options.sampling_temperature) .seed(options.seed) + .mode(options.mode.clone()) .build() .expect("Failed to build completion options"); for await new_text in self.imp.generate(prompt, options).await { + // WARN: for regular fetch /completion, the stop_condition is used to stop some where by the keyword, it shoundn't happen in NES mode let (should_stop, stop_length) = stop_condition.should_stop(&new_text); text += &new_text; if should_stop { diff --git a/crates/tabby-inference/src/completion.rs b/crates/tabby-inference/src/completion.rs index 1bd05c84d796..387d1510cefa 100644 --- a/crates/tabby-inference/src/completion.rs +++ b/crates/tabby-inference/src/completion.rs @@ -12,6 +12,9 @@ pub struct CompletionOptions { #[builder(default = "0.0")] pub presence_penalty: f32, + + #[builder(default = "\"standard\".to_string()")] + pub mode: String, } #[async_trait] diff --git a/crates/tabby-inference/src/lib.rs b/crates/tabby-inference/src/lib.rs index 1ad1947b241a..2ac156494629 100644 --- a/crates/tabby-inference/src/lib.rs +++ b/crates/tabby-inference/src/lib.rs @@ -1,6 +1,6 @@ //! Lays out the abstract definition of a text generation model, and utilities for encodings. mod chat; -mod code; +pub mod code; mod completion; mod decoding; mod embedding; diff --git a/crates/tabby/src/serve.rs b/crates/tabby/src/serve.rs index 252b915ced83..8a7a3928a285 100644 --- a/crates/tabby/src/serve.rs +++ b/crates/tabby/src/serve.rs @@ -35,6 +35,7 @@ use crate::{ to_local_config, Device, }; +// TODO: setup schemas for openapi.json here, for new nes #[derive(OpenApi)] #[openapi( info(title="Tabby Server", @@ -63,6 +64,10 @@ Install following IDE / Editor extensions to get started with [Tabby](https://gi completion::Snippet, completion::DebugOptions, completion::DebugData, + completion::EditHistory, + completion::CurrentVersion, + completion::CursorPosition, + completion::EditRange, health::HealthState, health::Version, api::server_setting::ServerSetting, diff --git a/crates/tabby/src/services/completion.rs b/crates/tabby/src/services/completion.rs index b771bf82d909..c78f3eda9c4a 100644 --- a/crates/tabby/src/services/completion.rs +++ b/crates/tabby/src/services/completion.rs @@ -1,4 +1,5 @@ mod completion_prompt; +mod next_edit_prompt; use std::sync::Arc; @@ -19,6 +20,7 @@ use tabby_inference::{ CompletionStream, }; use thiserror::Error; +use tracing::debug; use utoipa::ToSchema; use super::model; @@ -57,6 +59,48 @@ pub struct CompletionRequest { /// The seed used for randomly selecting tokens seed: Option, + + /// The mode for completion. Use 'standard' for normal code completions or 'next_edit_suggestion' + /// to predict the next edit the user will make. + #[serde(default = "default_standard_mode")] + mode: String, +} + +fn default_standard_mode() -> String { + "standard".to_string() +} + +/// Contains information about edit history for next edit suggestion mode +#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] +pub struct EditHistory { + /// Original code content before edits + original_code: String, + + /// Unified git-style diff of all edits made to the file + edits_diff: String, + + /// Current version of the code after all edits + current_version: CurrentVersion, +} + +/// Current version of the code after all edits +#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] +pub struct CurrentVersion { + /// Current content after all edits + content: String, + + /// Cursor position in the current version + cursor_position: CursorPosition, +} + +/// Cursor position in the current version +#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] +pub struct CursorPosition { + /// Line number (0-based) + line: u32, + + /// Character position within the line (0-based) + character: u32, } impl CompletionRequest { @@ -78,6 +122,11 @@ impl CompletionRequest { .as_ref() .is_some_and(|x| x.disable_retrieval_augmented_code_completion) } + + /// Returns true if the request is for next edit suggestion mode. + fn is_next_edit_suggestion_mode(&self) -> bool { + self.mode == "next_edit_suggestion" + } } #[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] @@ -150,6 +199,9 @@ pub struct Segments { /// Clipboard content when requesting code completion. clipboard: Option, + + /// Required when mode is 'next_edit_suggestion'. Contains information about edit history. + edit_history: Option, } impl From for api::event::Segments { @@ -194,11 +246,43 @@ impl From for api::event::Declaration { pub struct Choice { index: u32, text: String, + + /// Range information for next edit suggestion mode + #[serde(skip_serializing_if = "Option::is_none")] + edit_range: Option, +} + +/// Range information for next edit suggestion mode +#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)] +pub struct EditRange { + /// Start line of the edit (0-based) + start_line: u32, + + /// Start character position within the line (0-based) + start_character: u32, + + /// End line of the edit (0-based) + end_line: u32, + + /// End character position within the line (0-based) + end_character: u32, } impl Choice { pub fn new(text: String) -> Self { - Self { index: 0, text } + Self { + index: 0, + text, + edit_range: None, + } + } + + pub fn with_edit_range(text: String, edit_range: EditRange) -> Self { + Self { + index: 0, + text, + edit_range: Some(edit_range), + } } } @@ -249,6 +333,7 @@ pub struct CompletionService { engine: Arc, logger: Arc, prompt_builder: completion_prompt::PromptBuilder, + next_edit_prompt_builder: next_edit_prompt::NextEditPromptBuilder, } impl CompletionService { @@ -266,6 +351,7 @@ impl CompletionService { prompt_template, Some(code), ), + next_edit_prompt_builder: next_edit_prompt::NextEditPromptBuilder::new(), config, logger, } @@ -293,6 +379,7 @@ impl CompletionService { seed: Option, max_input_length: usize, max_output_tokens: usize, + mode: String, ) -> CodeGenerationOptions { let mut builder = CodeGenerationOptionsBuilder::default(); builder @@ -305,6 +392,10 @@ impl CompletionService { seed.inspect(|x| { builder.seed(*x); }); + + // Set mode + builder.mode(mode); + builder .build() .expect("Failed to create text generation options") @@ -318,14 +409,25 @@ impl CompletionService { ) -> Result { let completion_id = format!("cmpl-{}", uuid::Uuid::new_v4()); let language = request.language_or_unknown(); + + // Handle next edit suggestion mode + if request.is_next_edit_suggestion_mode() { + return self + .generate_next_edit_suggestion(request, completion_id, language, user_agent) + .await; + } + + // Standard mode let options = Self::text_generation_options( language.as_str(), request.temperature, request.seed, self.config.max_input_length, self.config.max_decoding_tokens, + request.mode.clone(), ); + debug!("Options: {:?}", options); let mut use_crlf = false; let (prompt, segments, snippets) = if let Some(prompt) = request.raw_prompt() { (prompt, None, vec![]) @@ -383,6 +485,91 @@ impl CompletionService { debug_data, )) } + + async fn generate_next_edit_suggestion( + &self, + request: &CompletionRequest, + completion_id: String, + language: String, + user_agent: Option<&str>, + ) -> Result { + // Validate that segments and edit_history are provided for next edit suggestion mode + let segments = request + .segments + .as_ref() + .ok_or_else(|| CompletionError::EmptyPrompt)?; + + let edit_history = segments + .edit_history + .as_ref() + .ok_or_else(|| CompletionError::EmptyPrompt)?; + + let prompt = self.next_edit_prompt_builder.build_prompt(edit_history); + + debug!("Next edit prompt: {}", prompt); + + // Generate completion with larger generation limit for edit suggestions + let options = Self::text_generation_options( + language.as_str(), + request.temperature, + request.seed, + self.config.max_input_length, + self.config.max_decoding_tokens * 2, + request.mode.clone(), + ); + + let generated_text = self.engine.generate(&prompt, options).await; + + debug!( + "Raw generated text (length: {}): '{}'", + generated_text.len(), + generated_text + ); + debug!("Is empty: {}", generated_text.is_empty()); + + debug!( + "Generated text for next edit suggestion: {}", + generated_text + ); + + let edit_range = self.next_edit_prompt_builder.parse_edit_range( + generated_text.as_str(), + edit_history.current_version.cursor_position.line, + edit_history.current_version.cursor_position.character, + ); + + // Log the completion + self.logger.log( + request.user.clone(), + Event::Completion { + completion_id: completion_id.clone(), + language, + prompt: prompt.clone(), + segments: None, + choices: vec![api::event::Choice { + index: 0, + text: generated_text.clone(), + }], + user_agent: user_agent.map(|x| x.to_owned()), + }, + ); + + // Create debug data if requested + let debug_data = request + .debug_options + .as_ref() + .map(|debug_options| DebugData { + snippets: None, + prompt: debug_options.return_prompt.then_some(prompt), + }); + + // Return completion response with edit range + Ok(CompletionResponse::new( + completion_id, + vec![Choice::with_edit_range(generated_text, edit_range)], + debug_data, + )) + } } fn contains_crlf(segments: &Segments) -> bool { @@ -515,6 +702,7 @@ mod tests { relevant_snippets_from_changed_files: None, relevant_snippets_from_recently_opened_files: None, clipboard: None, + edit_history: None, }; let request = CompletionRequest { language: Some("rust".into()), @@ -523,6 +711,7 @@ mod tests { debug_options: None, temperature: None, seed: None, + mode: "standard".into(), }; let allowed_code_repository = AllowedCodeRepository::default(); @@ -550,6 +739,7 @@ mod tests { relevant_snippets_from_changed_files: None, relevant_snippets_from_recently_opened_files: None, clipboard: None, + edit_history: None, }, Segments { prefix: "fn hello_world() -> &'static str {".into(), @@ -560,6 +750,7 @@ mod tests { relevant_snippets_from_changed_files: None, relevant_snippets_from_recently_opened_files: None, clipboard: None, + edit_history: None, }, Segments { prefix: "fn hello_world() -> &'static str {\r\n".into(), @@ -570,6 +761,7 @@ mod tests { relevant_snippets_from_changed_files: None, relevant_snippets_from_recently_opened_files: None, clipboard: None, + edit_history: None, }, ]; for segments in contained_crlf { @@ -585,6 +777,7 @@ mod tests { relevant_snippets_from_changed_files: None, relevant_snippets_from_recently_opened_files: None, clipboard: None, + edit_history: None, }]; for segments in not_contained_crlf { assert!(!contains_crlf(&segments)); diff --git a/crates/tabby/src/services/completion/next_edit_prompt.rs b/crates/tabby/src/services/completion/next_edit_prompt.rs new file mode 100644 index 000000000000..531f106de5e0 --- /dev/null +++ b/crates/tabby/src/services/completion/next_edit_prompt.rs @@ -0,0 +1,105 @@ +use super::{EditHistory, EditRange}; + +pub struct NextEditPromptBuilder; + +impl NextEditPromptBuilder { + pub fn new() -> Self { + Self + } + + pub fn build_prompt(&self, edit_history: &EditHistory) -> String { + let prompt = format!("<|original_code|>\n{}\n<|edits_diff|>\n{}\n<|current_version|>\n{}\n<|next_version|>\n", + edit_history.original_code, + edit_history.edits_diff, + edit_history.current_version.content + ); + + prompt + } + + pub fn parse_edit_range( + &self, + text: &str, + cursor_line: u32, + cursor_character: u32, + ) -> EditRange { + // For this basic implementation, we'll assume the edit happens at cursor position + // and extends to the end of the predicted text. In a real implementation, you would + // need more sophisticated parsing based on the model's output. + + let lines: Vec<&str> = text.lines().collect(); + let line_count = lines.len() as u32; + + let last_line = if line_count > 0 { + lines.last().unwrap() + } else { + "" + }; + + EditRange { + start_line: cursor_line, + start_character: cursor_character, + end_line: cursor_line + line_count.saturating_sub(1), + end_character: if line_count <= 1 { + cursor_character + last_line.len() as u32 + } else { + last_line.len() as u32 + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::services::completion::{CurrentVersion, CursorPosition}; + + #[test] + fn test_build_prompt() { + let edit_history = EditHistory { + original_code: "fn main() {\n println!(\"Hello, world!\");\n}".to_string(), + edits_diff: "---src/main.rs\n+++src/main.rs\n@@ -1,1 +1,2 @@\n println!(\"Hello, world!\");\n let x = 5;\n println!(\"Hello, world!\");".to_string(), + current_version: CurrentVersion { + content: "fn main() {\n let x = 5;\n println!(\"Hello, world!\");\n}".to_string(), + cursor_position: CursorPosition { + line: 2, + character: 0, + }, + }, + }; + + let builder = NextEditPromptBuilder::new(); + let prompt = builder.build_prompt(&edit_history); + + // Check that prompt contains all the important parts + assert!(prompt.contains("<|original_code|>")); + assert!(prompt.contains("<|edits_diff|>")); + assert!(prompt.contains("<|current_version|>")); + assert!(prompt.contains("fn main()")); + assert!(prompt.contains("let x = 5;")); + } + + #[test] + fn test_parse_edit_range() { + let builder = NextEditPromptBuilder::new(); + + // Test single line edit + let text = " println!(\"The value of x is: {}\", x);"; + let range = builder.parse_edit_range(text, 2, 4); + assert_eq!(range.start_line, 2); + assert_eq!(range.start_character, 4); + assert_eq!(range.end_line, 2); + assert_eq!(range.end_character, 4 + text.len() as u32); + + // Test multi-line edit + let text = " let y = 10;\n println!(\"The value of y is: {}\", y);"; + let range = builder.parse_edit_range(text, 2, 4); + assert_eq!(range.start_line, 2); + assert_eq!(range.start_character, 4); + assert_eq!(range.end_line, 3); + assert_eq!( + range.end_character, + " println!(\"The value of y is: {}\", y);".len() as u32 + ); + } +}