Skip to content

Commit a10a5cd

Browse files
astandrikastandrik
authored
feat: code assistant integration (#1902)
1 parent cc81623 commit a10a5cd

File tree

13 files changed

+450
-12
lines changed

13 files changed

+450
-12
lines changed

‎package-lock.json

+30-1
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json

+1
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
"@hookform/resolvers": "^3.10.0",
2929
"@reduxjs/toolkit": "^2.5.0",
3030
"@tanstack/react-table": "^8.20.6",
31+
"@ydb-platform/monaco-ghost": "^0.4.0",
3132
"axios": "^1.7.9",
3233
"axios-retry": "^4.5.0",
3334
"colord": "^2.9.3",

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

+33-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import NiceModal from '@ebay/nice-modal-react';
2+
import {useMonacoGhost} from '@ydb-platform/monaco-ghost';
23
import throttle from 'lodash/throttle';
34
import type Monaco from 'monaco-editor';
45

@@ -22,8 +23,9 @@ import {useUpdateErrorsHighlighting} from '../../../../utils/monaco/highlightErr
2223
import {QUERY_ACTIONS} from '../../../../utils/query';
2324
import {SAVE_QUERY_DIALOG} from '../SaveQuery/SaveQuery';
2425
import i18n from '../i18n';
26+
import {useSavedQueries} from '../utils/useSavedQueries';
2527

26-
import {useEditorOptions} from './helpers';
28+
import {useCodeAssist, useEditorOptions} from './helpers';
2729
import {getKeyBindings} from './keybindings';
2830

2931
const CONTEXT_MENU_GROUP_ID = 'navigation';
@@ -44,6 +46,7 @@ export function YqlEditor({
4446
const input = useTypedSelector(selectUserInput);
4547
const dispatch = useTypedDispatch();
4648
const historyQueries = useTypedSelector(selectQueriesHistory);
49+
const savedQueries = useSavedQueries();
4750
const editorOptions = useEditorOptions();
4851
const updateErrorsHighlighting = useUpdateErrorsHighlighting();
4952

@@ -68,6 +71,21 @@ export function YqlEditor({
6871
window.ydbEditor = undefined;
6972
};
7073

74+
const codeAssist = useCodeAssist();
75+
const {registerMonacoGhost} = useMonacoGhost({
76+
api: {
77+
getCodeAssistSuggestions: codeAssist.getCodeAssistSuggestions,
78+
},
79+
eventHandlers: {
80+
onCompletionAccept: codeAssist.onCompletionAccept,
81+
onCompletionDecline: codeAssist.onCompletionDecline,
82+
onCompletionIgnore: codeAssist.onCompletionIgnore,
83+
},
84+
config: {
85+
language: YQL_LANGUAGE_ID,
86+
},
87+
});
88+
7189
const editorDidMount = (editor: Monaco.editor.IStandaloneCodeEditor, monaco: typeof Monaco) => {
7290
window.ydbEditor = editor;
7391
const keybindings = getKeyBindings(monaco);
@@ -80,6 +98,20 @@ export function YqlEditor({
8098
contribution.insert(input);
8199
}
82100
});
101+
102+
if (window.api.codeAssist) {
103+
registerMonacoGhost(editor);
104+
codeAssist.prepareUserQueriesCache([
105+
...historyQueries.map((query, index) => ({
106+
name: `query${index}.yql`,
107+
text: query.queryText,
108+
})),
109+
...savedQueries.map((query) => ({
110+
name: query.name,
111+
text: query.body,
112+
})),
113+
]);
114+
}
83115
initResizeHandler(editor);
84116
initUserPrompt(editor, getLastQueryText);
85117
editor.focus();

‎src/containers/Tenant/Query/QueryEditor/helpers.ts

+54
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import React from 'react';
22

3+
import type {AcceptEvent, DeclineEvent, IgnoreEvent, PromptFile} from '@ydb-platform/monaco-ghost';
34
import type Monaco from 'monaco-editor';
45

6+
import {codeAssistApi} from '../../../../store/reducers/codeAssist/codeAssist';
7+
import type {TelemetryOpenTabs} from '../../../../types/api/codeAssist';
58
import {AUTOCOMPLETE_ON_ENTER, ENABLE_AUTOCOMPLETE} from '../../../../utils/constants';
69
import {useSetting} from '../../../../utils/hooks';
710

@@ -32,3 +35,54 @@ export function useEditorOptions() {
3235

3336
return options;
3437
}
38+
39+
export function useCodeAssist() {
40+
const [sendCodeAssistPrompt] = codeAssistApi.useLazyGetCodeAssistSuggestionsQuery();
41+
const [acceptSuggestion] = codeAssistApi.useAcceptSuggestionMutation();
42+
const [discardSuggestion] = codeAssistApi.useDiscardSuggestionMutation();
43+
const [ignoreSuggestion] = codeAssistApi.useIgnoreSuggestionMutation();
44+
const [sendUserQueriesData] = codeAssistApi.useSendUserQueriesDataMutation();
45+
46+
const getCodeAssistSuggestions = React.useCallback(
47+
async (promptFiles: PromptFile[]) => sendCodeAssistPrompt(promptFiles).unwrap(),
48+
[sendCodeAssistPrompt],
49+
);
50+
51+
const onCompletionAccept = React.useCallback(
52+
async (event: AcceptEvent) => acceptSuggestion(event).unwrap(),
53+
[acceptSuggestion],
54+
);
55+
56+
const onCompletionDecline = React.useCallback(
57+
async (event: DeclineEvent) => discardSuggestion(event).unwrap(),
58+
[discardSuggestion],
59+
);
60+
61+
const onCompletionIgnore = React.useCallback(
62+
async (event: IgnoreEvent) => ignoreSuggestion(event).unwrap(),
63+
[ignoreSuggestion],
64+
);
65+
66+
const prepareUserQueriesCache = React.useCallback(
67+
async (queries: {text: string; name?: string}[]) => {
68+
const preparedData: TelemetryOpenTabs = queries.map((query, index) => ({
69+
FileName: query.name || `query${index}.yql`,
70+
Text: query.text,
71+
}));
72+
try {
73+
return sendUserQueriesData(preparedData).unwrap();
74+
} catch {
75+
return {items: []};
76+
}
77+
},
78+
[sendUserQueriesData],
79+
);
80+
81+
return {
82+
getCodeAssistSuggestions,
83+
onCompletionAccept,
84+
onCompletionDecline,
85+
onCompletionIgnore,
86+
prepareUserQueriesCache,
87+
};
88+
}

‎src/services/api/codeAssist.ts

+130
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
// IMPORTANT!
2+
// In the future code assist api will be provided to ydb-ui component explicitly by consumer service.
3+
// Current solution is temporary and aimed to satisfy internal puproses.
4+
// It means this whole file will be moved to customer service
5+
6+
import type {PromptFile, Suggestions} from '@ydb-platform/monaco-ghost';
7+
8+
import {codeAssistBackend as CODE_ASSISTANT_BACKEND} from '../../store';
9+
import type {
10+
CodeAssistSuggestionsFiles,
11+
CodeAssistSuggestionsResponse,
12+
TelemetryEvent,
13+
TelemetryOpenTabs,
14+
} from '../../types/api/codeAssist';
15+
16+
import {BaseYdbAPI} from './base';
17+
const ideInfo = {
18+
Ide: 'ydb',
19+
IdeVersion: '1',
20+
PluginFamily: 'ydb',
21+
PluginVersion: '0.2',
22+
};
23+
24+
const limitForTab = 10_000;
25+
const limitBeforeCursor = 8_000;
26+
const limitAfterCursor = 1_000;
27+
28+
function prepareCodeAssistTabs(tabs: TelemetryOpenTabs): TelemetryOpenTabs {
29+
return tabs.map((tab) => {
30+
const text = tab.Text;
31+
if (text.length > limitForTab) {
32+
return {
33+
...tab,
34+
Text: text.slice(0, limitForTab),
35+
};
36+
}
37+
38+
return tab;
39+
});
40+
}
41+
42+
function prepareCodeAssistPrompt(promptFiles: PromptFile[]): CodeAssistSuggestionsFiles {
43+
return promptFiles.map((file) => {
44+
const cursorLine = file.cursorPosition.lineNumber;
45+
const cursorCol = file.cursorPosition.column;
46+
47+
return {
48+
Fragments: file.fragments.map((fragment) => {
49+
let text = fragment.text;
50+
const isBeforeCursor =
51+
fragment.end.lineNumber < cursorLine ||
52+
(fragment.end.lineNumber === cursorLine && fragment.end.column <= cursorCol);
53+
const isAfterCursor =
54+
fragment.start.lineNumber > cursorLine ||
55+
(fragment.start.lineNumber === cursorLine && fragment.start.column > cursorCol);
56+
57+
if (isBeforeCursor) {
58+
text = text.slice(-limitBeforeCursor);
59+
} else if (isAfterCursor) {
60+
text = text.slice(0, limitAfterCursor);
61+
}
62+
63+
return {
64+
Text: text,
65+
Start: {
66+
Ln: fragment.start.lineNumber,
67+
Col: fragment.start.column,
68+
},
69+
End: {
70+
Ln: fragment.end.lineNumber,
71+
Col: fragment.end.column,
72+
},
73+
};
74+
}),
75+
Cursor: {
76+
Ln: cursorLine,
77+
Col: cursorCol,
78+
},
79+
Path: `${file.path}.yql`,
80+
};
81+
});
82+
}
83+
84+
export class CodeAssistAPI extends BaseYdbAPI {
85+
getPath(path: string) {
86+
return `${CODE_ASSISTANT_BACKEND ?? ''}${path}`;
87+
}
88+
89+
async getCodeAssistSuggestions(data: PromptFile[]): Promise<Suggestions> {
90+
const request: CodeAssistSuggestionsFiles = prepareCodeAssistPrompt(data);
91+
92+
const response = await this.post<CodeAssistSuggestionsResponse>(
93+
this.getPath('/code-assist-suggestion'),
94+
{
95+
Files: request,
96+
ContextCreateType: 1,
97+
IdeInfo: ideInfo,
98+
},
99+
null,
100+
{
101+
concurrentId: 'code-assist-suggestion',
102+
collectRequest: false,
103+
},
104+
);
105+
106+
return {
107+
items: response.Suggests.map((suggestion) => suggestion.Text),
108+
requestId: response.RequestId,
109+
};
110+
}
111+
112+
sendCodeAssistTelemetry(data: TelemetryEvent) {
113+
return this.post('/code-assist-telemetry', data, null, {
114+
concurrentId: 'code-assist-telemetry',
115+
collectRequest: true,
116+
});
117+
}
118+
119+
sendCodeAssistOpenTabs(data: TelemetryOpenTabs) {
120+
return this.post(
121+
'/code-assist-telemetry',
122+
{OpenTabs: {Tabs: prepareCodeAssistTabs(data), IdeInfo: ideInfo}},
123+
null,
124+
{
125+
concurrentId: 'code-assist-telemetry',
126+
collectRequest: false,
127+
},
128+
);
129+
}
130+
}

‎src/services/api/index.ts

+3
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import type {AxiosRequestConfig} from 'axios';
22

33
import {AuthAPI} from './auth';
4+
import {CodeAssistAPI} from './codeAssist';
45
import {MetaAPI} from './meta';
56
import {OperationAPI} from './operation';
67
import {PDiskAPI} from './pdisk';
@@ -20,11 +21,13 @@ export class YdbEmbeddedAPI {
2021
vdisk: VDiskAPI;
2122
viewer: ViewerAPI;
2223
meta?: MetaAPI;
24+
codeAssist?: CodeAssistAPI;
2325

2426
constructor({config, webVersion}: {config: AxiosRequestConfig; webVersion?: boolean}) {
2527
this.auth = new AuthAPI({config});
2628
if (webVersion) {
2729
this.meta = new MetaAPI({config});
30+
this.codeAssist = new CodeAssistAPI({config});
2831
}
2932
this.operation = new OperationAPI({config});
3033
this.pdisk = new PDiskAPI({config});

‎src/store/configureStore.ts

+1
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ function _configureStore<
4646
export const webVersion = window.web_version;
4747
export const customBackend = window.custom_backend;
4848
export const metaBackend = window.meta_backend;
49+
export const codeAssistBackend = window.code_assist_backend;
4950

5051
const isSingleClusterMode = `${metaBackend}` === 'undefined';
5152

‎src/store/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ export {
55
configureStore,
66
customBackend,
77
metaBackend,
8+
codeAssistBackend,
89
webVersion,
910
} from './configureStore';
1011
export {rootReducer} from './reducers';

0 commit comments

Comments
 (0)