Allow specifying max lines used for AI code completion context (#14539)

* Allow specifying max lines used for AI code completion context fixed #14538 Signed-off-by: Jonas Helming <[email protected]>
eclipse-theia · Nov 28, 2024 · 8ad410c · 8ad410c
1 parent ab92800
commit 8ad410c
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 7 deletions.
diff --git a/packages/ai-code-completion/src/browser/ai-code-completion-frontend-module.ts b/packages/ai-code-completion/src/browser/ai-code-completion-frontend-module.ts
@@ -16,7 +16,7 @@
 
 import { ILogger } from '@theia/core';
 import { ContainerModule } from '@theia/core/shared/inversify';
-import { CodeCompletionAgent, CodeCompletionAgentImpl } from '../common/code-completion-agent';
+import { CodeCompletionAgent, CodeCompletionAgentImpl } from './code-completion-agent';
 import { AIFrontendApplicationContribution } from './ai-code-frontend-application-contribution';
 import { FrontendApplicationContribution, KeybindingContribution, PreferenceContribution } from '@theia/core/lib/browser';
 import { Agent } from '@theia/ai-core';

diff --git a/packages/ai-code-completion/src/browser/ai-code-completion-preference.ts b/packages/ai-code-completion/src/browser/ai-code-completion-preference.ts
@@ -19,6 +19,7 @@ import { AI_CORE_PREFERENCES_TITLE } from '@theia/ai-core/lib/browser/ai-core-pr
 
 export const PREF_AI_INLINE_COMPLETION_AUTOMATIC_ENABLE = 'ai-features.codeCompletion.automaticCodeCompletion';
 export const PREF_AI_INLINE_COMPLETION_EXCLUDED_EXTENSIONS = 'ai-features.codeCompletion.excludedFileExtensions';
+export const PREF_AI_INLINE_COMPLETION_MAX_CONTEXT_LINES = 'ai-features.codeCompletion.maxContextLines';
 
 export const AICodeCompletionPreferencesSchema: PreferenceSchema = {
     type: 'object',
@@ -39,6 +40,14 @@ export const AICodeCompletionPreferencesSchema: PreferenceSchema = {
                 type: 'string'
             },
             default: []
+        },
+        [PREF_AI_INLINE_COMPLETION_MAX_CONTEXT_LINES]: {
+            title: 'Maximum Context Lines',
+            type: 'number',
+            description: 'The maximum number of lines used as context, distributed among the lines before and after the cursor position (prefix and suffix).\
+            Set this to -1 to use the full file as context without any line limit and 0 to only use the current line.',
+            default: -1,
+            minimum: -1
         }
     }
 };
diff --git a/packages/ai-code-completion/src/browser/ai-code-inline-completion-provider.ts b/packages/ai-code-completion/src/browser/ai-code-inline-completion-provider.ts
@@ -17,7 +17,7 @@
 import * as monaco from '@theia/monaco-editor-core';
 
 import { inject, injectable } from '@theia/core/shared/inversify';
-import { CodeCompletionAgent } from '../common/code-completion-agent';
+import { CodeCompletionAgent } from './code-completion-agent';
 import { AgentService } from '@theia/ai-core';
 
 @injectable()

diff --git a/...etion/src/common/code-completion-agent.ts → ...tion/src/browser/code-completion-agent.ts b/...etion/src/common/code-completion-agent.ts → ...tion/src/browser/code-completion-agent.ts
@@ -21,6 +21,8 @@ import {
 import { generateUuid, ILogger, ProgressService } from '@theia/core';
 import { inject, injectable, named } from '@theia/core/shared/inversify';
 import * as monaco from '@theia/monaco-editor-core';
+import { PREF_AI_INLINE_COMPLETION_MAX_CONTEXT_LINES } from './ai-code-completion-preference';
+import { PreferenceService } from '@theia/core/lib/browser';
 
 export const CodeCompletionAgent = Symbol('CodeCompletionAgent');
 export interface CodeCompletionAgent extends Agent {
@@ -52,20 +54,46 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent {
                 return undefined;
             }
 
-            // Get text until the given position
+            const maxContextLines = this.preferences.get<number>(PREF_AI_INLINE_COMPLETION_MAX_CONTEXT_LINES, -1);
+
+            let prefixStartLine = 1;
+            let suffixEndLine = model.getLineCount();
+            // if maxContextLines is -1, use the full file as context without any line limit
+
+            if (maxContextLines === 0) {
+                // Only the cursor line
+                prefixStartLine = position.lineNumber;
+                suffixEndLine = position.lineNumber;
+            } else if (maxContextLines > 0) {
+                const linesBeforeCursor = position.lineNumber - 1;
+                const linesAfterCursor = model.getLineCount() - position.lineNumber;
+
+                // Allocate one more line to the prefix in case of an odd maxContextLines
+                const prefixLines = Math.min(
+                    Math.ceil(maxContextLines / 2),
+                    linesBeforeCursor
+                );
+                const suffixLines = Math.min(
+                    Math.floor(maxContextLines / 2),
+                    linesAfterCursor
+                );
+
+                prefixStartLine = Math.max(1, position.lineNumber - prefixLines);
+                suffixEndLine = Math.min(model.getLineCount(), position.lineNumber + suffixLines);
+            }
+
             const prefix = model.getValueInRange({
-                startLineNumber: 1,
+                startLineNumber: prefixStartLine,
                 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()),
+                endLineNumber: suffixEndLine,
+                endColumn: model.getLineMaxColumn(suffixEndLine),
             });
 
             const file = model.uri.toString(false);
@@ -137,6 +165,9 @@ export class CodeCompletionAgentImpl implements CodeCompletionAgent {
     @inject(ProgressService)
     protected progressService: ProgressService;
 
+    @inject(PreferenceService)
+    protected preferences: PreferenceService;
+
     id = 'Code Completion';
     name = 'Code Completion';
     description =