feat: code assistant integration

package.json

@@ -28,6 +28,7 @@
     "@hookform/resolvers": "^3.10.0",
     "@reduxjs/toolkit": "^2.5.0",
     "@tanstack/react-table": "^8.20.6",
+    "@ydb-platform/monaco-ghost": "^0.4.0",
     "axios": "^1.7.9",
     "axios-retry": "^4.5.0",
     "colord": "^2.9.3",

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

@@ -1,4 +1,5 @@
 import NiceModal from '@ebay/nice-modal-react';
+import {useMonacoGhost} from '@ydb-platform/monaco-ghost';
 import throttle from 'lodash/throttle';
 import type Monaco from 'monaco-editor';
 
@@ -22,8 +23,9 @@ import {useUpdateErrorsHighlighting} from '../../../../utils/monaco/highlightErr
 import {QUERY_ACTIONS} from '../../../../utils/query';
 import {SAVE_QUERY_DIALOG} from '../SaveQuery/SaveQuery';
 import i18n from '../i18n';
+import {useSavedQueries} from '../utils/useSavedQueries';
 
-import {useEditorOptions} from './helpers';
+import {useCodeAssist, useEditorOptions} from './helpers';
 import {getKeyBindings} from './keybindings';
 
 const CONTEXT_MENU_GROUP_ID = 'navigation';
@@ -44,6 +46,7 @@ export function YqlEditor({
     const input = useTypedSelector(selectUserInput);
     const dispatch = useTypedDispatch();
     const historyQueries = useTypedSelector(selectQueriesHistory);
+    const savedQueries = useSavedQueries();
     const editorOptions = useEditorOptions();
     const updateErrorsHighlighting = useUpdateErrorsHighlighting();
 
@@ -68,6 +71,21 @@ export function YqlEditor({
         window.ydbEditor = undefined;
     };
 
+    const codeAssist = useCodeAssist();
+    const {registerMonacoGhost} = useMonacoGhost({
+        api: {
+            getCodeAssistSuggestions: codeAssist.getCodeAssistSuggestions,
+        },
+        eventHandlers: {
+            onCompletionAccept: codeAssist.onCompletionAccept,
+            onCompletionDecline: codeAssist.onCompletionDecline,
+            onCompletionIgnore: codeAssist.onCompletionIgnore,
+        },
+        config: {
+            language: YQL_LANGUAGE_ID,
+        },
+    });
+
     const editorDidMount = (editor: Monaco.editor.IStandaloneCodeEditor, monaco: typeof Monaco) => {
         window.ydbEditor = editor;
         const keybindings = getKeyBindings(monaco);
@@ -80,6 +98,20 @@ export function YqlEditor({
                 contribution.insert(input);
             }
         });
+
+        if (window.api.codeAssist) {
+            registerMonacoGhost(editor);
+            codeAssist.prepareUserQueriesCache([
+                ...historyQueries.map((query, index) => ({
+                    name: `query${index}.yql`,
+                    text: query.queryText,
+                })),
+                ...savedQueries.map((query) => ({
+                    name: query.name,
+                    text: query.body,
+                })),
+            ]);
+        }
         initResizeHandler(editor);
         initUserPrompt(editor, getLastQueryText);
         editor.focus();

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

@@ -1,7 +1,10 @@
 import React from 'react';
 
+import type {AcceptEvent, DeclineEvent, IgnoreEvent, PromptFile} from '@ydb-platform/monaco-ghost';
 import type Monaco from 'monaco-editor';
 
+import {codeAssistApi} from '../../../../store/reducers/codeAssist/codeAssist';
+import type {TelemetryOpenTabs} from '../../../../types/api/codeAssist';
 import {AUTOCOMPLETE_ON_ENTER, ENABLE_AUTOCOMPLETE} from '../../../../utils/constants';
 import {useSetting} from '../../../../utils/hooks';
 
@@ -32,3 +35,54 @@ export function useEditorOptions() {
 
     return options;
 }
+
+export function useCodeAssist() {
+    const [sendCodeAssistPrompt] = codeAssistApi.useLazyGetCodeAssistSuggestionsQuery();
+    const [acceptSuggestion] = codeAssistApi.useAcceptSuggestionMutation();
+    const [discardSuggestion] = codeAssistApi.useDiscardSuggestionMutation();
+    const [ignoreSuggestion] = codeAssistApi.useIgnoreSuggestionMutation();
+    const [sendUserQueriesData] = codeAssistApi.useSendUserQueriesDataMutation();
+
+    const getCodeAssistSuggestions = React.useCallback(
+        async (promptFiles: PromptFile[]) => sendCodeAssistPrompt(promptFiles).unwrap(),
+        [sendCodeAssistPrompt],
+    );
+
+    const onCompletionAccept = React.useCallback(
+        async (event: AcceptEvent) => acceptSuggestion(event).unwrap(),
+        [acceptSuggestion],
+    );
+
+    const onCompletionDecline = React.useCallback(
+        async (event: DeclineEvent) => discardSuggestion(event).unwrap(),
+        [discardSuggestion],
+    );
+
+    const onCompletionIgnore = React.useCallback(
+        async (event: IgnoreEvent) => ignoreSuggestion(event).unwrap(),
+        [ignoreSuggestion],
+    );
+
+    const prepareUserQueriesCache = React.useCallback(
+        async (queries: {text: string; name?: string}[]) => {
+            const preparedData: TelemetryOpenTabs = queries.map((query, index) => ({
+                FileName: query.name || `query${index}.yql`,
+                Text: query.text,
+            }));
+            try {
+                return sendUserQueriesData(preparedData).unwrap();
+            } catch {
+                return {items: []};
+            }
+        },
+        [sendUserQueriesData],
+    );
+
+    return {
+        getCodeAssistSuggestions,
+        onCompletionAccept,
+        onCompletionDecline,
+        onCompletionIgnore,
+        prepareUserQueriesCache,
+    };
+}

src/services/api/codeAssist.ts

@@ -0,0 +1,130 @@
+// IMPORTANT!
+// In the future code assist api will be provided to ydb-ui component explicitly by consumer service.
+// Current solution is temporary and aimed to satisfy internal puproses.
+// It means this whole file will be moved to customer service
+
+import type {PromptFile, Suggestions} from '@ydb-platform/monaco-ghost';
+
+import {codeAssistBackend as CODE_ASSISTANT_BACKEND} from '../../store';
+import type {
+    CodeAssistSuggestionsFiles,
+    CodeAssistSuggestionsResponse,
+    TelemetryEvent,
+    TelemetryOpenTabs,
+} from '../../types/api/codeAssist';
+
+import {BaseYdbAPI} from './base';
+const ideInfo = {
+    Ide: 'ydb',
+    IdeVersion: '1',
+    PluginFamily: 'ydb',
+    PluginVersion: '0.2',
+};
+
+const limitForTab = 10_000;
+const limitBeforeCursor = 8_000;
+const limitAfterCursor = 1_000;
+
+function prepareCodeAssistTabs(tabs: TelemetryOpenTabs): TelemetryOpenTabs {
+    return tabs.map((tab) => {
+        const text = tab.Text;
+        if (text.length > limitForTab) {
+            return {
+                ...tab,
+                Text: text.slice(0, limitForTab),
+            };
+        }
+
+        return tab;
+    });
+}
+
+function prepareCodeAssistPrompt(promptFiles: PromptFile[]): CodeAssistSuggestionsFiles {
+    return promptFiles.map((file) => {
+        const cursorLine = file.cursorPosition.lineNumber;
+        const cursorCol = file.cursorPosition.column;
+
+        return {
+            Fragments: file.fragments.map((fragment) => {
+                let text = fragment.text;
+                const isBeforeCursor =
+                    fragment.end.lineNumber < cursorLine ||
+                    (fragment.end.lineNumber === cursorLine && fragment.end.column <= cursorCol);
+                const isAfterCursor =
+                    fragment.start.lineNumber > cursorLine ||
+                    (fragment.start.lineNumber === cursorLine && fragment.start.column > cursorCol);
+
+                if (isBeforeCursor) {
+                    text = text.slice(-limitBeforeCursor);
+                } else if (isAfterCursor) {
+                    text = text.slice(0, limitAfterCursor);
+                }
+
+                return {
+                    Text: text,
+                    Start: {
+                        Ln: fragment.start.lineNumber,
+                        Col: fragment.start.column,
+                    },
+                    End: {
+                        Ln: fragment.end.lineNumber,
+                        Col: fragment.end.column,
+                    },
+                };
+            }),
+            Cursor: {
+                Ln: cursorLine,
+                Col: cursorCol,
+            },
+            Path: `${file.path}.yql`,
+        };
+    });
+}
+
+export class CodeAssistAPI extends BaseYdbAPI {
+    getPath(path: string) {
+        return `${CODE_ASSISTANT_BACKEND ?? ''}${path}`;
+    }
+
+    async getCodeAssistSuggestions(data: PromptFile[]): Promise<Suggestions> {
+        const request: CodeAssistSuggestionsFiles = prepareCodeAssistPrompt(data);
+
+        const response = await this.post<CodeAssistSuggestionsResponse>(
+            this.getPath('/code-assist-suggestion'),
+            {
+                Files: request,
+                ContextCreateType: 1,
+                IdeInfo: ideInfo,
+            },
+            null,
+            {
+                concurrentId: 'code-assist-suggestion',
+                collectRequest: false,
+            },
+        );
+
+        return {
+            items: response.Suggests.map((suggestion) => suggestion.Text),
+            requestId: response.RequestId,
+        };
+    }
+
+    sendCodeAssistTelemetry(data: TelemetryEvent) {
+        return this.post('/code-assist-telemetry', data, null, {
+            concurrentId: 'code-assist-telemetry',
+            collectRequest: true,
+        });
+    }
+
+    sendCodeAssistOpenTabs(data: TelemetryOpenTabs) {
+        return this.post(
+            '/code-assist-telemetry',
+            {OpenTabs: {Tabs: prepareCodeAssistTabs(data), IdeInfo: ideInfo}},
+            null,
+            {
+                concurrentId: 'code-assist-telemetry',
+                collectRequest: false,
+            },
+        );
+    }
+}

src/services/api/index.ts

@@ -1,6 +1,7 @@
 import type {AxiosRequestConfig} from 'axios';
 
 import {AuthAPI} from './auth';
+import {CodeAssistAPI} from './codeAssist';
 import {MetaAPI} from './meta';
 import {OperationAPI} from './operation';
 import {PDiskAPI} from './pdisk';
@@ -22,11 +23,13 @@ export class YdbEmbeddedAPI {
     vdisk: VDiskAPI;
     viewer: ViewerAPI;
     meta?: MetaAPI;
+    codeAssist?: CodeAssistAPI;
 
     constructor({config, webVersion}: {config: AxiosRequestConfig; webVersion?: boolean}) {
         this.auth = new AuthAPI({config});
         if (webVersion) {
             this.meta = new MetaAPI({config});
+            this.codeAssist = new CodeAssistAPI({config});
         }
         this.operation = new OperationAPI({config});
         this.pdisk = new PDiskAPI({config});

src/store/configureStore.ts

@@ -46,6 +46,7 @@ function _configureStore<
 export const webVersion = window.web_version;
 export const customBackend = window.custom_backend;
 export const metaBackend = window.meta_backend;
+export const codeAssistBackend = window.code_assist_backend;
 
 const isSingleClusterMode = `${metaBackend}` === 'undefined';
 

src/store/index.ts

@@ -5,6 +5,7 @@ export {
     configureStore,
     customBackend,
     metaBackend,
+    codeAssistBackend,
     webVersion,
 } from './configureStore';
 export {rootReducer} from './reducers';