feat: code assistant

astandrik · astandrik · commit a24893d9936a · 2024-11-28T14:20:40.000+03:00
diff --git a/public/index.html b/public/index.html
@@ -25,6 +25,7 @@
       window.web_version = !'%REACT_APP_BACKEND%';
       window.custom_backend = '%NODE_ENV%' === 'development' && '%REACT_APP_BACKEND%';
       window.meta_backend = '%REACT_APP_META_BACKEND%'
+      window.code_assistant_backend = '%REACT_APP_CODE_ASSISTANT_BACKEND%'
     </script>
   </head>
   <body>
diff --git a/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx b/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx
@@ -42,6 +42,7 @@ import {
 import {useChangedQuerySettings} from '../../../../utils/hooks/useChangedQuerySettings';
 import {useLastQueryExecutionSettings} from '../../../../utils/hooks/useLastQueryExecutionSettings';
 import {YQL_LANGUAGE_ID} from '../../../../utils/monaco/constats';
+import {inlineCompletionProviderInstance} from '../../../../utils/monaco/yql/ydb.inlineCompletionProvider';
 import {QUERY_ACTIONS} from '../../../../utils/query';
 import type {InitialPaneState} from '../../utils/paneVisibilityToggleHelpers';
 import {
@@ -215,6 +216,23 @@ export default function QueryEditor(props: QueryEditorProps) {
                 contribution.insert(input);
             }
         });
+
+        if (window.api.codeAssistant) {
+            monaco.editor.registerCommand('acceptCodeAssistCompletion', (_accessor, ...args) => {
+                const data = args[0] ?? {};
+                if (!data || typeof data !== 'object') {
+                    return;
+                }
+                const {requestId, suggestionText} = data;
+                if (requestId && suggestionText) {
+                    inlineCompletionProviderInstance.handleAccept({requestId, suggestionText});
+                }
+            });
+
+            monaco.editor.registerCommand('declineCodeAssistCompletion', () => {
+                inlineCompletionProviderInstance.commandDiscard();
+            });
+        }
         initResizeHandler(editor);
         initUserPrompt(editor, getLastQueryText);
         editor.focus();
diff --git a/src/services/api/codeAssistant.ts b/src/services/api/codeAssistant.ts
@@ -0,0 +1,57 @@
+import {codeAssistantBackend as CODE_ASSISTANT_BACKEND} from '../../store';
+import type {
+    PromptFile,
+    Suggestions,
+    TelemetryEvent,
+    TelemetryOpenTabs,
+} from '../../types/api/codeAssistant';
+
+import {BaseYdbAPI} from './base';
+
+const ideInfo = {
+    Ide: 'yql',
+    IdeVersion: '1',
+    PluginFamily: 'yql',
+    PluginVersion: '0.2',
+};
+
+export class CodeAssistAPI extends BaseYdbAPI {
+    getPath(path: string) {
+        return `${CODE_ASSISTANT_BACKEND ?? ''}${path}`;
+    }
+
+    getCodeAssistSuggestions(data: PromptFile[]): Promise<Suggestions> {
+        return this.post<Suggestions>(
+            this.getPath('/code-assist-suggestion'),
+            {
+                Files: data,
+                ContextCreateType: 1,
+                IdeInfo: ideInfo,
+            },
+            null,
+            {
+                concurrentId: 'code-assist-suggestion',
+                collectRequest: false,
+            },
+        );
+    }
+
+    sendCodeAssistTelemetry(data: TelemetryEvent): Promise<unknown> {
+        return this.post(this.getPath('/code-assist-telemetry'), data, null, {
+            concurrentId: 'code-assist-telemetry',
+            collectRequest: false,
+        });
+    }
+
+    sendCodeAssistOpenTabs(data: TelemetryOpenTabs): Promise<unknown> {
+        return this.post(
+            this.getPath('/code-assist-telemetry'),
+            {OpenTabs: {Tabs: data, IdeInfo: ideInfo}},
+            null,
+            {
+                concurrentId: 'code-assist-telemetry',
+                collectRequest: false,
+            },
+        );
+    }
+}
diff --git a/src/services/api/index.ts b/src/services/api/index.ts
@@ -1,6 +1,7 @@
 import type {AxiosRequestConfig} from 'axios';
 
 import {AuthAPI} from './auth';
+import {CodeAssistAPI} from './codeAssistant';
 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;
+    codeAssistant?: CodeAssistAPI;
 
     constructor({config, webVersion}: {config: AxiosRequestConfig; webVersion?: boolean}) {
         this.auth = new AuthAPI({config});
         if (webVersion) {
             this.meta = new MetaAPI({config});
+            this.codeAssistant = new CodeAssistAPI({config});
         }
         this.operation = new OperationAPI({config});
         this.pdisk = new PDiskAPI({config});
diff --git a/src/store/configureStore.ts b/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 codeAssistantBackend = window.code_assistant_backend;
 
 const isSingleClusterMode = `${metaBackend}` === 'undefined';
 
diff --git a/src/store/index.ts b/src/store/index.ts
@@ -6,6 +6,7 @@ export {
     customBackend,
     metaBackend,
     webVersion,
+    codeAssistantBackend,
 } from './configureStore';
 export {rootReducer} from './reducers';
 
diff --git a/src/types/api/codeAssistant.ts b/src/types/api/codeAssistant.ts
@@ -0,0 +1,77 @@
+type IdeInfo = {
+    Ide: string;
+    IdeVersion: string;
+    PluginFamily: string;
+    PluginVersion: string;
+};
+
+export interface Prompt {
+    Files: PromptFile[];
+    ContextCreateType: ContextCreateType;
+    ForceSuggest?: boolean;
+    IdeInfo: IdeInfo;
+}
+
+export interface PromptPosition {
+    Ln: number;
+    Col: number;
+}
+
+export interface PromptFragment {
+    Text: string;
+    Start: PromptPosition;
+    End: PromptPosition;
+}
+
+export interface PromptFile {
+    Path: string;
+    Fragments: PromptFragment[];
+    Cursor: PromptPosition;
+}
+
+export type ContextCreateType = 1;
+
+export interface Suggestions {
+    Suggests: Suggestion[];
+    RequestId: string;
+}
+
+export type DiscardReason = 'OnCancel';
+
+export interface Suggestion {
+    Text: string;
+}
+
+export interface AcceptSuggestionEvent {
+    Accepted: {
+        RequestId: string;
+        Timestamp: number;
+        AcceptedText: string;
+        ConvertedText: string;
+    };
+}
+export interface DiscardSuggestionEvent {
+    Discarded: {
+        RequestId: string;
+        Timestamp: number;
+        DiscardReason: 'OnCancel';
+        DiscardedText: string;
+        CacheHitCount: number;
+    };
+}
+export interface IgnoreSuggestionEvent {
+    Ignored: {
+        RequestId: string;
+        Timestamp: number;
+        IgnoredText: string;
+    };
+}
+
+export type TelemetryEvent = AcceptSuggestionEvent | DiscardSuggestionEvent | IgnoreSuggestionEvent;
+
+type OpenTab = {
+    FileName: string;
+    Text: string;
+};
+
+export type TelemetryOpenTabs = OpenTab[];
diff --git a/src/types/window.d.ts b/src/types/window.d.ts
@@ -37,6 +37,7 @@ interface Window {
     web_version?: boolean;
     custom_backend?: string;
     meta_backend?: string;
+    code_assistant_backend?: string;
 
     userSettings?: import('../services/settings').SettingsObject;
     systemSettings?: import('../services/settings').SettingsObject;
diff --git a/src/utils/monaco/codeAssistTelemetry.ts b/src/utils/monaco/codeAssistTelemetry.ts
@@ -0,0 +1,134 @@
+import type Monaco from 'monaco-editor';
+import {v4} from 'uuid';
+
+import type {DiscardReason, PromptFile, TelemetryEvent} from '../../types/api/codeAssistant';
+
+const limitBeforeCursor = 8_000;
+const limitAfterCursor = 1_000;
+
+const sessionId = v4();
+
+export function getPromptFileContent(
+    model: Monaco.editor.ITextModel,
+    position: Monaco.Position,
+): PromptFile[] | undefined {
+    const linesContent = model.getLinesContent();
+    const prevTextInCurrentLine = linesContent[position.lineNumber - 1].slice(
+        0,
+        position.column - 1,
+    );
+    const postTextInCurrentLine = linesContent[position.lineNumber - 1].slice(position.column - 1);
+    const prevText = linesContent
+        .slice(0, position.lineNumber - 1)
+        .concat([prevTextInCurrentLine])
+        .join('\n');
+    const postText = [postTextInCurrentLine]
+        .concat(linesContent.slice(position.lineNumber))
+        .join('\n');
+    const cursorPostion = {Ln: position.lineNumber, Col: position.column};
+
+    const fragments = [];
+    if (prevText) {
+        fragments.push({
+            Text:
+                prevText.length > limitBeforeCursor
+                    ? prevText.slice(prevText.length - limitBeforeCursor)
+                    : prevText,
+            Start: {Ln: 1, Col: 1},
+            End: cursorPostion,
+        });
+    }
+    if (postText) {
+        fragments.push({
+            Text: postText.slice(0, limitAfterCursor),
+            Start: cursorPostion,
+            End: {
+                Ln: linesContent.length,
+                Col: linesContent[linesContent.length - 1].length,
+            },
+        });
+    }
+
+    return fragments.length
+        ? [
+              {
+                  Fragments: fragments,
+                  Cursor: cursorPostion,
+                  Path: `${sessionId}/query.yql`,
+              },
+          ]
+        : undefined;
+}
+
+interface SendCodeAssistTelemetry {
+    requestId: string;
+}
+
+interface SendCodeAssistTelemetryAccept extends SendCodeAssistTelemetry {
+    type: 'accept';
+
+    acceptedText: string;
+}
+interface SendCodeAssistTelemetryDecline extends SendCodeAssistTelemetry {
+    type: 'decline';
+    suggestionText: string;
+    reason: DiscardReason;
+    hitCount: number;
+}
+interface SendCodeAssistTelemetryIgnore extends SendCodeAssistTelemetry {
+    type: 'ignore';
+    suggestionText: string;
+}
+type SendCodeAssistTelemetryProps =
+    | SendCodeAssistTelemetryAccept
+    | SendCodeAssistTelemetryDecline
+    | SendCodeAssistTelemetryIgnore;
+
+function isAcceptType(data: SendCodeAssistTelemetryProps): data is SendCodeAssistTelemetryAccept {
+    return data.type === 'accept';
+}
+function isDeclineType(data: SendCodeAssistTelemetryProps): data is SendCodeAssistTelemetryDecline {
+    return data.type === 'decline';
+}
+function isIgnoreType(data: SendCodeAssistTelemetryProps): data is SendCodeAssistTelemetryIgnore {
+    return data.type === 'ignore';
+}
+
+export function sendCodeAssistTelemetry(props: SendCodeAssistTelemetryProps) {
+    let data: TelemetryEvent | null = null;
+    const timestamp = Date.now();
+    if (isAcceptType(props)) {
+        const {requestId, acceptedText} = props;
+        data = {
+            Accepted: {
+                RequestId: requestId,
+                Timestamp: timestamp,
+                AcceptedText: acceptedText,
+                ConvertedText: acceptedText,
+            },
+        };
+    } else if (isDeclineType(props)) {
+        const {requestId, suggestionText, reason, hitCount} = props;
+        data = {
+            Discarded: {
+                RequestId: requestId,
+                Timestamp: timestamp,
+                DiscardReason: reason,
+                DiscardedText: suggestionText,
+                CacheHitCount: hitCount,
+            },
+        };
+    } else if (isIgnoreType(props)) {
+        const {requestId, suggestionText} = props;
+        data = {
+            Ignored: {
+                RequestId: requestId,
+                Timestamp: timestamp,
+                IgnoredText: suggestionText,
+            },
+        };
+    }
+    if (data) {
+        window.api.codeAssistant?.sendCodeAssistTelemetry(data);
+    }
+}
diff --git a/src/utils/monaco/yql/ydb.inlineCompletionProvider.ts b/src/utils/monaco/yql/ydb.inlineCompletionProvider.ts
