From 48b6e0645156c60a954bd6dca0d660cb54c849cf Mon Sep 17 00:00:00 2001 From: Jonas Helming Date: Wed, 27 Nov 2024 00:49:47 +0100 Subject: [PATCH] Show progress while calculating AI code completion fixed #14536 Signed-off-by: Jonas Helming --- .../src/common/code-completion-agent.ts | 156 ++++++++++-------- 1 file changed, 83 insertions(+), 73 deletions(-) diff --git a/packages/ai-code-completion/src/common/code-completion-agent.ts b/packages/ai-code-completion/src/common/code-completion-agent.ts index c47f5b6089957..e081835d2ff5e 100644 --- a/packages/ai-code-completion/src/common/code-completion-agent.ts +++ b/packages/ai-code-completion/src/common/code-completion-agent.ts @@ -18,7 +18,7 @@ import { Agent, AgentSpecificVariables, CommunicationRecordingService, getTextOfResponse, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequirement, PromptService, PromptTemplate } from '@theia/ai-core/lib/common'; -import { generateUuid, ILogger } from '@theia/core'; +import { generateUuid, ILogger, ProgressService } from '@theia/core'; import { inject, injectable, named } from '@theia/core/shared/inversify'; import * as monaco from '@theia/monaco-editor-core'; @@ -36,82 +36,89 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent { context: monaco.languages.InlineCompletionContext, token: monaco.CancellationToken ): Promise { - const languageModel = - await this.languageModelRegistry.selectLanguageModel({ - agent: this.id, - ...this.languageModelRequirements[0], + const progress = await this.progressService.showProgress( + { text: 'Calculating AI code completion...', options: { location: 'window' } } + ); + try { + const languageModel = + await this.languageModelRegistry.selectLanguageModel({ + agent: this.id, + ...this.languageModelRequirements[0], + }); + if (!languageModel) { + this.logger.error( + 'No language model found for code-completion-agent' + ); + return undefined; + } + + // Get text until the given position + const prefix = model.getValueInRange({ + startLineNumber: 1, + startColumn: 1, + endLineNumber: position.lineNumber, + endColumn: position.column, }); - if (!languageModel) { - this.logger.error( - 'No language model found for code-completion-agent' - ); - return undefined; - } - // Get text until the given position - const prefix = model.getValueInRange({ - startLineNumber: 1, - startColumn: 1, - endLineNumber: position.lineNumber, - endColumn: position.column, - }); - - // Get text after the given position - const suffix = model.getValueInRange({ - startLineNumber: position.lineNumber, - startColumn: position.column, - endLineNumber: model.getLineCount(), - endColumn: model.getLineMaxColumn(model.getLineCount()), - }); - - const file = model.uri.toString(false); - const language = model.getLanguageId(); - - if (token.isCancellationRequested) { - return undefined; - } - const prompt = await this.promptService - .getPrompt('code-completion-prompt', { prefix, suffix, file, language }) - .then(p => p?.text); - if (!prompt) { - this.logger.error('No prompt found for code-completion-agent'); - return undefined; - } + // Get text after the given position + const suffix = model.getValueInRange({ + startLineNumber: position.lineNumber, + startColumn: position.column, + endLineNumber: model.getLineCount(), + endColumn: model.getLineMaxColumn(model.getLineCount()), + }); - // since we do not actually hold complete conversions, the request/response pair is considered a session - const sessionId = generateUuid(); - const requestId = generateUuid(); - const request: LanguageModelRequest = { - messages: [{ type: 'text', actor: 'user', query: prompt }], - }; - if (token.isCancellationRequested) { - return undefined; - } - this.recordingService.recordRequest({ - agentId: this.id, - sessionId, - requestId, - request: prompt, - }); - const response = await languageModel.request(request, token); - if (token.isCancellationRequested) { - return undefined; - } - const completionText = await getTextOfResponse(response); - if (token.isCancellationRequested) { - return undefined; + const file = model.uri.toString(false); + const language = model.getLanguageId(); + + if (token.isCancellationRequested) { + return undefined; + } + const prompt = await this.promptService + .getPrompt('code-completion-prompt', { prefix, suffix, file, language }) + .then(p => p?.text); + if (!prompt) { + this.logger.error('No prompt found for code-completion-agent'); + return undefined; + } + + // since we do not actually hold complete conversions, the request/response pair is considered a session + const sessionId = generateUuid(); + const requestId = generateUuid(); + const request: LanguageModelRequest = { + messages: [{ type: 'text', actor: 'user', query: prompt }], + }; + if (token.isCancellationRequested) { + return undefined; + } + this.recordingService.recordRequest({ + agentId: this.id, + sessionId, + requestId, + request: prompt, + }); + const response = await languageModel.request(request, token); + if (token.isCancellationRequested) { + return undefined; + } + const completionText = await getTextOfResponse(response); + if (token.isCancellationRequested) { + return undefined; + } + this.recordingService.recordResponse({ + agentId: this.id, + sessionId, + requestId, + response: completionText, + }); + + return { + items: [{ insertText: completionText }], + enableForwardStability: true, + }; + } finally { + progress.cancel(); } - this.recordingService.recordResponse({ - agentId: this.id, - sessionId, - requestId, - response: completionText, - }); - - return { - items: [{ insertText: completionText }], - enableForwardStability: true, - }; } @inject(ILogger) @@ -127,6 +134,9 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent { @inject(CommunicationRecordingService) protected recordingService: CommunicationRecordingService; + @inject(ProgressService) + protected progressService: ProgressService; + id = 'Code Completion'; name = 'Code Completion'; description =