From b4a160b680a4e76e7ba53d7f387bdb87c9bced67 Mon Sep 17 00:00:00 2001
From: Arvin Xu <arvinx@foxmail.com>
Date: Fri, 17 Jan 2025 20:45:07 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix:=20fix=20api=20key=20in=20ap?=
 =?UTF-8?q?i=20key=20form=20(#5498)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix api key

* fix client fetch
---
 .../Error/APIKeyForm/ProviderApiKeyForm.tsx   | 11 ++----
 .../Error/APIKeyForm/useApiKey.ts             | 34 +++++++++++++++++++
 .../aiInfra/slices/aiProvider/selectors.ts    | 28 +++++++--------
 src/types/aiProvider.ts                       |  2 +-
 4 files changed, 52 insertions(+), 23 deletions(-)
 create mode 100644 src/features/Conversation/Error/APIKeyForm/useApiKey.ts

diff --git a/src/features/Conversation/Error/APIKeyForm/ProviderApiKeyForm.tsx b/src/features/Conversation/Error/APIKeyForm/ProviderApiKeyForm.tsx
index df0b5156f7d8..ad38774fc10d 100644
--- a/src/features/Conversation/Error/APIKeyForm/ProviderApiKeyForm.tsx
+++ b/src/features/Conversation/Error/APIKeyForm/ProviderApiKeyForm.tsx
@@ -6,11 +6,10 @@ import { useTranslation } from 'react-i18next';
 
 import { useProviderName } from '@/hooks/useProviderName';
 import { featureFlagsSelectors, useServerConfigStore } from '@/store/serverConfig';
-import { useUserStore } from '@/store/user';
-import { keyVaultsConfigSelectors } from '@/store/user/selectors';
 import { GlobalLLMProviderKey } from '@/types/user/settings';
 
 import { FormAction } from '../style';
+import { useApiKey } from './useApiKey';
 
 interface ProviderApiKeyFormProps {
   apiKeyPlaceholder?: string;
@@ -25,11 +24,7 @@ const ProviderApiKeyForm = memo<ProviderApiKeyFormProps>(
     const { t: errorT } = useTranslation('error');
     const [showProxy, setShow] = useState(false);
 
-    const [apiKey, proxyUrl, setConfig] = useUserStore((s) => [
-      keyVaultsConfigSelectors.getVaultByProvider(provider)(s)?.apiKey,
-      keyVaultsConfigSelectors.getVaultByProvider(provider)(s)?.baseURL,
-      s.updateKeyVaultConfig,
-    ]);
+    const { apiKey, baseURL, setConfig } = useApiKey(provider);
     const { showOpenAIProxyUrl } = useServerConfigStore(featureFlagsSelectors);
     const providerName = useProviderName(provider);
 
@@ -58,7 +53,7 @@ const ProviderApiKeyForm = memo<ProviderApiKeyFormProps>(
               }}
               placeholder={'https://api.openai.com/v1'}
               type={'block'}
-              value={proxyUrl}
+              value={baseURL}
             />
           ) : (
             <Button
diff --git a/src/features/Conversation/Error/APIKeyForm/useApiKey.ts b/src/features/Conversation/Error/APIKeyForm/useApiKey.ts
new file mode 100644
index 000000000000..8e252136b800
--- /dev/null
+++ b/src/features/Conversation/Error/APIKeyForm/useApiKey.ts
@@ -0,0 +1,34 @@
+import isEqual from 'fast-deep-equal';
+
+import { isDeprecatedEdition } from '@/const/version';
+import { aiProviderSelectors, useAiInfraStore } from '@/store/aiInfra';
+import { useUserStore } from '@/store/user';
+import { keyVaultsConfigSelectors } from '@/store/user/selectors';
+
+export const useApiKey = (provider: string) => {
+  const [apiKey, baseURL, setConfig] = useUserStore((s) => [
+    keyVaultsConfigSelectors.getVaultByProvider(provider as any)(s)?.apiKey,
+    keyVaultsConfigSelectors.getVaultByProvider(provider as any)(s)?.baseURL,
+    s.updateKeyVaultConfig,
+  ]);
+
+  const updateAiProviderConfig = useAiInfraStore((s) => s.updateAiProviderConfig);
+  const data = useAiInfraStore(aiProviderSelectors.providerConfigById(provider), isEqual);
+
+  // TODO: remove this in V2
+  if (isDeprecatedEdition) return { apiKey, baseURL, setConfig };
+  //
+
+  return {
+    apiKey: data?.keyVaults.apiKey,
+    baseURL: data?.keyVaults?.baseURL,
+    setConfig: async (id: string, params: Record<string, string>) => {
+      await updateAiProviderConfig(id, {
+        keyVaults: {
+          ...data?.keyVaults,
+          ...params,
+        },
+      });
+    },
+  };
+};
diff --git a/src/store/aiInfra/slices/aiProvider/selectors.ts b/src/store/aiInfra/slices/aiProvider/selectors.ts
index e2de2c626e38..354557a2b424 100644
--- a/src/store/aiInfra/slices/aiProvider/selectors.ts
+++ b/src/store/aiInfra/slices/aiProvider/selectors.ts
@@ -37,6 +37,17 @@ const isActiveProviderApiKeyNotEmpty = (s: AIProviderStoreState) => {
   return !!vault?.apiKey || !!vault?.accessKeyId || !!vault?.secretAccessKey;
 };
 
+const providerConfigById =
+  (id: string) =>
+  (s: AIProviderStoreState): AiProviderRuntimeConfig | undefined => {
+    if (!id) return undefined;
+
+    return s.aiProviderRuntimeConfig?.[id];
+  };
+
+const isProviderConfigUpdating = (id: string) => (s: AIProviderStoreState) =>
+  s.aiProviderConfigUpdatingIds.includes(id);
+
 /**
  * @description The conditions to enable client fetch
  * 1. If no baseUrl and apikey input, force on Server.
@@ -46,7 +57,7 @@ const isActiveProviderApiKeyNotEmpty = (s: AIProviderStoreState) => {
  */
 const isProviderFetchOnClient =
   (provider: GlobalLLMProviderKey | string) => (s: AIProviderStoreState) => {
-    const config = activeProviderConfig(s);
+    const config = providerConfigById(provider)(s);
 
     // If the provider already disable broswer request in model config, force on Server.
     if (isProviderDisableBroswerRequest(provider)) return false;
@@ -56,8 +67,8 @@ const isProviderFetchOnClient =
       return config?.fetchOnClient;
 
     // 1. If no baseUrl and apikey input, force on Server.
-    const isProviderEndpointNotEmpty = isActiveProviderEndpointNotEmpty(s);
-    const isProviderApiKeyNotEmpty = isActiveProviderApiKeyNotEmpty(s);
+    const isProviderEndpointNotEmpty = !!config?.keyVaults.baseURL;
+    const isProviderApiKeyNotEmpty = !!config?.keyVaults.apiKey;
     if (!isProviderEndpointNotEmpty && !isProviderApiKeyNotEmpty) return false;
 
     // 2. If only contains baseUrl, force on Client
@@ -76,17 +87,6 @@ const providerKeyVaults = (provider: string | undefined) => (s: AIProviderStoreS
   return s.aiProviderRuntimeConfig?.[provider]?.keyVaults;
 };
 
-const providerConfigById =
-  (id: string) =>
-  (s: AIProviderStoreState): AiProviderRuntimeConfig | undefined => {
-    if (!id) return undefined;
-
-    return s.aiProviderRuntimeConfig?.[id];
-  };
-
-const isProviderConfigUpdating = (id: string) => (s: AIProviderStoreState) =>
-  s.aiProviderConfigUpdatingIds.includes(id);
-
 export const aiProviderSelectors = {
   activeProviderConfig,
   disabledAiProviderList,
diff --git a/src/types/aiProvider.ts b/src/types/aiProvider.ts
index d97cc7669336..00cf99432661 100644
--- a/src/types/aiProvider.ts
+++ b/src/types/aiProvider.ts
@@ -201,7 +201,7 @@ export interface EnabledAiModel {
 
 export interface AiProviderRuntimeConfig {
   fetchOnClient?: boolean;
-  keyVaults: Record<string, object>;
+  keyVaults: Record<string, string>;
   settings: AiProviderSettings;
 }