From 62a70b30f33c2b83269ed2f1b655c76dc09d3426 Mon Sep 17 00:00:00 2001
From: Lin Wang <wonglam@amazon.com>
Date: Fri, 22 Dec 2023 16:29:50 +0800
Subject: [PATCH] Address PR comments

Signed-off-by: Lin Wang <wonglam@amazon.com>
---
 .../register_model/model_deployment.tsx       | 37 +++-------
 .../register_model/model_source.tsx           | 17 +++--
 .../pretrained_model_select.tsx               | 71 +++++++++----------
 .../register_model/register_model.tsx         |  5 +-
 4 files changed, 57 insertions(+), 73 deletions(-)

diff --git a/public/components/register_model/model_deployment.tsx b/public/components/register_model/model_deployment.tsx
index 27855849..97eceefd 100644
--- a/public/components/register_model/model_deployment.tsx
+++ b/public/components/register_model/model_deployment.tsx
@@ -3,48 +3,33 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useCallback, useState } from 'react';
+import React from 'react';
 import { EuiCheckbox, EuiText, EuiFormRow } from '@elastic/eui';
 import { useController, useFormContext } from 'react-hook-form';
 import { useSearchParams } from '../../hooks/use_search_params';
 export const ModelDeployment = () => {
   const searchParams = useSearchParams();
   const typeParams = searchParams.get('type');
-  const [checked, setChecked] = useState(false);
   const { control } = useFormContext<{ deployment: boolean }>();
   const modelDeploymentController = useController({
     name: 'deployment',
     control,
   });
+  const isRegisterExternal = typeParams === 'external';
 
   const { ref: deploymentInputRef, ...deploymentField } = modelDeploymentController.field;
-  const onDeploymentChange = useCallback(
-    (e) => {
-      setChecked(e.target.checked);
-      deploymentField.onChange(checked);
-    },
-    [deploymentField, checked]
-  );
   return (
-    <EuiFormRow label={typeParams === 'external' ? 'Activation' : 'Deployment'}>
+    <EuiFormRow label={isRegisterExternal ? 'Activation' : 'Deployment'}>
       <div>
         {<EuiText size="xs">Needs a description</EuiText>}
-        {(typeParams === 'upload' || typeParams === 'import') && (
-          <EuiCheckbox
-            id="deployment"
-            label="Start deployment automatically"
-            checked={checked}
-            onChange={onDeploymentChange}
-          />
-        )}
-        {typeParams === 'external' && (
-          <EuiCheckbox
-            id="activation"
-            label="Activate on registration"
-            checked={checked}
-            onChange={onDeploymentChange}
-          />
-        )}
+        <EuiCheckbox
+          id="deployment"
+          aria-label={
+            isRegisterExternal ? 'Activate on registration' : 'Start deployment automatically'
+          }
+          checked={deploymentField.value}
+          onChange={deploymentField.onChange}
+        />
       </div>
     </EuiFormRow>
   );
diff --git a/public/components/register_model/model_source.tsx b/public/components/register_model/model_source.tsx
index a9bdc4f5..7c7d6407 100644
--- a/public/components/register_model/model_source.tsx
+++ b/public/components/register_model/model_source.tsx
@@ -12,15 +12,20 @@ import {
   EuiFormRow,
   EuiComboBox,
 } from '@elastic/eui';
-
 import { useController, useFormContext } from 'react-hook-form';
-import { useMonitoring } from '../monitoring/use_monitoring';
+
+import { useFetcher } from '../../hooks';
+import { APIProvider } from '../../apis/api_provider';
 
 export const ModelSource = () => {
-  const { allExternalConnectors } = useMonitoring();
-  const connectorOptions = allExternalConnectors?.map((item) => {
-    return Object.assign({}, { label: item.name, value: item.id });
-  });
+  const { data: allConnectorsData } = useFetcher(APIProvider.getAPI('connector').getAll);
+  const connectorOptions = useMemo(
+    () =>
+      allConnectorsData?.data?.map((item) => {
+        return Object.assign({}, { label: item.name, value: item.id });
+      }),
+    [allConnectorsData]
+  );
   const { control } = useFormContext<{ modelConnector: string }>();
 
   const modelConnectorController = useController({
diff --git a/public/components/register_model/pretrained_model_select.tsx b/public/components/register_model/pretrained_model_select.tsx
index f4d75485..763c03e7 100644
--- a/public/components/register_model/pretrained_model_select.tsx
+++ b/public/components/register_model/pretrained_model_select.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React, { useCallback, useState, Fragment, useEffect } from 'react';
+import React, { useCallback, Fragment } from 'react';
 import {
   EuiSpacer,
   EuiTextColor,
@@ -12,10 +12,12 @@ import {
   EuiSelectableOption,
   EuiHighlight,
 } from '@elastic/eui';
-import { useHistory } from 'react-router-dom';
-import { generatePath } from 'react-router-dom';
+import { useHistory, generatePath } from 'react-router-dom';
+import { useObservable } from 'react-use';
+
 import { modelRepositoryManager } from '../../utils/model_repository_manager';
 import { routerPaths } from '../../../common/router_paths';
+
 interface IItem {
   label: string;
   checked?: 'on' | undefined;
@@ -34,43 +36,36 @@ const renderModelOption = (option: IItem, searchValue: string) => {
     </>
   );
 };
-export const PreTrainedModelSelect = () => {
-  useEffect(() => {
-    const subscribe = modelRepositoryManager.getPreTrainedModels$().subscribe((models) => {
-      setModelRepoSelection(
-        Object.keys(models).map((name) => ({
-          label: name,
-          description: models[name].description,
-          checked: undefined,
-        }))
-      );
-    });
-    return () => {
-      subscribe.unsubscribe();
-    };
-  }, []);
-  const [modelRepoSelection, setModelRepoSelection] = useState<Array<EuiSelectableOption<IItem>>>(
-    []
-  );
+export const PreTrainedModelSelect = ({
+  checkedPreTrainedModel,
+}: {
+  checkedPreTrainedModel?: string;
+}) => {
+  const preTrainedModels = useObservable(modelRepositoryManager.getPreTrainedModels$());
+  const preTrainedModelOptions = preTrainedModels
+    ? Object.keys(preTrainedModels).map((name) => ({
+        label: name,
+        description: preTrainedModels[name].description,
+        checked: checkedPreTrainedModel === name ? ('on' as const) : undefined,
+      }))
+    : [];
+
   const history = useHistory();
   const onChange = useCallback(
-    (modelSelection: Array<EuiSelectableOption<IItem>>) => {
-      setModelRepoSelection(modelSelection);
-      // ShowRest(true);
+    (options: Array<EuiSelectableOption<IItem>>) => {
+      const selectedOption = options.find((option) => option.checked === 'on');
+
+      if (selectedOption?.label) {
+        history.push(
+          `${generatePath(routerPaths.registerModel, { id: undefined })}/?type=import&name=${
+            selectedOption.label
+          }`
+        );
+      }
     },
-    // [ShowRest]
-    []
+    [history]
   );
-  useEffect(() => {
-    const selectedOption = modelRepoSelection.find((option) => option.checked === 'on');
-    if (selectedOption?.label) {
-      history.push(
-        `${generatePath(routerPaths.registerModel, { id: undefined })}/?type=import&name=${
-          selectedOption?.label
-        }`
-      );
-    }
-  }, [modelRepoSelection, history]);
+
   return (
     <div>
       <small>
@@ -95,7 +90,7 @@ export const PreTrainedModelSelect = () => {
           'data-test-subj': 'findModel',
           placeholder: 'Find model',
         }}
-        options={modelRepoSelection}
+        options={preTrainedModelOptions}
         onChange={onChange}
         singleSelection={true}
         noMatchesMessage="No model found"
@@ -106,7 +101,7 @@ export const PreTrainedModelSelect = () => {
           'data-test-subj': 'opensearchModelList',
           showIcons: true,
         }}
-        isLoading={modelRepoSelection.length === 0}
+        isLoading={!preTrainedModels}
       >
         {(list, search) => (
           <Fragment>
diff --git a/public/components/register_model/register_model.tsx b/public/components/register_model/register_model.tsx
index 7c47ea2a..49880983 100644
--- a/public/components/register_model/register_model.tsx
+++ b/public/components/register_model/register_model.tsx
@@ -295,7 +295,7 @@ export const RegisterModelForm = ({ defaultValues = DEFAULT_VALUES }: RegisterMo
   const errorCount = Object.keys(form.formState.errors).length;
   const formHeader = (
     <>
-      <EuiPageHeader pageTitle={getPageTitle} />
+      <EuiPageHeader pageTitle={getPageTitle()} />
       <EuiText style={{ maxWidth: 725 }}>
         <small>
           {registerToModelId && (
@@ -353,7 +353,6 @@ export const RegisterModelForm = ({ defaultValues = DEFAULT_VALUES }: RegisterMo
                 )}
               </React.Fragment>
             ))}
-            {/* {formType === 'import' ? nameParams && formFooter : formFooter} */}
           </EuiPanel>
           <EuiSpacer size="xxl" />
           <EuiSpacer size="xxl" />
@@ -380,7 +379,7 @@ export const RegisterModelForm = ({ defaultValues = DEFAULT_VALUES }: RegisterMo
               <EuiFlexItem grow={false}>
                 <EuiButton
                   form={FORM_ID}
-                  disabled={form.formState.isSubmitting}
+                  disabled={(formType === 'import' && !nameParams) || form.formState.isSubmitting}
                   isLoading={form.formState.isSubmitting}
                   type="submit"
                   onClick={() => setIsSubmitted(true)}