feat: add openai api key validation (#1057)

QuivrHQ · Aug 30, 2023 · 082d9f8 · 082d9f8
1 parent 5515a60
commit 082d9f8
Show file tree

Hide file tree

Showing 9 changed files with 138 additions and 25 deletions.
diff --git a/...t/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts b/...t/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts
@@ -14,6 +14,8 @@ import { Brain } from "@/lib/context/BrainProvider/types";
 import { defineMaxTokens } from "@/lib/helpers/defineMaxTokens";
 import { useToast } from "@/lib/hooks";
 
+import { validateOpenAIKey } from "../utils/validateOpenAIKey";
+
 type UseSettingsTabProps = {
   brainId: UUID;
 };
@@ -95,8 +97,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
       if (brain.model !== undefined) {
         setValue("model", brain.model);
       }
-    },50);
-
+    }, 50);
   };
   useEffect(() => {
     void fetchBrain();
@@ -142,7 +143,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
       await setAsDefaultBrain(brainId);
       publish({
         variant: "success",
-        text: t("defaultBrainSet",{ns:"config"}),
+        text: t("defaultBrainSet", { ns: "config" }),
       });
       void fetchAllBrains();
       void fetchDefaultBrain();
@@ -180,12 +181,12 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
       void fetchBrain();
       publish({
         variant: "success",
-        text: t("promptRemoved",{ns:"config"}),
+        text: t("promptRemoved", { ns: "config" }),
       });
     } catch (err) {
       publish({
         variant: "danger",
-        text: t("errorRemovingPrompt",{ns:"config"}),
+        text: t("errorRemovingPrompt", { ns: "config" }),
       });
     } finally {
       setIsUpdating(false);
@@ -203,39 +204,47 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
     }
   };
 
-  const handleSubmit = async (checkDirty:boolean) => {
+  const handleSubmit = async (checkDirty: boolean) => {
     const hasChanges = Object.keys(dirtyFields).length > 0;
     if (!hasChanges && checkDirty) {
       return;
     }
     const { name: isNameDirty } = dirtyFields;
-    const { name } = getValues();
+    const { name, openAiKey: openai_api_key } = getValues();
     if (isNameDirty !== undefined && isNameDirty && name.trim() === "") {
       publish({
         variant: "danger",
-        text: t("nameRequired",{ns:"config"}),
+        text: t("nameRequired", { ns: "config" }),
       });
 
       return;
     }
 
+    if (
+      openai_api_key !== undefined &&
+      !(await validateOpenAIKey(
+        openai_api_key,
+        {
+          badApiKeyError: t("incorrectApiKey", { ns: "config" }),
+          invalidApiKeyError: t("invalidApiKeyError", { ns: "config" }),
+        },
+        publish
+      ))
+    ) {
+      return;
+    }
+
     try {
       setIsUpdating(true);
-
-      const {
-        maxTokens: max_tokens,
-        openAiKey: openai_api_key,
-        prompt,
-        ...otherConfigs
-      } = getValues();
+      const { maxTokens: max_tokens, prompt, ...otherConfigs } = getValues();
 
       if (
         dirtyFields["prompt"] &&
         (prompt.content === "" || prompt.title === "")
       ) {
         publish({
           variant: "warning",
-          text: t("promptFieldsRequired",{ns:"config"}),
+          text: t("promptFieldsRequired", { ns: "config" }),
         });
 
         return;
@@ -279,7 +288,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
 
       publish({
         variant: "success",
-        text: t("brainUpdated",{ns:"config"}),
+        text: t("brainUpdated", { ns: "config" }),
       });
       void fetchAllBrains();
     } catch (err) {
@@ -318,7 +327,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
     setValue("prompt.content", content, {
       shouldDirty: true,
     });
-  }; 
+  };
 
   return {
     handleSubmit,

diff --git a/...brainId]/components/BrainManagementTabs/components/SettingsTab/utils/validateOpenAIKey.ts b/...brainId]/components/BrainManagementTabs/components/SettingsTab/utils/validateOpenAIKey.ts
@@ -0,0 +1,71 @@
+import axios from "axios";
+
+import { ToastData } from "@/lib/components/ui/Toast/domain/types";
+import { getAxiosErrorParams } from "@/lib/helpers/getAxiosErrorParams";
+
+export const getOpenAIKeyValidationStatusCode = async (
+  key: string
+): Promise<number> => {
+  const url = "https://api.openai.com/v1/chat/completions";
+  const headers = {
+    Authorization: `Bearer ${key}`,
+    "Content-Type": "application/json",
+  };
+
+  const data = JSON.stringify({
+    model: "gpt-3.5-turbo",
+    messages: [
+      {
+        role: "user",
+        content: "Hello!",
+      },
+    ],
+  });
+
+  try {
+    await axios.post(url, data, { headers });
+
+    return 200;
+  } catch (error) {
+    return getAxiosErrorParams(error)?.status ?? 400;
+  }
+};
+
+type ErrorMessages = {
+  badApiKeyError: string;
+  invalidApiKeyError: string;
+};
+
+export const validateOpenAIKey = async (
+  openai_api_key: string | undefined,
+  errorMessages: ErrorMessages,
+  publish: (toast: ToastData) => void
+): Promise<boolean> => {
+  if (openai_api_key !== undefined) {
+    const keyValidationStatusCode = await getOpenAIKeyValidationStatusCode(
+      openai_api_key
+    );
+
+    if (keyValidationStatusCode !== 200) {
+      if (keyValidationStatusCode === 401) {
+        publish({
+          variant: "danger",
+          text: errorMessages.badApiKeyError,
+        });
+      }
+
+      if (keyValidationStatusCode === 429) {
+        publish({
+          variant: "danger",
+          text: errorMessages.invalidApiKeyError,
+        });
+      }
+
+      return false;
+    }
+
+    return true;
+  }
+
+  return false;
+};
diff --git a/frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts b/frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts
@@ -1,6 +1,8 @@
 /* eslint-disable max-lines */
 import { useEffect, useState } from "react";
+import { useTranslation } from "react-i18next";
 
+import { validateOpenAIKey } from "@/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/utils/validateOpenAIKey";
 import { useAuthApi } from "@/lib/api/auth/useAuthApi";
 import { useUserApi } from "@/lib/api/user/useUserApi";
 import { UserIdentity } from "@/lib/api/user/user";
@@ -20,6 +22,7 @@ export const useApiKeyConfig = () => {
   const { createApiKey } = useAuthApi();
   const { publish } = useToast();
   const [userIdentity, setUserIdentity] = useState<UserIdentity>();
+  const { t } = useTranslation(["config"]);
 
   const fetchUserIdentity = async () => {
     setUserIdentity(await getUserIdentity());
@@ -56,6 +59,24 @@ export const useApiKeyConfig = () => {
   const changeOpenAiApiKey = async () => {
     try {
       setChangeOpenAiApiKeyRequestPending(true);
+
+      if (
+        openAiApiKey !== undefined &&
+        openAiApiKey !== null &&
+        !(await validateOpenAIKey(
+          openAiApiKey,
+          {
+            badApiKeyError: t("incorrectApiKey", { ns: "config" }),
+            invalidApiKeyError: t("invalidApiKeyError", { ns: "config" }),
+          },
+          publish
+        ))
+      ) {
+        setChangeOpenAiApiKeyRequestPending(false);
+
+        return;
+      }
+
       await updateUserIdentity({
         openai_api_key: openAiApiKey,
       });

diff --git a/frontend/public/locales/en/config.json b/frontend/public/locales/en/config.json
@@ -46,5 +46,7 @@
     "roleRequired": "You don't have the necessary role to access this tab 🧠💡🥲.",
     "requireAccess": "Please require access from the owner.",
     "ohno": "Oh no!",
-    "noUser": "No user"
-}
+    "noUser": "No user",
+    "incorrectApiKey": "Incorrect API Key",
+    "invalidApiKeyError": "Invalid API Key"
+}   
diff --git a/frontend/public/locales/es/config.json b/frontend/public/locales/es/config.json
@@ -46,5 +46,7 @@
   "supabaseURLPlaceHolder": "URL de Supabase",
   "temperature": "Temperatura",
   "title": "Configuración",
-  "updatingBrainSettings": "Actualizando configuración del cerebro..."
+  "updatingBrainSettings": "Actualizando configuración del cerebro...",
+  "incorrectApiKey": "Clave de API incorrecta",
+  "invalidApiKeyError": "Clave de API inválida"
 }
diff --git a/frontend/public/locales/fr/config.json b/frontend/public/locales/fr/config.json
@@ -46,5 +46,7 @@
   "supabaseURLPlaceHolder": "URL Supabase",
   "temperature": "Température",
   "title": "Configuration",
-  "updatingBrainSettings": "Mise à jour des paramètres du cerveau..."
+  "updatingBrainSettings": "Mise à jour des paramètres du cerveau...",
+  "incorrectApiKey": "Clé API incorrecte",
+  "invalidApiKeyError": "Clé API invalide"
 }
diff --git a/frontend/public/locales/pt-br/config.json b/frontend/public/locales/pt-br/config.json
@@ -46,5 +46,7 @@
     "roleRequired": "Você não possui a função necessária para acessar esta aba 🧠💡🥲.",
     "requireAccess": "Por favor, solicite acesso ao proprietário.",
     "ohno": "Oh, não!",
-    "noUser": "Nenhum usuário"
+    "noUser": "Nenhum usuário",
+    "incorrectApiKey": "Chave de API incorreta",
+    "invalidApiKeyError": "Chave de API inválida"
 }
diff --git a/frontend/public/locales/ru/config.json b/frontend/public/locales/ru/config.json
@@ -46,5 +46,7 @@
   "roleRequired": "У вас нет необходимой роли для доступа к этой вкладке 🧠💡🥲.",
   "requireAccess": "Пожалуйста, запросите доступ у владельца.",
   "ohno": "О нет!",
-  "noUser": "Пользователь не найден"
+  "noUser": "Пользователь не найден",
+  "incorrectApiKey": "Неверный ключ API",
+  "invalidApiKeyError": "Недействительный ключ API"
 }
diff --git a/frontend/public/locales/zh-cn/config.json b/frontend/public/locales/zh-cn/config.json
@@ -46,5 +46,7 @@
     "roleRequired": "您没有访问此选项卡所需的权限 🧠💡🥲.",
     "requireAccess": "请向所有者申请访问权限.",
     "ohno": "哎呀!",
-    "noUser": "没有用户"
+    "noUser": "没有用户",
+    "incorrectApiKey": "无效的API密钥",
+    "invalidApiKeyError": "无效的API密钥"
 }