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

Issue: https://github.com/StanGirard/quivr/issues/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
This commit is contained in:
Mamadou DICKO 2023-11-16 18:57:12 +01:00 committed by GitHub
parent f54c2a19c1
commit 9522d6b71f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
13 changed files with 126 additions and 202 deletions

View File

@ -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;

View File

@ -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>
);
};

View File

@ -1,7 +1,6 @@
/* 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";
@ -9,13 +8,12 @@ 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,7 +27,6 @@ export const GeneralInformation = (
): JSX.Element => {
const { t } = useTranslation(["translation", "brain", "config"]);
const {
register,
hasEditRights,
isPublicBrain,
isOwnedByCurrentUser,
@ -37,6 +34,8 @@ export const GeneralInformation = (
isSettingAsDefault,
setAsDefaultBrainHandler,
} = props;
const { register } = useBrainFormState();
const { brainStatusOptions, brainTypeOptions } = useGeneralInformation();
return (

View File

@ -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 (
<>

View File

@ -2,14 +2,14 @@ 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 = ({
@ -17,9 +17,9 @@ export const AccessConfirmationModal = ({
onClose,
onConfirm,
onCancel,
selectedStatus,
}: AccessConfirmationModalProps): JSX.Element => {
const { t } = useTranslation(["brain"]);
const { status: selectedStatus } = useBrainFormState();
const isPrivateStatus = selectedStatus === "private";

View File

@ -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(() => {

View File

@ -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,
};
};

View File

@ -1,39 +1,17 @@
/* eslint-disable max-lines */
import axios from "axios";
import { UUID } from "crypto";
import { useEffect, useState } from "react";
import {
UseFormGetValues,
UseFormRegister,
UseFormReset,
UseFormResetField,
UseFormSetValue,
} from "react-hook-form";
import { useTranslation } from "react-i18next";
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
import { usePromptApi } from "@/lib/api/prompt/usePromptApi";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { useToast } from "@/lib/hooks";
import { BrainConfig } from "@/lib/types/brainConfig";
type DirtyFields<T> = {
[K in keyof T]?: T[K] extends object
? DirtyFields<T[K]>
: boolean | undefined;
};
import { useBrainFormState } from "./useBrainFormState";
export type UsePromptProps = {
brainId: UUID;
getValues: UseFormGetValues<BrainConfig>;
promptId: string | undefined;
register: UseFormRegister<BrainConfig>;
reset: UseFormReset<BrainConfig>;
setValue: UseFormSetValue<BrainConfig>;
setIsUpdating: (isUpdating: boolean) => void;
resetField: UseFormResetField<BrainConfig>;
updateFormValues: () => void;
dirtyFields: DirtyFields<BrainConfig>;
};
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
@ -44,19 +22,19 @@ export const usePrompt = (props: UsePromptProps) => {
const { getPrompt, updatePrompt, createPrompt } = usePromptApi();
const [isRemovingPrompt, setIsRemovingPrompt] = useState(false);
const { fetchAllBrains } = useBrainContext();
const {
brainId,
dirtyFields,
getValues,
promptId,
register,
reset,
setValue,
resetField,
updateFormValues,
setIsUpdating,
} = props;
promptId,
refetchBrain,
brainId,
} = useBrainFormState();
const { setIsUpdating } = props;
const [currentPromptId, setCurrentPromptId] = useState<string | undefined>(
promptId
@ -82,6 +60,9 @@ export const usePrompt = (props: UsePromptProps) => {
}, [currentPromptId]);
const removeBrainPrompt = async () => {
if (brainId === undefined) {
return;
}
try {
setIsRemovingPrompt(true);
await updateBrain(brainId, {
@ -92,7 +73,7 @@ export const usePrompt = (props: UsePromptProps) => {
content: "",
});
reset();
void updateFormValues();
refetchBrain();
publish({
variant: "success",
text: t("promptRemoved", { ns: "config" }),
@ -141,6 +122,10 @@ export const usePrompt = (props: UsePromptProps) => {
return;
}
if (brainId === undefined) {
return;
}
try {
if (promptId === "" || promptId === undefined) {
otherConfigs["prompt_id"] = (
@ -149,13 +134,13 @@ export const usePrompt = (props: UsePromptProps) => {
content: prompt.content,
})
).id;
console.log("OTHER CONFIGS", otherConfigs);
await updateBrain(brainId, {
...otherConfigs,
max_tokens,
openai_api_key,
});
void updateFormValues();
refetchBrain();
} else {
await Promise.all([
updateBrain(brainId, {

View File

@ -1,13 +1,11 @@
/* eslint-disable complexity */
/* eslint-disable max-lines */
import axios from "axios";
import { UUID } from "crypto";
import { useCallback, useEffect, useRef, useState } from "react";
import { useEffect, useRef, useState } from "react";
import { useTranslation } from "react-i18next";
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { Brain } from "@/lib/context/BrainProvider/types";
import { defineMaxTokens } from "@/lib/helpers/defineMaxTokens";
import { getAccessibleModels } from "@/lib/helpers/getAccessibleModels";
import { useToast } from "@/lib/hooks";
@ -33,66 +31,20 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
const { userData } = useUserData();
const {
brain,
dirtyFields,
getValues,
maxTokens,
promptId,
register,
reset,
setValue,
openAiKey,
model,
temperature,
status,
isDefaultBrain,
resetField,
} = useBrainFormState({
brainId,
});
} = useBrainFormState();
const accessibleModels = getAccessibleModels({
openAiKey,
userData,
});
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]);
useEffect(() => {
setValue("maxTokens", Math.min(maxTokens, defineMaxTokens(model)));
}, [maxTokens, model, setValue]);
@ -196,25 +148,12 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
return {
handleSubmit,
register,
setAsDefaultBrainHandler,
setValue,
brain,
model,
temperature,
maxTokens,
isUpdating,
isSettingAsDefault,
isDefaultBrain,
formRef,
promptId,
accessibleModels,
status,
dirtyFields,
resetField,
updateFormValues,
reset,
getValues,
setIsUpdating,
};
};

View File

@ -1,4 +1,4 @@
import { useQuery } from "@tanstack/react-query";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { UUID } from "crypto";
import { useRouter } from "next/navigation";
@ -12,6 +12,7 @@ type UseBrainFetcherProps = {
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const useBrainFetcher = ({ brainId }: UseBrainFetcherProps) => {
const { getBrain } = useBrainApi();
const queryClient = useQueryClient();
const router = useRouter();
const fetchBrain = async () => {
@ -32,7 +33,14 @@ export const useBrainFetcher = ({ brainId }: UseBrainFetcherProps) => {
enabled: brainId !== undefined,
});
const invalidateBrainQuery = () => {
void queryClient.invalidateQueries({
queryKey: [getBrainDataKey(brainId!)],
});
};
return {
brain,
refetchBrain: invalidateBrainQuery,
};
};

View File

@ -39,11 +39,9 @@ export const AddBrainConfig = ({
isPending,
pickPublicPrompt,
accessibleModels,
status,
isPublicAccessConfirmationModalOpened,
onCancelPublicAccess,
onConfirmPublicAccess,
brainType,
register,
handleSubmit,
} = useAddBrainConfig();
@ -100,7 +98,6 @@ export const AddBrainConfig = ({
<Radio
items={brainStatusOptions}
label={t("brain_status_label", { ns: "brain" })}
value={status}
className="flex-1 justify-between w-[50%]"
{...register("status")}
/>
@ -111,7 +108,6 @@ export const AddBrainConfig = ({
<Radio
items={knowledgeSourceOptions}
label={t("knowledge_source_label", { ns: "brain" })}
value={brainType}
className="flex-1 justify-between w-[50%]"
{...register("brain_type")}
/>

View File

@ -202,11 +202,9 @@ export const useAddBrainConfig = () => {
isPending,
accessibleModels,
pickPublicPrompt,
status,
isPublicAccessConfirmationModalOpened,
onConfirmPublicAccess,
onCancelPublicAccess,
brainType,
register,
};
};

View File

@ -1,4 +1,4 @@
import { DetailedHTMLProps, forwardRef, InputHTMLAttributes } from "react";
import { forwardRef } from "react";
import { cn } from "@/lib/utils";
@ -7,11 +7,7 @@ type RadioItem = {
label: string;
};
interface RadioProps
extends DetailedHTMLProps<
InputHTMLAttributes<HTMLInputElement>,
HTMLInputElement
> {
interface RadioProps {
name: string;
items: RadioItem[];
label?: string;
@ -19,7 +15,7 @@ interface RadioProps
}
export const Radio = forwardRef<HTMLInputElement, RadioProps>(
({ items, label, className, value, ...props }, ref) => (
({ items, label, className, ...props }, ref) => (
<div className={cn("flex flex-col", className)}>
{label !== undefined && (
<label className="text-sm font-medium leading-6 mb-2">{label}</label>
@ -32,7 +28,6 @@ export const Radio = forwardRef<HTMLInputElement, RadioProps>(
type="radio"
className="form-radio h-4 w-4 text-indigo-600 border-indigo-600"
value={item.value}
checked={value === item.value}
{...props}
/>
<label className="ml-2" htmlFor={item.value}>