From 77b3d6bedb97ae223e2c5876be16daa67adf67fa Mon Sep 17 00:00:00 2001
From: mamadoudicko <mamadoudicko100@gmail.com>
Date: Tue, 25 Jul 2023 16:49:04 +0200
Subject: [PATCH] add brain settings tab

---
 backend/core/routes/brain_routes.py           |   6 +-
 .../BrainManagementTabs.tsx                   |   6 +-
 .../components/SettingsTab/SettingsTab.tsx    | 121 ++++++++++++++
 .../SettingsTab/hooks/useSettingsTab.ts       | 154 ++++++++++++++++++
 .../api/brain/__tests__/useBrainApi.test.ts   |   7 +-
 frontend/lib/api/brain/brain.ts               |   7 +-
 frontend/lib/api/brain/types.ts               |   4 +-
 frontend/lib/api/brain/useBrainApi.ts         |   7 +-
 .../BrainProvider/hooks/useBrainProvider.ts   |   4 +-
 frontend/lib/context/BrainProvider/types.ts   |   9 +-
 10 files changed, 305 insertions(+), 20 deletions(-)
 create mode 100644 frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx
 create mode 100644 frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts

diff --git a/backend/core/routes/brain_routes.py b/backend/core/routes/brain_routes.py
index 0588a2ce2..de17b69c0 100644
--- a/backend/core/routes/brain_routes.py
+++ b/backend/core/routes/brain_routes.py
@@ -10,6 +10,7 @@ from models.brains import (
 )
 from models.settings import BrainRateLimiting
 from models.users import User
+
 from routes.authorizations.brain_authorization import RoleEnum, has_brain_authorization
 
 logger = get_logger(__name__)
@@ -74,10 +75,7 @@ async def get_brain_endpoint(
     brain = Brain(id=brain_id)
     brains = brain.get_brain_details()
     if len(brains) > 0:
-        return {
-            "id": brain_id,
-            "name": brains[0]["name"],
-        }
+        return brains[0]
     else:
         return HTTPException(
             status_code=404,
diff --git a/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/BrainManagementTabs.tsx b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/BrainManagementTabs.tsx
index 7ec26a520..1dc4675f2 100644
--- a/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/BrainManagementTabs.tsx
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/BrainManagementTabs.tsx
@@ -1,6 +1,7 @@
 import { Content, List, Root } from "@radix-ui/react-tabs";
 
 import { BrainTabTrigger, PeopleTab } from "./components";
+import { SettingsTab } from "./components/SettingsTab/SettingsTab";
 import { useBrainManagementTabs } from "./hooks/useBrainManagementTabs";
 
 export const BrainManagementTabs = (): JSX.Element => {
@@ -38,11 +39,14 @@ export const BrainManagementTabs = (): JSX.Element => {
 
       <div className="p-20 pt-5">
         <Content value="settings">
-          <p>coming soon</p>
+          <SettingsTab brainId={brainId} />
         </Content>
         <Content value="people">
           <PeopleTab brainId={brainId} />
         </Content>
+        <Content value="knowledge">
+          <p>Coming soon</p>
+        </Content>
       </div>
     </Root>
   );
diff --git a/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx
new file mode 100644
index 000000000..8f0e00108
--- /dev/null
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx
@@ -0,0 +1,121 @@
+/* eslint-disable max-lines */
+
+import { UUID } from "crypto";
+
+import Button from "@/lib/components/ui/Button";
+import { Divider } from "@/lib/components/ui/Divider";
+import Field from "@/lib/components/ui/Field";
+import { TextArea } from "@/lib/components/ui/TextField";
+import { models, paidModels } from "@/lib/context/BrainConfigProvider/types";
+import { defineMaxTokens } from "@/lib/helpers/defineMexTokens";
+
+import { ApiKeyConfig } from "./components";
+import { useSettingsTab } from "./hooks/useSettingsTab";
+
+type SettingsTabProps = {
+  brainId: UUID;
+};
+
+export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
+  const {
+    handleSubmit,
+    register,
+    hasChanges,
+    openAiKey,
+    temperature,
+    maxTokens,
+    model,
+    setAsDefaultBrainHandler,
+    isSettingAsDefault,
+    isUpdating,
+  } = useSettingsTab({ brainId });
+
+  return (
+    <form
+      onSubmit={(e) => void handleSubmit(e)}
+      className="my-10 mb-0 flex flex-col items-center gap-2"
+    >
+      <Field
+        label="Name"
+        placeholder="E.g. History notes"
+        autoComplete="off"
+        className="flex-1"
+        {...register("name")}
+      />
+      <TextArea
+        label="Description"
+        placeholder="My new brain is about..."
+        autoComplete="off"
+        className="flex-1 m-3"
+        {...register("description")}
+      />
+      <Divider text="Model config" />
+      <Field
+        label="OpenAI API Key"
+        placeholder="sk-xxx"
+        autoComplete="off"
+        className="flex-1"
+        {...register("openAiKey")}
+      />
+      <fieldset className="w-full flex flex-col mt-2">
+        <label className="flex-1 text-sm" htmlFor="model">
+          Model
+        </label>
+        <select
+          id="model"
+          {...register("model")}
+          className="px-5 py-2 dark:bg-gray-700 bg-gray-200 rounded-md"
+        >
+          {(openAiKey !== undefined ? paidModels : models).map(
+            (availableModel) => (
+              <option value={availableModel} key={availableModel}>
+                {availableModel}
+              </option>
+            )
+          )}
+        </select>
+      </fieldset>
+      <fieldset className="w-full flex mt-4">
+        <label className="flex-1" htmlFor="temp">
+          Temperature: {temperature}
+        </label>
+        <input
+          id="temp"
+          type="range"
+          min="0"
+          max="1"
+          step="0.01"
+          value={temperature}
+          {...register("temperature")}
+        />
+      </fieldset>
+      <fieldset className="w-full flex mt-4">
+        <label className="flex-1" htmlFor="tokens">
+          Max tokens: {maxTokens}
+        </label>
+        <input
+          type="range"
+          min="10"
+          max={defineMaxTokens(model)}
+          value={maxTokens}
+          {...register("maxTokens")}
+        />
+      </fieldset>
+      <div className="flex flex-row justify-end flex-1 w-full mt-8">
+        <Button isLoading={isUpdating} disabled={!hasChanges}>
+          Save changes
+        </Button>
+      </div>
+      <Divider text="Default brain" className="mt-4" />
+      <Button
+        variant={"secondary"}
+        isLoading={isSettingAsDefault}
+        onClick={() => void setAsDefaultBrainHandler()}
+      >
+        Set as default brain
+      </Button>
+
+      <ApiKeyConfig />
+    </form>
+  );
+};
diff --git a/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts
new file mode 100644
index 000000000..8379e6769
--- /dev/null
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts
@@ -0,0 +1,154 @@
+/* eslint-disable complexity */
+/* eslint-disable max-lines */
+import axios from "axios";
+import { UUID } from "crypto";
+import { FormEvent, useEffect, useState } from "react";
+import { useForm } from "react-hook-form";
+
+import { useBrainApi } from "@/lib/api/brain/useBrainApi";
+import { useBrainConfig } from "@/lib/context/BrainConfigProvider";
+import { defineMaxTokens } from "@/lib/helpers/defineMexTokens";
+import { useToast } from "@/lib/hooks";
+
+type UseSettingsTabProps = {
+  brainId: UUID;
+};
+
+// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
+export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
+  const [isUpdating, setIsUpdating] = useState(false);
+  const [isSettingAsDefault, setIsSettingHasDefault] = useState(false);
+  const { publish } = useToast();
+
+  const { setAsDefaultBrain, getBrain, updateBrain } = useBrainApi();
+  const { config } = useBrainConfig();
+
+  const defaultValues = {
+    ...config,
+    name: "",
+    description: "",
+    setDefault: false,
+  };
+
+  const {
+    register,
+    getValues,
+    reset,
+    watch,
+    setValue,
+
+    formState: { dirtyFields },
+  } = useForm({
+    defaultValues,
+  });
+
+  useEffect(() => {
+    const fetchBrain = async () => {
+      const brain = await getBrain(brainId);
+      if (brain === undefined) {
+        return;
+      }
+      reset({
+        ...brain,
+        maxTokens: brain.max_tokens,
+      });
+    };
+    void fetchBrain();
+  }, []);
+
+  const openAiKey = watch("openAiKey");
+  const model = watch("model");
+  const temperature = watch("temperature");
+  const maxTokens = watch("maxTokens");
+
+  useEffect(() => {
+    setValue("maxTokens", Math.min(maxTokens, defineMaxTokens(model)));
+  }, [maxTokens, model, setValue]);
+
+  const setAsDefaultBrainHandler = async () => {
+    try {
+      setIsSettingHasDefault(true);
+      await setAsDefaultBrain(brainId);
+      publish({
+        variant: "success",
+        text: "Brain set as default successfully",
+      });
+    } catch (err) {
+      if (axios.isAxiosError(err) && err.response?.status === 429) {
+        publish({
+          variant: "danger",
+          text: `${JSON.stringify(
+            (
+              err.response as {
+                data: { detail: string };
+              }
+            ).data.detail
+          )}`,
+        });
+
+        return;
+      }
+    } finally {
+      setIsSettingHasDefault(false);
+    }
+  };
+
+  const handleSubmit = async (e: FormEvent) => {
+    e.preventDefault();
+    const { name: isNameDirty } = dirtyFields;
+    const { name } = getValues();
+    if (isNameDirty !== undefined && isNameDirty && name.trim() === "") {
+      publish({
+        variant: "danger",
+        text: "Name is required",
+      });
+
+      return;
+    }
+
+    try {
+      setIsUpdating(true);
+
+      await updateBrain(brainId, getValues());
+
+      reset(defaultValues);
+      publish({
+        variant: "success",
+        text: "Brain created successfully",
+      });
+    } catch (err) {
+      if (axios.isAxiosError(err) && err.response?.status === 429) {
+        publish({
+          variant: "danger",
+          text: `${JSON.stringify(
+            (
+              err.response as {
+                data: { detail: string };
+              }
+            ).data.detail
+          )}`,
+        });
+      } else {
+        publish({
+          variant: "danger",
+          text: `${JSON.stringify(err)}`,
+        });
+      }
+    } finally {
+      setIsUpdating(false);
+    }
+  };
+
+  return {
+    handleSubmit,
+    register,
+    openAiKey,
+    model,
+    temperature,
+    maxTokens,
+    isUpdating,
+    hasChanges: Object.keys(dirtyFields).length > 0,
+    setAsDefaultBrainHandler,
+    isSettingAsDefault,
+  };
+};
diff --git a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
index b281ced43..e2ee7c588 100644
--- a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
+++ b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
@@ -4,8 +4,9 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 
 import { Subscription } from "../brain";
 import {
-  CreateOrUpdateBrainInput,
+  CreateBrainInput,
   SubscriptionUpdatableProperties,
+  UpdateBrainInput,
 } from "../types";
 import { useBrainApi } from "../useBrainApi";
 
@@ -61,7 +62,7 @@ describe("useBrainApi", () => {
       },
     } = renderHook(() => useBrainApi());
 
-    const brain: CreateOrUpdateBrainInput = {
+    const brain: CreateBrainInput = {
       name: "Test Brain",
       description: "This is a description",
       status: "public",
@@ -212,7 +213,7 @@ describe("useBrainApi", () => {
       },
     } = renderHook(() => useBrainApi());
     const brainId = "123";
-    const brain: CreateOrUpdateBrainInput = {
+    const brain: UpdateBrainInput = {
       name: "Test Brain",
       description: "This is a description",
       status: "public",
diff --git a/frontend/lib/api/brain/brain.ts b/frontend/lib/api/brain/brain.ts
index 9f2e3ed37..9589103b5 100644
--- a/frontend/lib/api/brain/brain.ts
+++ b/frontend/lib/api/brain/brain.ts
@@ -10,8 +10,9 @@ import {
 import { Document } from "@/lib/types/Document";
 
 import {
-  CreateOrUpdateBrainInput,
+  CreateBrainInput,
   SubscriptionUpdatableProperties,
+  UpdateBrainInput,
 } from "./types";
 import { mapBackendMinimalBrainToMinimalBrain } from "./utils/mapBackendMinimalBrainToMinimalBrain";
 import {
@@ -32,7 +33,7 @@ export const getBrainDocuments = async (
 };
 
 export const createBrain = async (
-  brain: CreateOrUpdateBrainInput,
+  brain: CreateBrainInput,
   axiosInstance: AxiosInstance
 ): Promise<MinimalBrainForUser> => {
   return mapBackendMinimalBrainToMinimalBrain(
@@ -130,7 +131,7 @@ export const setAsDefaultBrain = async (
 
 export const updateBrain = async (
   brainId: string,
-  brain: CreateOrUpdateBrainInput,
+  brain: UpdateBrainInput,
   axiosInstance: AxiosInstance
 ): Promise<void> => {
   await axiosInstance.put(`/brains/${brainId}/`, brain);
diff --git a/frontend/lib/api/brain/types.ts b/frontend/lib/api/brain/types.ts
index 0e145454e..eac7f31d6 100644
--- a/frontend/lib/api/brain/types.ts
+++ b/frontend/lib/api/brain/types.ts
@@ -4,7 +4,7 @@ export type SubscriptionUpdatableProperties = {
   role: BrainRoleType | null;
 };
 
-export type CreateOrUpdateBrainInput = {
+export type CreateBrainInput = {
   name: string;
   description?: string;
   status?: string;
@@ -13,3 +13,5 @@ export type CreateOrUpdateBrainInput = {
   max_tokens?: number;
   openai_api_key?: string;
 };
+
+export type UpdateBrainInput = Partial<CreateBrainInput>;
diff --git a/frontend/lib/api/brain/useBrainApi.ts b/frontend/lib/api/brain/useBrainApi.ts
index 42c53fa74..ad4a1bbc6 100644
--- a/frontend/lib/api/brain/useBrainApi.ts
+++ b/frontend/lib/api/brain/useBrainApi.ts
@@ -15,8 +15,9 @@ import {
   updateBrainAccess,
 } from "./brain";
 import {
-  CreateOrUpdateBrainInput,
+  CreateBrainInput,
   SubscriptionUpdatableProperties,
+  UpdateBrainInput,
 } from "./types";
 
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
@@ -26,7 +27,7 @@ export const useBrainApi = () => {
   return {
     getBrainDocuments: async (brainId: string) =>
       getBrainDocuments(brainId, axiosInstance),
-    createBrain: async (brain: CreateOrUpdateBrainInput) =>
+    createBrain: async (brain: CreateBrainInput) =>
       createBrain(brain, axiosInstance),
     deleteBrain: async (id: string) => deleteBrain(id, axiosInstance),
     getDefaultBrain: async () => getDefaultBrain(axiosInstance),
@@ -45,7 +46,7 @@ export const useBrainApi = () => {
     ) => updateBrainAccess(brainId, userEmail, subscription, axiosInstance),
     setAsDefaultBrain: async (brainId: string) =>
       setAsDefaultBrain(brainId, axiosInstance),
-    updateBrain: async (brainId: string, brain: CreateOrUpdateBrainInput) =>
+    updateBrain: async (brainId: string, brain: UpdateBrainInput) =>
       updateBrain(brainId, brain, axiosInstance),
   };
 };
diff --git a/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts b/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
index 1db25facf..93ceaff30 100644
--- a/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
+++ b/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
@@ -2,7 +2,7 @@
 import { UUID } from "crypto";
 import { useCallback, useState } from "react";
 
-import { CreateOrUpdateBrainInput } from "@/lib/api/brain/types";
+import { CreateBrainInput } from "@/lib/api/brain/types";
 import { useBrainApi } from "@/lib/api/brain/useBrainApi";
 import { useToast } from "@/lib/hooks";
 import { useEventTracking } from "@/services/analytics/useEventTracking";
@@ -29,7 +29,7 @@ export const useBrainProvider = () => {
   const currentBrain = allBrains.find((brain) => brain.id === currentBrainId);
 
   const createBrainHandler = async (
-    brain: CreateOrUpdateBrainInput
+    brain: CreateBrainInput
   ): Promise<UUID | undefined> => {
     const createdBrain = await createBrain(brain);
     try {
diff --git a/frontend/lib/context/BrainProvider/types.ts b/frontend/lib/context/BrainProvider/types.ts
index 6644fc774..c7342d7e8 100644
--- a/frontend/lib/context/BrainProvider/types.ts
+++ b/frontend/lib/context/BrainProvider/types.ts
@@ -4,15 +4,18 @@ import { BrainRoleType } from "@/lib/components/NavBar/components/NavItems/compo
 import { Document } from "@/lib/types/Document";
 
 import { useBrainProvider } from "./hooks/useBrainProvider";
+import { Model } from "../BrainConfigProvider/types";
 
 export type Brain = {
   id: UUID;
   name: string;
   documents?: Document[];
   status?: string;
-  model?: string;
-  max_tokens?: string;
-  temperature?: string;
+  model?: Model;
+  max_tokens?: number;
+  temperature?: number;
+  openai_api_key?: string;
+  description?: string;
 };
 
 export type MinimalBrainForUser = {