fix: allow user to set a brain as public after creation (#1646)

Issue: #1647


- Refactor brain management page settings tabs hooks: use context
- Fix brain status change 

Demo:



https://github.com/StanGirard/quivr/assets/63923024/073be02f-394c-4887-8572-ff293792c023

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.test.tsx

@@ -35,21 +35,24 @@ vi.mock("@/lib/api/brain/useBrainApi", () => ({
   }),
 }));
 
+vi.mock("next/navigation", () => ({
+  useRouter: () => ({ replace: vi.fn() }),
+  useParams: () => ({}),
+}));
+
 vi.mock("@tanstack/react-query", async () => {
   const actual = await vi.importActual<typeof import("@tanstack/react-query")>(
     "@tanstack/react-query"
   );
 
-  vi.mock("next/navigation", () => ({
-    useRouter: () => ({ replace: vi.fn() }),
-    useParams: () => ({}),
-  }));
-
   return {
     ...actual,
     useQuery: () => ({
       data: {},
     }),
+    useQueryClient: () => ({
+      invalidateQueries: vi.fn(),
+    }),
   };
 });
 
@@ -83,13 +86,10 @@ describe("Settings tab in brains-management", () => {
       </SupabaseProviderMock>
     );
 
-    expect(
-      screen.getByRole("button", { name: "setDefaultBrain" })
-    ).toBeVisible();
+    expect(screen.getByText("defaultBrain")).toBeVisible();
     expect(screen.getByText("brainName")).toBeVisible();
     expect(screen.getByLabelText("brainDescription")).toBeVisible();
 
     expect(screen.getByLabelText("promptName")).toBeVisible();
   });
 });
-2;

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/SettingsTab.tsx

@@ -1,15 +1,19 @@
 /* eslint max-lines:["error", 135] */
 
 import { UUID } from "crypto";
+import { FormProvider, useForm } from "react-hook-form";
 import { useTranslation } from "react-i18next";
 import { FaSpinner } from "react-icons/fa";
 
 import { Divider } from "@/lib/components/ui/Divider";
+import { defaultBrainConfig } from "@/lib/config/defaultBrainConfig";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
+import { BrainConfig } from "@/lib/types/brainConfig";
 
 import { GeneralInformation, ModelSelection, Prompt } from "./components";
 import { AccessConfirmationModal } from "./components/PrivateAccessConfirmationModal/AccessConfirmationModal";
 import { useAccessConfirmationModal } from "./components/PrivateAccessConfirmationModal/hooks/useAccessConfirmationModal";
+import { UsePromptProps } from "./hooks/usePrompt";
 import { useSettingsTab } from "./hooks/useSettingsTab";
 import { getBrainPermissions } from "../../utils/getBrainPermissions";
 
@@ -18,51 +22,27 @@ type SettingsTabProps = {
 };
 
 // eslint-disable-next-line complexity
-export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
+export const SettingsTabContent = ({
+  brainId,
+}: SettingsTabProps): JSX.Element => {
   const { t } = useTranslation(["translation", "brain", "config"]);
   const {
     handleSubmit,
-    register,
-    temperature,
-    maxTokens,
-    model,
     setAsDefaultBrainHandler,
     isSettingAsDefault,
     isUpdating,
     isDefaultBrain,
     formRef,
     accessibleModels,
-    status,
-    setValue,
-    dirtyFields,
-    resetField,
     setIsUpdating,
-    promptId,
-    getValues,
-    reset,
-    updateFormValues,
   } = useSettingsTab({ brainId });
 
-  const promptProps = {
-    brainId,
-    getValues,
-    promptId,
-    register,
-    reset,
-    setValue,
-    resetField,
-    updateFormValues,
-    dirtyFields,
+  const promptProps: UsePromptProps = {
     setIsUpdating,
   };
 
   const { onCancel, isAccessModalOpened, closeModal } =
-    useAccessConfirmationModal({
-      status,
-      setValue,
-      isStatusDirty: Boolean(dirtyFields.status),
-      resetField,
-    });
+    useAccessConfirmationModal();
 
   const { allBrains } = useBrainContext();
 
@@ -88,16 +68,11 @@ export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
           isOwnedByCurrentUser={isOwnedByCurrentUser}
           isPublicBrain={isPublicBrain}
           isSettingAsDefault={isSettingAsDefault}
-          register={register}
           setAsDefaultBrainHandler={setAsDefaultBrainHandler}
         />
         <Divider text={t("modelSection", { ns: "config" })} />
         <ModelSelection
           accessibleModels={accessibleModels}
-          model={model}
-          maxTokens={maxTokens}
-          temperature={temperature}
-          register={register}
           hasEditRights={hasEditRights}
           brainId={brainId}
           handleSubmit={handleSubmit}
@@ -122,8 +97,19 @@ export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
         onClose={onCancel}
         onCancel={onCancel}
         onConfirm={closeModal}
-        selectedStatus={status}
       />
     </>
   );
 };
+
+export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
+  const methods = useForm<BrainConfig>({
+    defaultValues: defaultBrainConfig,
+  });
+
+  return (
+    <FormProvider {...methods}>
+      <SettingsTabContent brainId={brainId} />
+    </FormProvider>
+  );
+};

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/GeneralInformation/GeneralInformation.tsx

@@ -1,21 +1,19 @@
 /* eslint max-lines:["error", 150] */
 // TODO: useFormContext to avoid passing too many props
 
-import { UseFormRegister } from "react-hook-form";
 import { useTranslation } from "react-i18next";
 
 import Button from "@/lib/components/ui/Button";
 import { Chip } from "@/lib/components/ui/Chip";
 import Field from "@/lib/components/ui/Field";
 import { Radio } from "@/lib/components/ui/Radio";
 import { TextArea } from "@/lib/components/ui/TextArea";
-import { BrainConfig } from "@/lib/types/brainConfig";
 
 import { ApiBrainDefinition } from "./components/ApiBrainDefinition";
 import { useGeneralInformation } from "./hooks/useGeneralInformation";
+import { useBrainFormState } from "../../hooks/useBrainFormState";
 
 type GeneralInformationProps = {
-  register: UseFormRegister<BrainConfig>;
   hasEditRights: boolean;
   isPublicBrain: boolean;
   isOwnedByCurrentUser: boolean;
@@ -29,14 +27,15 @@ export const GeneralInformation = (
 ): JSX.Element => {
   const { t } = useTranslation(["translation", "brain", "config"]);
   const {
-    register,
     hasEditRights,
     isPublicBrain,
     isOwnedByCurrentUser,
     isDefaultBrain,
     isSettingAsDefault,
     setAsDefaultBrainHandler,
   } = props;
+  const { register } = useBrainFormState();
+
   const { brainStatusOptions, brainTypeOptions } = useGeneralInformation();
 
   return (

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/ModelSelection/ModelSelection.tsx

@@ -1,34 +1,23 @@
 import { UUID } from "crypto";
-import { UseFormRegister } from "react-hook-form";
 import { useTranslation } from "react-i18next";
 
 import Field from "@/lib/components/ui/Field";
 import { defineMaxTokens } from "@/lib/helpers/defineMaxTokens";
-import { BrainConfig } from "@/lib/types/brainConfig";
 import { SaveButton } from "@/shared/SaveButton";
 
+import { useBrainFormState } from "../../hooks/useBrainFormState";
+
 type ModelSelectionProps = {
   brainId: UUID;
-  temperature: number;
-  maxTokens: number;
-  model: "gpt-3.5-turbo" | "gpt-3.5-turbo-16k";
   handleSubmit: (checkDirty: boolean) => Promise<void>;
-  register: UseFormRegister<BrainConfig>;
   hasEditRights: boolean;
   accessibleModels: string[];
 };
 
 export const ModelSelection = (props: ModelSelectionProps): JSX.Element => {
+  const { model, maxTokens, temperature, register } = useBrainFormState();
   const { t } = useTranslation(["translation", "brain", "config"]);
-  const {
-    handleSubmit,
-    register,
-    temperature,
-    maxTokens,
-    model,
-    hasEditRights,
-    accessibleModels,
-  } = props;
+  const { handleSubmit, hasEditRights, accessibleModels } = props;
 
   return (
     <>

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/PrivateAccessConfirmationModal/AccessConfirmationModal.tsx

@@ -2,24 +2,24 @@ import { useTranslation } from "react-i18next";
 
 import Button from "@/lib/components/ui/Button";
 import { Modal } from "@/lib/components/ui/Modal";
-import { BrainAccessStatus } from "@/lib/context/BrainProvider/types";
+
+import { useBrainFormState } from "../../hooks/useBrainFormState";
 
 type AccessConfirmationModalProps = {
   opened: boolean;
   onClose: () => void;
   onConfirm: () => void;
   onCancel: () => void;
-  selectedStatus: BrainAccessStatus;
 };
 
 export const AccessConfirmationModal = ({
   opened,
   onClose,
   onConfirm,
   onCancel,
-  selectedStatus,
 }: AccessConfirmationModalProps): JSX.Element => {
   const { t } = useTranslation(["brain"]);
+  const { status: selectedStatus } = useBrainFormState();
 
   const isPrivateStatus = selectedStatus === "private";
 

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/PrivateAccessConfirmationModal/hooks/useAccessConfirmationModal.ts

@@ -1,21 +1,11 @@
 import { useEffect, useState } from "react";
-import { UseFormResetField } from "react-hook-form";
 
-import { BrainConfig, BrainStatus } from "@/lib/types/brainConfig";
-
-type UseAccessConfirmationModalProps = {
-  status: BrainStatus;
-  setValue: (name: "status", value: "private" | "public") => void;
-  isStatusDirty: boolean;
-  resetField: UseFormResetField<BrainConfig>;
-};
+import { useBrainFormState } from "../../../hooks/useBrainFormState";
 
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
-export const useAccessConfirmationModal = ({
-  status,
-  isStatusDirty,
-  resetField,
-}: UseAccessConfirmationModalProps) => {
+export const useAccessConfirmationModal = () => {
+  const { dirtyFields, resetField, status } = useBrainFormState();
+  const isStatusDirty = Boolean(dirtyFields.status);
   const [isAccessModalOpened, setIsAccessModalOpened] = useState(false);
 
   useEffect(() => {

frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useBrainFormState.ts

@@ -1,18 +1,19 @@
-import { UUID } from "crypto";
-import { useForm } from "react-hook-form";
+/* eslint-disable complexity */
+
+import { useCallback, useEffect } from "react";
+import { useFormContext } from "react-hook-form";
 
-import { defaultBrainConfig } from "@/lib/config/defaultBrainConfig";
 import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
+import { Brain } from "@/lib/context/BrainProvider/types";
+import { useUrlBrain } from "@/lib/hooks/useBrainIdFromUrl";
 import { BrainConfig } from "@/lib/types/brainConfig";
 
 import { useBrainFetcher } from "../../../hooks/useBrainFetcher";
 
-type UseBrainFormStateProps = {
-  brainId: UUID;
-};
-
 // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
-export const useBrainFormState = ({ brainId }: UseBrainFormStateProps) => {
+export const useBrainFormState = () => {
+  const { brainId } = useUrlBrain();
+
   const { defaultBrainId } = useBrainContext();
 
   const {
@@ -23,10 +24,9 @@ export const useBrainFormState = ({ brainId }: UseBrainFormStateProps) => {
     reset,
     resetField,
     formState: { dirtyFields },
-  } = useForm<BrainConfig>({
-    defaultValues: { ...defaultBrainConfig, status: undefined },
-  });
-  const { brain } = useBrainFetcher({
+  } = useFormContext<BrainConfig>();
+
+  const { brain, refetchBrain } = useBrainFetcher({
     brainId,
   });
 
@@ -38,8 +38,46 @@ export const useBrainFormState = ({ brainId }: UseBrainFormStateProps) => {
   const maxTokens = watch("maxTokens");
   const status = watch("status");
 
+  const updateFormValues = useCallback(() => {
+    if (brain === undefined) {
+      return;
+    }
+
+    for (const key in brain) {
+      const brainKey = key as keyof Brain;
+      if (!(key in brain)) {
+        return;
+      }
+
+      if (brainKey === "max_tokens" && brain["max_tokens"] !== undefined) {
+        setValue("maxTokens", brain["max_tokens"]);
+        continue;
+      }
+
+      if (brainKey === "openai_api_key") {
+        setValue("openAiKey", brain["openai_api_key"] ?? "");
+        continue;
+      }
+
+      // @ts-expect-error bad type inference from typescript
+      // eslint-disable-next-line
+      if (Boolean(brain[key])) setValue(key, brain[key]);
+    }
+
+    setTimeout(() => {
+      if (brain.model !== undefined && brain.model !== null) {
+        setValue("model", brain.model);
+      }
+    }, 50);
+  }, [brain, setValue]);
+
+  useEffect(() => {
+    updateFormValues();
+  }, [brain, updateFormValues]);
+
   return {
     brain,
+    brainId,
     model,
     temperature,
     maxTokens,
@@ -53,5 +91,6 @@ export const useBrainFormState = ({ brainId }: UseBrainFormStateProps) => {
     setValue,
     reset,
     resetField,
+    refetchBrain,
   };
 };