chore(Queries): add neural network config [YTFRONT-4612]

ytsaurus · Feb 3, 2025 · a0c94d0 · a0c94d0
1 parent 9118af4
commit a0c94d0
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 2 deletions.
diff --git a/packages/ui/src/server/ServerFactory.ts b/packages/ui/src/server/ServerFactory.ts
@@ -19,7 +19,7 @@ export interface ServerFactory {
         vcs: Omit<VCSSettings, 'type'>,
         token?: string,
     ): VcsApi | undefined;
-    createNeuralNetworkApi(): NeuralNetworkApi | undefined;
+    createNeuralNetworkApi(config: unknown): NeuralNetworkApi | undefined;
 }
 
 let app: ExpressKit;

diff --git a/packages/ui/src/server/components/layout-config.ts b/packages/ui/src/server/components/layout-config.ts
@@ -6,6 +6,7 @@ import {isUserColumnPresetsEnabled} from '../controllers/table-column-preset';
 import {getInterfaceVersion, isProductionEnv} from '../utils';
 import {getAuthWay} from '../utils/authorization';
 import {getOAuthSettings, isOAuthAllowed} from './oauth';
+import ServerFactory from '../ServerFactory';
 
 interface Params {
     name?: string;
@@ -70,6 +71,7 @@ export async function getLayoutConfig(req: Request, params: Params): Promise<App
             allowUserColumnPresets: isUserColumnPresetsEnabled(req),
             odinPageEnabled: Boolean(odinBaseUrl),
             allowTabletErrorsAPI: Boolean(tabletErrorsBaseUrl),
+            neuralNetwork: Boolean(ServerFactory.createNeuralNetworkApi(uiVersion.neuralNetwork)),
         },
         pluginsOptions: {
             yandexMetrika: {

diff --git a/packages/ui/src/server/controllers/neuralNetwork.ts b/packages/ui/src/server/controllers/neuralNetwork.ts
@@ -4,7 +4,9 @@ import {ErrorWithCode, sendApiError} from '../utils';
 
 export const getQuerySuggestions = async (req: Request, res: Response) => {
     try {
-        const nNetwork = ServerFactory.createNeuralNetworkApi();
+        const nNetwork = ServerFactory.createNeuralNetworkApi(
+            req.ctx.config.uiSettings.neuralNetworkConfig,
+        );
 
         if (!nNetwork) {
             throw new ErrorWithCode(500, 'Neural network is not configured');

diff --git a/packages/ui/src/shared/ui-settings.ts b/packages/ui/src/shared/ui-settings.ts
@@ -204,6 +204,16 @@ export interface UISettings {
      */
     vcsSettings?: VCSSettings[];
 
+    /**
+     * Adds your neural network configuration.
+     * The configuration will be passed to the 'createNeuralNetworkApi' method in ServerFactory
+     * @example
+     * neuralNetworkConfig: {
+     *     baseUrl: 'my_neural_network_api_url'
+     * }
+     */
+    neuralNetworkConfig?: unknown;
+
     /**
      * Allows to override idm object type to Effective ACL for objects that match the regex.
      * @example reUseEffectiveAclForPath: '//sys/access_control_object_namespaces[^/+]{0,}'

diff --git a/packages/ui/src/shared/yt-types.d.ts b/packages/ui/src/shared/yt-types.d.ts
@@ -300,6 +300,7 @@ export interface ConfigData {
     allowUserColumnPresets?: boolean;
     odinPageEnabled: boolean;
     allowTabletErrorsAPI: boolean;
+    neuralNetwork?: boolean;
 }
 
 export type PipelineParams = {

diff --git a/packages/ui/src/ui/libs/monaco-yql-languages/helpers/createInlineSuggestions.ts b/packages/ui/src/ui/libs/monaco-yql-languages/helpers/createInlineSuggestions.ts
@@ -4,6 +4,7 @@ import axios from 'axios';
 import {QueryEngine} from '../../../pages/query-tracker/module/engines';
 import debounce_ from 'lodash/debounce';
 import {getWindowStore} from '../../../store/window-store';
+import {getConfigData} from '../../../config/ui-settings';
 
 const getSuggestions = async (data: {
     contextId: string;
@@ -36,6 +37,13 @@ export const createInlineSuggestions =
         _token: CancellationToken,
     ): Promise<{items: languages.InlineCompletion[]}> => {
         const state = getWindowStore().getState();
+        const hasNeuralNetwork = getConfigData().neuralNetwork;
+
+        if (!hasNeuralNetwork) {
+            return {
+                items: [],
+            };
+        }
 
         const suggestions = await debouncedGetSuggestions({
             contextId: state.queryTracker.query.draft.annotations?.contextId || '',