Skip to content

Commit 993ba91

Browse files
astandrikastandrik
committed
feat: code completion module
1 parent c67d8e3 commit 993ba91

File tree

10 files changed

+636
-250
lines changed

10 files changed

+636
-250
lines changed

‎src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx‎

Lines changed: 7 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {v4 as uuidv4} from 'uuid';
99
import {CodeAssistantTelemetry} from '../../../../components/CodeAssistantTelemetry/CodeAssistantTelemetry';
1010
import {MonacoEditor} from '../../../../components/MonacoEditor/MonacoEditor';
1111
import SplitPane from '../../../../components/SplitPane';
12+
import {registerCompletionCommands} from '../../../../services/codeCompletion/registerCommands';
1213
import {useTracingLevelOptionAvailable} from '../../../../store/reducers/capabilities/hooks';
1314
import {
1415
goToNextQuery,
@@ -43,7 +44,7 @@ import {
4344
import {useChangedQuerySettings} from '../../../../utils/hooks/useChangedQuerySettings';
4445
import {useLastQueryExecutionSettings} from '../../../../utils/hooks/useLastQueryExecutionSettings';
4546
import {YQL_LANGUAGE_ID} from '../../../../utils/monaco/constats';
46-
import {inlineCompletionProviderInstance} from '../../../../utils/monaco/yql/ydb.inlineCompletionProvider';
47+
import {getCompletionProvider} from '../../../../utils/monaco/yql/ydb.inlineCompletionProvider';
4748
import {QUERY_ACTIONS} from '../../../../utils/query';
4849
import type {InitialPaneState} from '../../utils/paneVisibilityToggleHelpers';
4950
import {
@@ -219,21 +220,12 @@ export default function QueryEditor(props: QueryEditorProps) {
219220
});
220221

221222
if (window.api.codeAssistant) {
222-
monaco.editor.registerCommand('acceptCodeAssistCompletion', (_accessor, ...args) => {
223-
const data = args[0] ?? {};
224-
if (!data || typeof data !== 'object') {
225-
return;
226-
}
227-
const {requestId, suggestionText} = data;
228-
if (requestId && suggestionText) {
229-
inlineCompletionProviderInstance.handleAccept({requestId, suggestionText});
230-
}
231-
});
232-
233-
monaco.editor.registerCommand('declineCodeAssistCompletion', () => {
234-
inlineCompletionProviderInstance.commandDiscard();
235-
});
223+
const provider = getCompletionProvider();
224+
if (provider) {
225+
registerCompletionCommands(editor, monaco, provider);
226+
}
236227
}
228+
237229
initResizeHandler(editor);
238230
initUserPrompt(editor, getLastQueryText);
239231
editor.focus();

‎src/containers/Tenant/Tenant.tsx‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,13 +71,24 @@ export function Tenant(props: TenantProps) {
7171
const previousTenant = React.useRef<string>();
7272
React.useEffect(() => {
7373
if (previousTenant.current !== tenantName) {
74-
const register = async () => {
74+
const registerSuggestCompletion = async () => {
7575
const {registerYQLCompletionItemProvider} = await import(
7676
'../../utils/monaco/yql/yql.completionItemProvider'
7777
);
7878
registerYQLCompletionItemProvider(tenantName);
7979
};
80-
register().catch(console.error);
80+
81+
const registerInlineCompletion = async () => {
82+
const {registerInlineCompletionProvider} = await import(
83+
'../../utils/monaco/yql/ydb.inlineCompletionProvider'
84+
);
85+
if (window.api.codeAssistant) {
86+
registerInlineCompletionProvider(window.api.codeAssistant);
87+
}
88+
};
89+
90+
registerSuggestCompletion().catch(console.error);
91+
registerInlineCompletion().catch(console.error);
8192
previousTenant.current = tenantName;
8293
}
8394
}, [tenantName]);

‎src/services/codeCompletion/CodeCompletionService.ts‎

Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
import * as monaco from 'monaco-editor/esm/vs/editor/editor.api';
2+
3+
import type {DiscardReason} from '../../types/api/codeAssistant';
4+
import {getPromptFileContent} from '../../utils/monaco/codeAssistTelemetry';
5+
6+
import type {
7+
EnrichedCompletion,
8+
ICodeCompletionAPI,
9+
ICodeCompletionService,
10+
ITelemetryService,
11+
InternalSuggestion,
12+
} from './types';
13+
14+
export class CodeCompletionService implements ICodeCompletionService {
15+
private prevSuggestions: InternalSuggestion[] = [];
16+
private timer: number | null = null;
17+
private readonly api: ICodeCompletionAPI;
18+
private readonly telemetry: ITelemetryService;
19+
private readonly editor?: monaco.editor.IStandaloneCodeEditor;
20+
21+
constructor(
22+
api: ICodeCompletionAPI,
23+
telemetry: ITelemetryService,
24+
editor?: monaco.editor.IStandaloneCodeEditor,
25+
) {
26+
this.api = api;
27+
this.telemetry = telemetry;
28+
this.editor = editor;
29+
}
30+
31+
handleItemDidShow(
32+
_completions: monaco.languages.InlineCompletions<EnrichedCompletion>,
33+
item: EnrichedCompletion,
34+
) {
35+
for (const suggests of this.prevSuggestions) {
36+
for (const completion of suggests.items) {
37+
if (completion.pristine === item.pristine) {
38+
suggests.shownCount++;
39+
break;
40+
}
41+
}
42+
}
43+
}
44+
45+
async provideInlineCompletions(
46+
model: monaco.editor.ITextModel,
47+
position: monaco.Position,
48+
_context: monaco.languages.InlineCompletionContext,
49+
_token: monaco.CancellationToken,
50+
) {
51+
const cachedCompletions = this.getCachedCompletion(model, position);
52+
if (cachedCompletions.length) {
53+
return {items: cachedCompletions};
54+
}
55+
while (this.prevSuggestions.length > 0) {
56+
this.dismissCompletion(this.prevSuggestions.pop());
57+
}
58+
const {suggestions, requestId} = await this.getSuggestions(model, position);
59+
60+
this.prevSuggestions = [{items: suggestions, shownCount: 0, requestId}];
61+
return {
62+
items: suggestions,
63+
};
64+
}
65+
66+
handlePartialAccept(
67+
_completions: monaco.languages.InlineCompletions,
68+
item: monaco.languages.InlineCompletion,
69+
acceptedLetters: number,
70+
) {
71+
const {command} = item;
72+
const commandArguments = command?.arguments?.[0] ?? {};
73+
const {suggestionText, requestId, prevWordLength = 0} = commandArguments;
74+
const cachedSuggestions = this.prevSuggestions.find((el) => {
75+
return el.items.some((item) => item.pristine === suggestionText);
76+
});
77+
if (requestId && suggestionText && typeof item.insertText === 'string') {
78+
const acceptedText = item.insertText.slice(prevWordLength, acceptedLetters);
79+
if (acceptedText) {
80+
if (cachedSuggestions) {
81+
cachedSuggestions.wasAccepted = true;
82+
}
83+
this.telemetry.sendAcceptTelemetry(requestId, acceptedText);
84+
}
85+
}
86+
}
87+
88+
handleAccept({requestId, suggestionText}: {requestId: string; suggestionText: string}) {
89+
this.emptyCache();
90+
this.telemetry.sendAcceptTelemetry(requestId, suggestionText);
91+
}
92+
93+
commandDiscard(reason: DiscardReason = 'OnCancel'): void {
94+
while (this.prevSuggestions.length > 0) {
95+
this.discardCompletion(reason, this.prevSuggestions.pop());
96+
}
97+
if (this.editor) {
98+
this.editor.trigger(undefined, 'editor.action.inlineSuggest.hide', undefined);
99+
}
100+
}
101+
102+
emptyCache() {
103+
this.prevSuggestions = [];
104+
}
105+
106+
freeInlineCompletions(): void {
107+
// This method is required by Monaco's InlineCompletionsProvider interface
108+
// but we don't need to do anything here since we handle cleanup in other methods
109+
}
110+
111+
private getCachedCompletion(
112+
model: monaco.editor.ITextModel,
113+
position: monaco.Position,
114+
): EnrichedCompletion[] {
115+
const completions: EnrichedCompletion[] = [];
116+
for (const suggests of this.prevSuggestions) {
117+
for (const completion of suggests.items) {
118+
if (!completion.range) {
119+
continue;
120+
}
121+
if (
122+
position.lineNumber < completion.range.startLineNumber ||
123+
position.column < completion.range.startColumn
124+
) {
125+
continue;
126+
}
127+
const startCompletionPosition = new monaco.Position(
128+
completion.range.startLineNumber,
129+
completion.range.startColumn,
130+
);
131+
const startOffset = model.getOffsetAt(startCompletionPosition);
132+
const endOffset = startOffset + completion.insertText.toString().length;
133+
const positionOffset = model.getOffsetAt(position);
134+
if (positionOffset > endOffset) {
135+
continue;
136+
}
137+
138+
const completionReplaceText = completion.insertText
139+
.toString()
140+
.slice(0, positionOffset - startOffset);
141+
142+
const newRange = new monaco.Range(
143+
completion.range.startLineNumber,
144+
completion.range.startColumn,
145+
position.lineNumber,
146+
position.column,
147+
);
148+
const currentReplaceText = model.getValueInRange(newRange);
149+
if (completionReplaceText.toLowerCase() === currentReplaceText.toLowerCase()) {
150+
completions.push({
151+
insertText:
152+
currentReplaceText +
153+
completion.insertText.toString().slice(positionOffset - startOffset),
154+
range: newRange,
155+
command: completion.command,
156+
pristine: completion.pristine,
157+
});
158+
}
159+
}
160+
}
161+
return completions;
162+
}
163+
164+
private async getSuggestions(model: monaco.editor.ITextModel, position: monaco.Position) {
165+
if (this.timer) {
166+
window.clearTimeout(this.timer);
167+
}
168+
await new Promise((r) => {
169+
this.timer = window.setTimeout(r, 200);
170+
});
171+
let suggestions: EnrichedCompletion[] = [];
172+
let requestId = '';
173+
try {
174+
const data = getPromptFileContent(model, position);
175+
if (!data) {
176+
return {suggestions: []};
177+
}
178+
179+
const codeAssistSuggestions = await this.api.getCodeAssistSuggestions(data);
180+
requestId = codeAssistSuggestions.RequestId;
181+
const {word, startColumn: lastWordStartColumn} = model.getWordUntilPosition(position);
182+
suggestions = codeAssistSuggestions.Suggests.map((el) => {
183+
const suggestionText = el.Text;
184+
const label = word + suggestionText;
185+
return {
186+
label: label,
187+
sortText: 'a',
188+
insertText: label,
189+
pristine: suggestionText,
190+
range: new monaco.Range(
191+
position.lineNumber,
192+
lastWordStartColumn,
193+
position.lineNumber,
194+
position.column,
195+
),
196+
command: {
197+
id: 'acceptCodeAssistCompletion',
198+
title: '',
199+
arguments: [
200+
{
201+
requestId,
202+
suggestionText: suggestionText,
203+
prevWordLength: word.length,
204+
},
205+
],
206+
},
207+
};
208+
});
209+
} catch (err) {}
210+
return {suggestions, requestId};
211+
}
212+
213+
private discardCompletion(reason: DiscardReason, completion?: InternalSuggestion): void {
214+
if (completion === undefined) {
215+
return;
216+
}
217+
const {requestId, items, shownCount} = completion;
218+
if (!requestId || !items.length) {
219+
return;
220+
}
221+
for (const item of items) {
222+
this.telemetry.sendDeclineTelemetry(requestId, item.pristine, reason, shownCount);
223+
}
224+
}
225+
226+
private dismissCompletion(completion?: InternalSuggestion): void {
227+
if (completion === undefined) {
228+
return;
229+
}
230+
const {requestId, items, shownCount, wasAccepted} = completion;
231+
232+
if (!requestId || !items.length || !shownCount || wasAccepted) {
233+
return;
234+
}
235+
for (const item of items) {
236+
this.telemetry.sendIgnoreTelemetry(requestId, item.pristine);
237+
}
238+
}
239+
}

0 commit comments

Comments
 (0)