From 3529222b95f2213e819c4ebc907f42a1e059e70b Mon Sep 17 00:00:00 2001
From: Mamadou DICKO <63923024+mamadoudicko@users.noreply.github.com>
Date: Tue, 25 Jul 2023 23:12:46 +0200
Subject: [PATCH] Brain management 4 (#762)

* feat: add <ApiKeyConfig/>

* feat(SDK): add update brain

* feat: add removeUndefined helper

* feat: remove unnecessary autofocus flag

* add brain settings tab

* ui: add tab delimitor

* feat: improve ux
---
 backend/core/routes/brain_routes.py           |   6 +-
 .../BrainManagementTabs.tsx                   |   6 +-
 .../components/BrainTabTrigger.tsx            |   4 +-
 .../components/SettingsTab/SettingsTab.tsx    | 130 +++++++++++++++
 .../SettingsTab/hooks/useSettingsTab.ts       | 156 ++++++++++++++++++
 .../components/ApiKeyConfig/ApiKeyConfig.tsx  |  41 +++++
 .../__tests__/ApiKeyConfig.test.tsx           |  52 ++++++
 .../hooks/__tests__/useApiKeyConfig.test.ts   |  67 ++++++++
 .../ApiKeyConfig/hooks/useApiKeyConfig.ts     |  42 +++++
 .../app/user/components/ApiKeyConfig/index.ts |   1 +
 .../app/user/components/UserStatistics.tsx    |   2 +
 .../api/brain/__tests__/useBrainApi.test.ts   |  29 +++-
 frontend/lib/api/brain/brain.ts               |  23 +--
 frontend/lib/api/brain/types.ts               |  12 ++
 frontend/lib/api/brain/useBrainApi.ts         |  10 +-
 .../AddBrainModal/AddBrainModal.tsx           |   2 -
 .../brain-config-provider.tsx                 |   5 +-
 .../BrainProvider/hooks/useBrainProvider.ts   |  22 ++-
 frontend/lib/context/BrainProvider/types.ts   |   9 +-
 frontend/lib/helpers/removeUndefined.ts       |  12 ++
 .../lib/helpers/setEmptyStringsUndefined.ts   |  11 --
 21 files changed, 595 insertions(+), 47 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
 create mode 100644 frontend/app/user/components/ApiKeyConfig/ApiKeyConfig.tsx
 create mode 100644 frontend/app/user/components/ApiKeyConfig/__tests__/ApiKeyConfig.test.tsx
 create mode 100644 frontend/app/user/components/ApiKeyConfig/hooks/__tests__/useApiKeyConfig.test.ts
 create mode 100644 frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts
 create mode 100644 frontend/app/user/components/ApiKeyConfig/index.ts
 create mode 100644 frontend/lib/helpers/removeUndefined.ts
 delete mode 100644 frontend/lib/helpers/setEmptyStringsUndefined.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/BrainTabTrigger.tsx b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/BrainTabTrigger.tsx
index 8371a271a..5c3cc2fc4 100644
--- a/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/BrainTabTrigger.tsx
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/BrainTabTrigger.tsx
@@ -16,8 +16,8 @@ export const BrainTabTrigger = ({
 }: BrainTabTriggerProps): JSX.Element => {
   return (
     <Trigger
-      className={`tracking-wide flex-1 text-lg align-center ${
-        selected ? "font-bold" : ""
+      className={`tracking-wide flex-1 pb-4 border-gray-500	 text-lg align-center ${
+        selected ? "font-bold border-b-2" : ""
       }`}
       value={value}
       onClick={() => onChange(value)}
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..223368037
--- /dev/null
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx
@@ -0,0 +1,130 @@
+/* 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 { 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,
+    isDefaultBrain,
+  } = useSettingsTab({ brainId });
+
+  return (
+    <form
+      onSubmit={(e) => void handleSubmit(e)}
+      className="my-10 mb-0 flex flex-col items-center gap-2"
+    >
+      <div className="flex flex-row flex-1 justify-between w-full">
+        <div>
+          <Field
+            label="Name"
+            placeholder="E.g. History notes"
+            autoComplete="off"
+            className="flex-1"
+            {...register("name")}
+          />
+        </div>
+        <div className="mt-4">
+          {isDefaultBrain ? (
+            <div className="border rounded-lg border-dashed border-black dark:border-white bg-white dark:bg-black text-black dark:text-white focus:bg-black dark:focus:bg-white dark dark focus:text-white dark:focus:text-black transition-colors py-2 px-4 shadow-none">
+              Default brain
+            </div>
+          ) : (
+            <Button
+              variant={"secondary"}
+              isLoading={isSettingAsDefault}
+              onClick={() => void setAsDefaultBrainHandler()}
+            >
+              Set as default brain
+            </Button>
+          )}
+        </div>
+      </div>
+      <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>
+    </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..81de0748d
--- /dev/null
+++ b/frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts
@@ -0,0 +1,156 @@
+/* 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 { useBrainProvider } from "@/lib/context/BrainProvider/hooks/useBrainProvider";
+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());
+
+      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);
+    }
+  };
+  const { defaultBrainId } = useBrainProvider();
+  const isDefaultBrain = defaultBrainId === brainId;
+
+  return {
+    handleSubmit,
+    register,
+    openAiKey,
+    model,
+    temperature,
+    maxTokens,
+    isUpdating,
+    hasChanges: Object.keys(dirtyFields).length > 0,
+    setAsDefaultBrainHandler,
+    isSettingAsDefault,
+    isDefaultBrain,
+  };
+};
diff --git a/frontend/app/user/components/ApiKeyConfig/ApiKeyConfig.tsx b/frontend/app/user/components/ApiKeyConfig/ApiKeyConfig.tsx
new file mode 100644
index 000000000..28ab0a20c
--- /dev/null
+++ b/frontend/app/user/components/ApiKeyConfig/ApiKeyConfig.tsx
@@ -0,0 +1,41 @@
+"use client";
+
+import Button from "@/lib/components/ui/Button";
+import { Divider } from "@/lib/components/ui/Divider";
+
+import { useApiKeyConfig } from "./hooks/useApiKeyConfig";
+
+export const ApiKeyConfig = (): JSX.Element => {
+  const { apiKey, handleCopyClick, handleCreateClick } = useApiKeyConfig();
+
+  return (
+    <>
+      <Divider text="API Key Config" className="mt-4" />
+      <div className="flex justify-center items-center mt-4">
+        <div className="flex items-center space-x-4">
+          {apiKey === "" && (
+            <Button
+              data-testid="create-new-key"
+              variant="secondary"
+              onClick={() => void handleCreateClick()}
+            >
+              Create New Key
+            </Button>
+          )}
+        </div>
+        {apiKey !== "" && (
+          <div className="flex items-center space-x-4">
+            <span className="text-gray-600">{apiKey}</span>
+            <Button
+              data-testid="copy-api-key-button"
+              variant="secondary"
+              onClick={handleCopyClick}
+            >
+              📋
+            </Button>
+          </div>
+        )}
+      </div>
+    </>
+  );
+};
diff --git a/frontend/app/user/components/ApiKeyConfig/__tests__/ApiKeyConfig.test.tsx b/frontend/app/user/components/ApiKeyConfig/__tests__/ApiKeyConfig.test.tsx
new file mode 100644
index 000000000..e720c558c
--- /dev/null
+++ b/frontend/app/user/components/ApiKeyConfig/__tests__/ApiKeyConfig.test.tsx
@@ -0,0 +1,52 @@
+import { fireEvent, render } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+import { ApiKeyConfig } from "../ApiKeyConfig";
+
+const handleCreateClickMock = vi.fn(() => ({}));
+const handleCopyClickMock = vi.fn(() => ({}));
+
+const useApiKeyConfigMock = vi.fn(() => ({
+  apiKey: "",
+  handleCreateClick: () => handleCreateClickMock(),
+  handleCopyClick: () => handleCopyClickMock(),
+}));
+
+vi.mock("../hooks/useApiKeyConfig", () => ({
+  useApiKeyConfig: () => useApiKeyConfigMock(),
+}));
+
+describe("ApiKeyConfig", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should render ApiConfig Component", () => {
+    const { getByText } = render(<ApiKeyConfig />);
+    expect(getByText("API Key Config")).toBeDefined();
+  });
+
+  it("renders 'Create New Key' button when apiKey is empty", () => {
+    const { getByTestId } = render(<ApiKeyConfig />);
+
+    const createButton = getByTestId("create-new-key");
+    expect(createButton).toBeDefined();
+
+    fireEvent.click(createButton);
+    expect(handleCreateClickMock).toHaveBeenCalledTimes(1);
+    expect(handleCreateClickMock).toHaveBeenCalledWith();
+  });
+  it('renders "Copy" button when apiKey is not empty', () => {
+    useApiKeyConfigMock.mockReturnValue({
+      apiKey: "123456789",
+      handleCreateClick: () => handleCreateClickMock(),
+      handleCopyClick: () => handleCopyClickMock(),
+    });
+
+    const { getByTestId } = render(<ApiKeyConfig />);
+    const copyButton = getByTestId("copy-api-key-button");
+    expect(copyButton).toBeDefined();
+    fireEvent.click(copyButton);
+    expect(handleCopyClickMock).toHaveBeenCalledTimes(1);
+  });
+});
diff --git a/frontend/app/user/components/ApiKeyConfig/hooks/__tests__/useApiKeyConfig.test.ts b/frontend/app/user/components/ApiKeyConfig/hooks/__tests__/useApiKeyConfig.test.ts
new file mode 100644
index 000000000..ae755f11f
--- /dev/null
+++ b/frontend/app/user/components/ApiKeyConfig/hooks/__tests__/useApiKeyConfig.test.ts
@@ -0,0 +1,67 @@
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+import { useApiKeyConfig } from "../useApiKeyConfig";
+
+const createApiKeyMock = vi.fn(() => "dummyApiKey");
+const trackMock = vi.fn((props: unknown) => ({ props }));
+
+const useAuthApiMock = vi.fn(() => ({
+  createApiKey: () => createApiKeyMock(),
+}));
+
+const useEventTrackingMock = vi.fn(() => ({
+  track: (props: unknown) => trackMock(props),
+}));
+
+vi.mock("@/lib/api/auth/useAuthApi", () => ({
+  useAuthApi: () => useAuthApiMock(),
+}));
+vi.mock("@/services/analytics/useEventTracking", () => ({
+  useEventTracking: () => useEventTrackingMock(),
+}));
+
+describe("useApiKeyConfig", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should set the apiKey when handleCreateClick is called", async () => {
+    const { result } = renderHook(() => useApiKeyConfig());
+
+    await act(async () => {
+      await result.current.handleCreateClick();
+    });
+
+    expect(createApiKeyMock).toHaveBeenCalledTimes(1);
+    expect(trackMock).toHaveBeenCalledWith("CREATE_API_KEY");
+    expect(result.current.apiKey).toBe("dummyApiKey");
+  });
+
+  it("should call copyToClipboard when handleCopyClick is called with a non-empty apiKey", () => {
+    vi.mock("react", async () => {
+      const actual = await vi.importActual<typeof import("react")>("react");
+
+      return {
+        ...actual,
+        useState: () => ["dummyApiKey", vi.fn()],
+      };
+    });
+    //@ts-ignore - clipboard is not actually readonly
+    global.navigator.clipboard = {
+      writeText: vi.fn(),
+    };
+
+    const { result } = renderHook(() => useApiKeyConfig());
+
+    act(() => result.current.handleCopyClick());
+
+    expect(trackMock).toHaveBeenCalledTimes(1);
+    expect(trackMock).toHaveBeenCalledWith("COPY_API_KEY");
+
+    // eslint-disable-next-line @typescript-eslint/unbound-method
+    expect(global.navigator.clipboard.writeText).toHaveBeenCalledWith(
+      "dummyApiKey"
+    );
+  });
+});
diff --git a/frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts b/frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts
new file mode 100644
index 000000000..6d9b3d064
--- /dev/null
+++ b/frontend/app/user/components/ApiKeyConfig/hooks/useApiKeyConfig.ts
@@ -0,0 +1,42 @@
+import { useState } from "react";
+
+import { useAuthApi } from "@/lib/api/auth/useAuthApi";
+import { useEventTracking } from "@/services/analytics/useEventTracking";
+
+// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
+export const useApiKeyConfig = () => {
+  const [apiKey, setApiKey] = useState("");
+  const { track } = useEventTracking();
+  const { createApiKey } = useAuthApi();
+
+  const handleCreateClick = async () => {
+    try {
+      void track("CREATE_API_KEY");
+      const createdApiKey = await createApiKey();
+      setApiKey(createdApiKey);
+    } catch (error) {
+      console.error("Error creating API key: ", error);
+    }
+  };
+
+  const copyToClipboard = async (text: string) => {
+    try {
+      void track("COPY_API_KEY");
+      await navigator.clipboard.writeText(text);
+    } catch (err) {
+      console.error("Failed to copy:", err);
+    }
+  };
+
+  const handleCopyClick = () => {
+    if (apiKey !== "") {
+      void copyToClipboard(apiKey);
+    }
+  };
+
+  return {
+    handleCreateClick,
+    apiKey,
+    handleCopyClick,
+  };
+};
diff --git a/frontend/app/user/components/ApiKeyConfig/index.ts b/frontend/app/user/components/ApiKeyConfig/index.ts
new file mode 100644
index 000000000..785aed328
--- /dev/null
+++ b/frontend/app/user/components/ApiKeyConfig/index.ts
@@ -0,0 +1 @@
+export * from "./ApiKeyConfig";
diff --git a/frontend/app/user/components/UserStatistics.tsx b/frontend/app/user/components/UserStatistics.tsx
index 45956c6e4..4df761b38 100644
--- a/frontend/app/user/components/UserStatistics.tsx
+++ b/frontend/app/user/components/UserStatistics.tsx
@@ -8,6 +8,7 @@ import Button from "@/lib/components/ui/Button";
 import { UserStats } from "@/lib/types/User";
 import { cn } from "@/lib/utils";
 
+import { ApiKeyConfig } from "./ApiKeyConfig";
 import { BrainConsumption } from "./BrainConsumption";
 import { DateComponent } from "./Date";
 import BrainSpaceChart from "./Graphs/BrainSpaceChart";
@@ -67,6 +68,7 @@ export const UserStatistics = (userStats: UserStats): JSX.Element => {
           </div>
         </UserStatisticsCard>
       </div>
+      <ApiKeyConfig />
     </>
   );
 };
diff --git a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
index 9d0e08e8b..e2ee7c588 100644
--- a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
+++ b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts
@@ -2,8 +2,12 @@
 import { renderHook } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 
-import { CreateBrainInput, Subscription } from "../brain";
-import { SubscriptionUpdatableProperties } from "../types";
+import { Subscription } from "../brain";
+import {
+  CreateBrainInput,
+  SubscriptionUpdatableProperties,
+  UpdateBrainInput,
+} from "../types";
 import { useBrainApi } from "../useBrainApi";
 
 const axiosGetMock = vi.fn(() => ({
@@ -201,4 +205,25 @@ describe("useBrainApi", () => {
     expect(axiosPostMock).toHaveBeenCalledTimes(1);
     expect(axiosPostMock).toHaveBeenCalledWith(`/brains/${brainId}/default`);
   });
+
+  it("should call updateBrain with correct brainId and brain", async () => {
+    const {
+      result: {
+        current: { updateBrain },
+      },
+    } = renderHook(() => useBrainApi());
+    const brainId = "123";
+    const brain: UpdateBrainInput = {
+      name: "Test Brain",
+      description: "This is a description",
+      status: "public",
+      model: "gpt-3.5-turbo-0613",
+      temperature: 0.0,
+      max_tokens: 256,
+      openai_api_key: "123",
+    };
+    await updateBrain(brainId, brain);
+    expect(axiosPutMock).toHaveBeenCalledTimes(1);
+    expect(axiosPutMock).toHaveBeenCalledWith(`/brains/${brainId}/`, brain);
+  });
 });
diff --git a/frontend/lib/api/brain/brain.ts b/frontend/lib/api/brain/brain.ts
index 236237b9f..9589103b5 100644
--- a/frontend/lib/api/brain/brain.ts
+++ b/frontend/lib/api/brain/brain.ts
@@ -9,7 +9,11 @@ import {
 } from "@/lib/context/BrainProvider/types";
 import { Document } from "@/lib/types/Document";
 
-import { SubscriptionUpdatableProperties } from "./types";
+import {
+  CreateBrainInput,
+  SubscriptionUpdatableProperties,
+  UpdateBrainInput,
+} from "./types";
 import { mapBackendMinimalBrainToMinimalBrain } from "./utils/mapBackendMinimalBrainToMinimalBrain";
 import {
   BackendSubscription,
@@ -28,15 +32,6 @@ export const getBrainDocuments = async (
   return response.data.documents;
 };
 
-export type CreateBrainInput = {
-  name: string;
-  description?: string;
-  status?: string;
-  model?: string;
-  temperature?: number;
-  max_tokens?: number;
-  openai_api_key?: string;
-};
 export const createBrain = async (
   brain: CreateBrainInput,
   axiosInstance: AxiosInstance
@@ -133,3 +128,11 @@ export const setAsDefaultBrain = async (
 ): Promise<void> => {
   await axiosInstance.post(`/brains/${brainId}/default`);
 };
+
+export const updateBrain = async (
+  brainId: string,
+  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 3a0546d3c..eac7f31d6 100644
--- a/frontend/lib/api/brain/types.ts
+++ b/frontend/lib/api/brain/types.ts
@@ -3,3 +3,15 @@ import { BrainRoleType } from "@/lib/components/NavBar/components/NavItems/compo
 export type SubscriptionUpdatableProperties = {
   role: BrainRoleType | null;
 };
+
+export type CreateBrainInput = {
+  name: string;
+  description?: string;
+  status?: string;
+  model?: string;
+  temperature?: number;
+  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 ee6928559..ad4a1bbc6 100644
--- a/frontend/lib/api/brain/useBrainApi.ts
+++ b/frontend/lib/api/brain/useBrainApi.ts
@@ -3,7 +3,6 @@ import { useAxios } from "@/lib/hooks";
 import {
   addBrainSubscriptions,
   createBrain,
-  CreateBrainInput,
   deleteBrain,
   getBrain,
   getBrainDocuments,
@@ -12,9 +11,14 @@ import {
   getDefaultBrain,
   setAsDefaultBrain,
   Subscription,
+  updateBrain,
   updateBrainAccess,
 } from "./brain";
-import { SubscriptionUpdatableProperties } from "./types";
+import {
+  CreateBrainInput,
+  SubscriptionUpdatableProperties,
+  UpdateBrainInput,
+} from "./types";
 
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
 export const useBrainApi = () => {
@@ -42,5 +46,7 @@ export const useBrainApi = () => {
     ) => updateBrainAccess(brainId, userEmail, subscription, axiosInstance),
     setAsDefaultBrain: async (brainId: string) =>
       setAsDefaultBrain(brainId, axiosInstance),
+    updateBrain: async (brainId: string, brain: UpdateBrainInput) =>
+      updateBrain(brainId, brain, axiosInstance),
   };
 };
diff --git a/frontend/lib/components/AddBrainModal/AddBrainModal.tsx b/frontend/lib/components/AddBrainModal/AddBrainModal.tsx
index fe7a2b19b..e3b9b57a3 100644
--- a/frontend/lib/components/AddBrainModal/AddBrainModal.tsx
+++ b/frontend/lib/components/AddBrainModal/AddBrainModal.tsx
@@ -52,7 +52,6 @@ export const AddBrainModal = (): JSX.Element => {
 
         <TextArea
           label="Enter a brain description"
-          autoFocus
           placeholder="My new brain is about..."
           autoComplete="off"
           className="flex-1 m-3"
@@ -61,7 +60,6 @@ export const AddBrainModal = (): JSX.Element => {
 
         <Field
           label="OpenAI API Key"
-          autoFocus
           placeholder="sk-xxx"
           autoComplete="off"
           className="flex-1"
diff --git a/frontend/lib/context/BrainConfigProvider/brain-config-provider.tsx b/frontend/lib/context/BrainConfigProvider/brain-config-provider.tsx
index d6f0fac89..aa10a8d57 100644
--- a/frontend/lib/context/BrainConfigProvider/brain-config-provider.tsx
+++ b/frontend/lib/context/BrainConfigProvider/brain-config-provider.tsx
@@ -3,8 +3,7 @@
 
 import { createContext, useEffect, useState } from "react";
 
-import { setEmptyStringsUndefined } from "@/lib/helpers/setEmptyStringsUndefined";
-
+import { removeUndefined } from "@/lib/helpers/removeUndefined";
 import {
   getBrainConfigFromLocalStorage,
   saveBrainConfigInLocalStorage,
@@ -39,7 +38,7 @@ export const BrainConfigProvider = ({
     setBrainConfig((config) => {
       const updatedConfig: BrainConfig = {
         ...config,
-        ...setEmptyStringsUndefined(newConfig),
+        ...removeUndefined(newConfig),
       };
       saveBrainConfigInLocalStorage(updatedConfig);
 
diff --git a/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts b/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
index 610b34910..17c512969 100644
--- a/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
+++ b/frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
@@ -1,8 +1,8 @@
 /* eslint-disable max-lines */
 import { UUID } from "crypto";
-import { useCallback, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 
-import { CreateBrainInput } from "@/lib/api/brain/brain";
+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";
@@ -24,10 +24,10 @@ export const useBrainProvider = () => {
 
   const [allBrains, setAllBrains] = useState<MinimalBrainForUser[]>([]);
   const [currentBrainId, setCurrentBrainId] = useState<null | UUID>(null);
+  const [defaultBrainId, setDefaultBrainId] = useState<UUID>();
   const [isFetchingBrains, setIsFetchingBrains] = useState(false);
 
   const currentBrain = allBrains.find((brain) => brain.id === currentBrainId);
-
   const createBrainHandler = async (
     brain: CreateBrainInput
   ): Promise<UUID | undefined> => {
@@ -79,10 +79,10 @@ export const useBrainProvider = () => {
   );
 
   const setDefaultBrain = useCallback(async () => {
-    const defaultBrain = await getDefaultBrain();
-    if (defaultBrain !== undefined) {
-      saveBrainInLocalStorage(defaultBrain);
-      setActiveBrain({ ...defaultBrain });
+    const userDefaultBrain = await getDefaultBrain();
+    if (userDefaultBrain !== undefined) {
+      saveBrainInLocalStorage(userDefaultBrain);
+      setActiveBrain(userDefaultBrain);
     } else {
       console.warn("No brains found");
     }
@@ -97,6 +97,13 @@ export const useBrainProvider = () => {
     }
   }, [setDefaultBrain, setActiveBrain]);
 
+  useEffect(() => {
+    const fetchDefaultBrain = async () => {
+      setDefaultBrainId((await getDefaultBrain())?.id);
+    };
+    void fetchDefaultBrain();
+  }, []);
+
   return {
     currentBrain,
     currentBrainId,
@@ -108,5 +115,6 @@ export const useBrainProvider = () => {
     setDefaultBrain,
     fetchAndSetActiveBrain,
     isFetchingBrains,
+    defaultBrainId,
   };
 };
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 = {
diff --git a/frontend/lib/helpers/removeUndefined.ts b/frontend/lib/helpers/removeUndefined.ts
new file mode 100644
index 000000000..dc6fc6565
--- /dev/null
+++ b/frontend/lib/helpers/removeUndefined.ts
@@ -0,0 +1,12 @@
+export const removeUndefined = <T extends Record<string, unknown>>(
+  obj: T
+): Partial<T> => {
+  const newObj = {} as Partial<T>;
+  for (const key in obj) {
+    if (obj[key] !== undefined) {
+      newObj[key] = obj[key];
+    }
+  }
+
+  return newObj;
+};
diff --git a/frontend/lib/helpers/setEmptyStringsUndefined.ts b/frontend/lib/helpers/setEmptyStringsUndefined.ts
deleted file mode 100644
index aaa0c12a0..000000000
--- a/frontend/lib/helpers/setEmptyStringsUndefined.ts
+++ /dev/null
@@ -1,11 +0,0 @@
-export const setEmptyStringsUndefined = (
-  obj: Record<string, unknown>
-): Record<string, unknown> => {
-  Object.keys(obj).forEach((key) => {
-    if (obj[key] === "") {
-      obj[key] = undefined;
-    }
-  });
-
-  return obj;
-};