StanGirard/quivr
1
1
Fork 0
mirror of https://github.com/StanGirard/quivr.git synced 2024-11-23 12:26:03 +03:00
Code Issues Packages Projects Releases Wiki Activity

feat: allow setting public brain status to private (#1258)

Browse Source 
* feat: refetch brains list on when new brain is added

* feat: update BrainConfig type

* feat: update useSettingsTab add usebrainFormState and useSettings tab

* feat: add <PrivateAccessConfirmationModal/> modal

* feat: update translations

* feat: handle brain status change to private

* feat: validate chat access

* test: fix failaing tests and remove deprecated
This commit is contained in:
Mamadou DICKO 2023-09-26 10:35:52 +02:00 committed by GitHub
parent 4b88c89814
commit a4a2d769b3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
25 changed files with 582 additions and 282 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
backend/models/databases/supabase/brains.py
View File

@ -160,6 +160,17 @@ class Brain(Repository):
return results
def delete_brain_subscribers(self, brain_id: UUID):
results = (
self.db.table("brains_users")
.delete()
.match({"brain_id": str(brain_id)})
.match({"rights": "Viewer"})
.execute()
).data
return results
def delete_brain(self, brain_id: str):
results = (
self.db.table("brains").delete().match({"brain_id": brain_id}).execute()

10
backend/repository/brain/delete_brain_users.py Normal file
View File

@ -0,0 +1,10 @@
from uuid import UUID
from models.settings import get_supabase_db
def delete_brain_users(brain_id: UUID) -> None:
supabase_db = get_supabase_db()
supabase_db.delete_brain_subscribers(
brain_id=brain_id,
)

2
backend/routes/authorizations/brain_authorization.py
View File

@ -59,7 +59,7 @@ def validate_brain_authorization(
user_brain = get_brain_for_user(user_id, brain_id)
if user_brain is None:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
status_code=status.HTTP_403_FORBIDDEN,
detail="You don't have permission for this brain",
)

23
backend/routes/brain_routes.py
View File

@ -21,6 +21,7 @@ from repository.brain import (
set_as_default_brain_for_user,
update_brain_by_id,
)
from repository.brain.delete_brain_users import delete_brain_users
from repository.brain.get_public_brains import get_public_brains
from repository.prompt import delete_prompt_by_id, get_prompt_by_id
@ -182,27 +183,31 @@ async def create_brain_endpoint(
)
async def update_brain_endpoint(
brain_id: UUID,
input_brain: BrainUpdatableProperties,
brain_to_update: BrainUpdatableProperties,
):
"""
Update an existing brain with new brain configuration
"""
# Remove prompt if it is private and no longer used by brain
if input_brain.prompt_id is None:
existing_brain = get_brain_details(brain_id)
if existing_brain is None:
raise HTTPException(
status_code=404,
detail="Brain not found",
)
existing_brain = get_brain_details(brain_id)
if existing_brain is None:
raise HTTPException(
status_code=404,
detail="Brain not found",
)
if brain_to_update.prompt_id is None:
prompt_id = existing_brain.prompt_id
if prompt_id is not None:
prompt = get_prompt_by_id(prompt_id)
if prompt is not None and prompt.status == "private":
delete_prompt_by_id(prompt_id)
update_brain_by_id(brain_id, input_brain)
if brain_to_update.status == "private" and existing_brain.status == "public":
delete_brain_users(brain_id)
update_brain_by_id(brain_id, brain_to_update)
return {"message": f"Brain {brain_id} has been updated."}

9
backend/routes/chat_routes.py
View File

@ -35,6 +35,9 @@ from repository.chat.get_chat_history_with_notifications import (
from repository.notification.remove_chat_notifications import remove_chat_notifications
from repository.user_identity import get_user_identity
from routes.authorizations.brain_authorization import has_brain_authorization
from routes.authorizations.types import RoleEnum
chat_router = APIRouter()
@ -167,6 +170,9 @@ async def create_chat_handler(
Depends(
AuthBearer(),
),
Depends(
has_brain_authorization([RoleEnum.Viewer, RoleEnum.Owner, RoleEnum.Editor])
),
],
tags=["Chat"],
)
@ -255,6 +261,9 @@ async def create_question_handler(
Depends(
AuthBearer(),
),
Depends(
has_brain_authorization([RoleEnum.Viewer, RoleEnum.Owner, RoleEnum.Editor])
),
],
tags=["Chat"],
)

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

@ -8,11 +8,14 @@ import Button from "@/lib/components/ui/Button";
import { Chip } from "@/lib/components/ui/Chip";
import { Divider } from "@/lib/components/ui/Divider";
import Field from "@/lib/components/ui/Field";
import { Radio } from "@/lib/components/ui/Radio";
import { TextArea } from "@/lib/components/ui/TextArea";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { defineMaxTokens } from "@/lib/helpers/defineMaxTokens";
import { SaveButton } from "@/shared/SaveButton";
import { PrivateAccessConfirmationModal } from "./components/PrivateAccessConfirmationModal/PrivateAccessConfirmationModal";
import { usePrivateAccessConfirmationModal } from "./components/PrivateAccessConfirmationModal/hooks/usePrivateAccessConfirmationModal";
import { PublicPrompts } from "./components/PublicPrompts/PublicPrompts";
import { useSettingsTab } from "./hooks/useSettingsTab";
import { getBrainPermissions } from "../../utils/getBrainPermissions";
@ -39,166 +42,198 @@ export const SettingsTab = ({ brainId }: SettingsTabProps): JSX.Element => {
pickPublicPrompt,
removeBrainPrompt,
accessibleModels,
brainStatusOptions,
status,
setValue,
dirtyFields,
} = useSettingsTab({ brainId });
const { onCancel, isPrivateAccessModalOpened, closeModal } =
usePrivateAccessConfirmationModal({
status,
setValue,
isStatusDirty: Boolean(dirtyFields.status),
});
const { allBrains } = useBrainContext();
const { hasEditRights, isPublicBrain } = getBrainPermissions({
brainId,
userAccessibleBrains: allBrains,
});
const { hasEditRights, isOwnedByCurrentUser, isPublicBrain } =
getBrainPermissions({
brainId,
userAccessibleBrains: allBrains,
});
return (
<form
onSubmit={(e) => {
e.preventDefault();
void handleSubmit(true);
}}
className="my-10 mb-0 flex flex-col items-center gap-2"
ref={formRef}
>
<div className="flex flex-row flex-1 justify-between w-full items-end">
<div>
<Field
label={t("brainName", { ns: "brain" })}
placeholder={t("brainNamePlaceholder", { ns: "brain" })}
autoComplete="off"
className="flex-1"
required
disabled={!hasEditRights}
{...register("name")}
/>
</div>
<>
<form
onSubmit={(e) => {
e.preventDefault();
void handleSubmit(true);
}}
className="my-10 mb-0 flex flex-col items-center gap-2"
ref={formRef}
>
<div className="flex flex-row flex-1 justify-between w-full items-end">
<div>
<Field
label={t("brainName", { ns: "brain" })}
placeholder={t("brainNamePlaceholder", { ns: "brain" })}
autoComplete="off"
className="flex-1"
required
disabled={!hasEditRights}
{...register("name")}
/>
</div>
<div className="mt-4">
<div className="flex flex-1 items-center flex-col">
{isPublicBrain && (
<Chip className="mb-3 bg-purple-600 text-white w-full">
{t("brain:public_brain_label")}
</Chip>
)}
{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">
{t("defaultBrain", { ns: "brain" })}
</div>
) : (
<Button
variant={"secondary"}
isLoading={isSettingAsDefault}
onClick={() => void setAsDefaultBrainHandler()}
type="button"
disabled={!hasEditRights}
>
{t("setDefaultBrain", { ns: "brain" })}
</Button>
)}
<div className="mt-4">
<div className="flex flex-1 items-center flex-col">
{isPublicBrain && !isOwnedByCurrentUser && (
<Chip className="mb-3 bg-purple-600 text-white w-full">
{t("brain:public_brain_label")}
</Chip>
)}
{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">
{t("defaultBrain", { ns: "brain" })}
</div>
) : (
<Button
variant={"secondary"}
isLoading={isSettingAsDefault}
onClick={() => void setAsDefaultBrainHandler()}
type="button"
disabled={!hasEditRights}
>
{t("setDefaultBrain", { ns: "brain" })}
</Button>
)}
</div>
</div>
</div>
</div>
<TextArea
label={t("brainDescription", { ns: "brain" })}
placeholder={t("brainDescriptionPlaceholder", { ns: "brain" })}
autoComplete="off"
className="flex-1 m-3"
disabled={!hasEditRights}
{...register("description")}
/>
<Divider text={t("modelSection", { ns: "config" })} />
<Field
label={t("openAiKeyLabel", { ns: "config" })}
placeholder={t("openAiKeyPlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("openAiKey")}
/>
<fieldset className="w-full flex flex-col mt-2">
<label className="flex-1 text-sm" htmlFor="model">
{t("modelLabel", { ns: "config" })}
</label>
<select
id="model"
disabled={!hasEditRights}
{...register("model")}
className="px-5 py-2 dark:bg-gray-700 bg-gray-200 rounded-md"
onChange={() => {
void handleSubmit(false); // Trigger form submission
}}
>
{accessibleModels.map((availableModel) => (
<option value={availableModel} key={availableModel}>
{availableModel}
</option>
))}
</select>
</fieldset>
<fieldset className="w-full flex mt-4">
<label className="flex-1" htmlFor="temp">
{t("temperature", { ns: "config" })}: {temperature}
</label>
<input
id="temp"
type="range"
min="0"
max="1"
step="0.01"
value={temperature}
disabled={!hasEditRights}
{...register("temperature")}
/>
</fieldset>
<fieldset className="w-full flex mt-4">
<label className="flex-1" htmlFor="tokens">
{t("maxTokens", { ns: "config" })}: {maxTokens}
</label>
<input
type="range"
min="10"
max={defineMaxTokens(model)}
value={maxTokens}
disabled={!hasEditRights}
{...register("maxTokens")}
/>
</fieldset>
<div className="flex w-full justify-end py-4">
<SaveButton disabled={!hasEditRights} handleSubmit={handleSubmit} />
</div>
<Divider text={t("customPromptSection", { ns: "config" })} />
{hasEditRights && <PublicPrompts onSelect={pickPublicPrompt} />}
<Field
label={t("promptName", { ns: "config" })}
placeholder={t("promptNamePlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("prompt.title")}
/>
<TextArea
label={t("promptContent", { ns: "config" })}
placeholder={t("promptContentPlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("prompt.content")}
/>
<div className="flex w-full justify-end py-4">
<SaveButton disabled={!hasEditRights} handleSubmit={handleSubmit} />
</div>
{promptId !== "" && (
<Button
disabled={isUpdating || !hasEditRights}
onClick={() => void removeBrainPrompt()}
>
{t("removePrompt", { ns: "config" })}
</Button>
)}
<div className="flex flex-row justify-end flex-1 w-full mt-8">
{isUpdating && <FaSpinner className="animate-spin" />}
{isUpdating && (
<span className="ml-2 text-sm">
{t("updatingBrainSettings", { ns: "config" })}
</span>
{isPublicBrain && isOwnedByCurrentUser && (
<div className="w-full mt-4">
<Radio
items={brainStatusOptions}
label={t("brain_status_label", { ns: "brain" })}
value={status}
className="flex-1 justify-between w-[50%]"
{...register("status")}
/>
</div>
)}
</div>
</form>
<TextArea
label={t("brainDescription", { ns: "brain" })}
placeholder={t("brainDescriptionPlaceholder", { ns: "brain" })}
autoComplete="off"
className="flex-1 m-3"
disabled={!hasEditRights}
{...register("description")}
/>
<Divider text={t("modelSection", { ns: "config" })} />
<Field
label={t("openAiKeyLabel", { ns: "config" })}
placeholder={t("openAiKeyPlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("openAiKey")}
/>
<fieldset className="w-full flex flex-col mt-2">
<label className="flex-1 text-sm" htmlFor="model">
{t("modelLabel", { ns: "config" })}
</label>
<select
id="model"
disabled={!hasEditRights}
{...register("model")}
className="px-5 py-2 dark:bg-gray-700 bg-gray-200 rounded-md"
onChange={() => {
void handleSubmit(false); // Trigger form submission
}}
>
{accessibleModels.map((availableModel) => (
<option value={availableModel} key={availableModel}>
{availableModel}
</option>
))}
</select>
</fieldset>
<fieldset className="w-full flex mt-4">
<label className="flex-1" htmlFor="temp">
{t("temperature", { ns: "config" })}: {temperature}
</label>
<input
id="temp"
type="range"
min="0"
max="1"
step="0.01"
value={temperature}
disabled={!hasEditRights}
{...register("temperature")}
/>
</fieldset>
<fieldset className="w-full flex mt-4">
<label className="flex-1" htmlFor="tokens">
{t("maxTokens", { ns: "config" })}: {maxTokens}
</label>
<input
type="range"
min="10"
max={defineMaxTokens(model)}
value={maxTokens}
disabled={!hasEditRights}
{...register("maxTokens")}
/>
</fieldset>
<div className="flex w-full justify-end py-4">
<SaveButton disabled={!hasEditRights} handleSubmit={handleSubmit} />
</div>
<Divider text={t("customPromptSection", { ns: "config" })} />
{hasEditRights && <PublicPrompts onSelect={pickPublicPrompt} />}
<Field
label={t("promptName", { ns: "config" })}
placeholder={t("promptNamePlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("prompt.title")}
/>
<TextArea
label={t("promptContent", { ns: "config" })}
placeholder={t("promptContentPlaceholder", { ns: "config" })}
autoComplete="off"
className="flex-1"
disabled={!hasEditRights}
{...register("prompt.content")}
/>
<div className="flex w-full justify-end py-4">
<SaveButton disabled={!hasEditRights} handleSubmit={handleSubmit} />
</div>
{promptId !== "" && (
<Button
disabled={isUpdating || !hasEditRights}
onClick={() => void removeBrainPrompt()}
>
{t("removePrompt", { ns: "config" })}
</Button>
)}
<div className="flex flex-row justify-end flex-1 w-full mt-8">
{isUpdating && <FaSpinner className="animate-spin" />}
{isUpdating && (
<span className="ml-2 text-sm">
{t("updatingBrainSettings", { ns: "config" })}
</span>
)}
</div>
</form>
<PrivateAccessConfirmationModal
opened={isPrivateAccessModalOpened}
onClose={onCancel}
onCancel={onCancel}
onConfirm={closeModal}
/>
</>
);
};

42
frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/PrivateAccessConfirmationModal/PrivateAccessConfirmationModal.tsx Normal file
View File

@ -0,0 +1,42 @@
import { useTranslation } from "react-i18next";
import Button from "@/lib/components/ui/Button";
import { Modal } from "@/lib/components/ui/Modal";
type PrivateAccessConfirmationModalProps = {
opened: boolean;
onClose: () => void;
onConfirm: () => void;
onCancel: () => void;
};
export const PrivateAccessConfirmationModal = ({
opened,
onClose,
onConfirm,
onCancel,
}: PrivateAccessConfirmationModalProps): JSX.Element => {
const { t } = useTranslation(["brain"]);
return (
<Modal isOpen={opened} setOpen={onClose} CloseTrigger={<div />}>
<div
dangerouslySetInnerHTML={{
__html: t("set_brain_status_to_private_modal_title"),
}}
/>
<div
dangerouslySetInnerHTML={{
__html: t("set_brain_status_to_private_modal_description"),
}}
/>
<div className="flex flex-row justify-between pt-10 px-10">
<Button type="button" onClick={onConfirm} variant="secondary">
{t("confirm_set_brain_status_to_private")}
</Button>
<Button type="button" onClick={onCancel}>
{t("cancel_set_brain_status_to_private")}
</Button>
</div>
</Modal>
);
};

40
frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/PrivateAccessConfirmationModal/hooks/usePrivateAccessConfirmationModal.ts Normal file
View File

@ -0,0 +1,40 @@
import { useEffect, useState } from "react";
import { BrainStatus } from "@/lib/types/brainConfig";
type UsePrivateAccessModalProps = {
status: BrainStatus;
setValue: (name: "status", value: "private" | "public") => void;
isStatusDirty: boolean;
};
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const usePrivateAccessConfirmationModal = ({
status,
setValue,
isStatusDirty,
}: UsePrivateAccessModalProps) => {
const [isPrivateAccessModalOpened, setIsPrivateAccessModalOpened] =
useState(false);
useEffect(() => {
if (status === "private" && isStatusDirty) {
setIsPrivateAccessModalOpened(true);
}
}, [status]);
const closeModal = () => {
setIsPrivateAccessModalOpened(false);
};
const onCancel = () => {
closeModal();
setValue("status", "public");
};
return {
isPrivateAccessModalOpened,
closeModal,
onCancel,
};
};

1
frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/components/PrivateAccessConfirmationModal/index.ts Normal file
View File

@ -0,0 +1 @@
export * from "./PrivateAccessConfirmationModal";

55
frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useBrainFormState.ts Normal file
View File

@ -0,0 +1,55 @@
import { UUID } from "crypto";
import { useForm } from "react-hook-form";
import { defaultBrainConfig } from "@/lib/config/defaultBrainConfig";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
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) => {
const { defaultBrainId } = useBrainContext();
const {
register,
getValues,
watch,
setValue,
reset,
formState: { dirtyFields },
} = useForm<BrainConfig>({
defaultValues: { ...defaultBrainConfig, status: undefined },
});
const { brain } = useBrainFetcher({
brainId,
});
const isDefaultBrain = defaultBrainId === brainId;
const promptId = watch("prompt_id");
const openAiKey = watch("openAiKey");
const model = watch("model");
const temperature = watch("temperature");
const maxTokens = watch("maxTokens");
const status = watch("status");
return {
brain,
model,
temperature,
maxTokens,
isDefaultBrain,
promptId,
openAiKey,
dirtyFields,
status,
register,
getValues,
setValue,
reset,
};
};

84
frontend/app/brains-management/[brainId]/components/BrainManagementTabs/components/SettingsTab/hooks/useSettingsTab.ts
View File

@ -1,24 +1,21 @@
/* eslint-disable complexity */
/* eslint-disable max-lines */
import { useQuery } from "@tanstack/react-query";
import axios from "axios";
import { UUID } from "crypto";
import { useCallback, useEffect, useRef, useState } from "react";
import { useForm } from "react-hook-form";
import { useTranslation } from "react-i18next";
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
import { usePromptApi } from "@/lib/api/prompt/usePromptApi";
import { USER_DATA_KEY } from "@/lib/api/user/config";
import { useUserApi } from "@/lib/api/user/useUserApi";
import { defaultBrainConfig } from "@/lib/config/defaultBrainConfig";
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";
import { useUserData } from "@/lib/hooks/useUserData";
import { BrainStatus } from "@/lib/types/brainConfig";
import { useBrainFetcher } from "../../../hooks/useBrainFetcher";
import { useBrainFormState } from "./useBrainFormState";
import { validateOpenAIKey } from "../utils/validateOpenAIKey";
type UseSettingsTabProps = {
@ -33,48 +30,41 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
const { publish } = useToast();
const formRef = useRef<HTMLFormElement>(null);
const { setAsDefaultBrain, updateBrain } = useBrainApi();
const { fetchAllBrains, fetchDefaultBrain, defaultBrainId } =
useBrainContext();
const { fetchAllBrains, fetchDefaultBrain } = useBrainContext();
const { getPrompt, updatePrompt, createPrompt } = usePromptApi();
const { getUser } = useUserApi();
const { data: userData } = useQuery({
queryKey: [USER_DATA_KEY],
queryFn: getUser,
});
const defaultValues = {
...defaultBrainConfig,
name: "",
description: "",
setDefault: false,
prompt_id: "",
prompt: {
title: "",
content: "",
},
};
const { userData } = useUserData();
const {
register,
brain,
dirtyFields,
getValues,
watch,
setValue,
maxTokens,
promptId,
register,
reset,
formState: { dirtyFields },
} = useForm({
defaultValues,
});
const { brain } = useBrainFetcher({
setValue,
openAiKey,
model,
temperature,
status,
isDefaultBrain,
} = useBrainFormState({
brainId,
});
const isDefaultBrain = defaultBrainId === brainId;
const promptId = watch("prompt_id");
const openAiKey = watch("openAiKey");
const model = watch("model");
const temperature = watch("temperature");
const maxTokens = watch("maxTokens");
const brainStatusOptions: {
label: string;
value: BrainStatus;
}[] = [
{
label: t("private_brain_label", { ns: "brain" }),
value: "private",
},
{
label: t("public_brain_label", { ns: "brain" }),
value: "public",
},
];
const accessibleModels = getAccessibleModels({
openAiKey,
@ -138,7 +128,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
}, [formRef.current]);
const fetchPrompt = async () => {
if (promptId === "") {
if (promptId === "" || promptId === undefined) {
return;
}
@ -211,7 +201,7 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
const promptHandler = async () => {
const { prompt } = getValues();
if (dirtyFields["prompt"]) {
if (dirtyFields["prompt"] && promptId !== undefined) {
await updatePrompt(promptId, {
title: prompt.title,
content: prompt.content,
@ -348,18 +338,22 @@ export const useSettingsTab = ({ brainId }: UseSettingsTabProps) => {
return {
handleSubmit,
register,
removeBrainPrompt,
pickPublicPrompt,
setAsDefaultBrainHandler,
setValue,
brain,
model,
temperature,
maxTokens,
isUpdating,
setAsDefaultBrainHandler,
isSettingAsDefault,
isDefaultBrain,
formRef,
promptId,
removeBrainPrompt,
pickPublicPrompt,
accessibleModels,
brainStatusOptions,
status,
dirtyFields,
};
};

19
frontend/app/chat/[chatId]/__tests__/page.test.tsx
View File

@ -1,3 +1,4 @@
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { render } from "@testing-library/react";
import { describe, expect, it, vi } from "vitest";
@ -16,6 +17,8 @@ import {
import SelectedChatPage from "../page";
const queryClient = new QueryClient();
vi.mock("@/lib/context/ChatProvider/ChatProvider", () => ({
ChatContext: ChatContextMock,
ChatProvider: ChatProviderMock,
@ -55,13 +58,15 @@ vi.mock("@tanstack/react-query", async () => {
describe("Chat page", () => {
it("should render chat page correctly", () => {
const { getByTestId } = render(
<ChatProviderMock>
<SupabaseProviderMock>
<BrainProviderMock>
<SelectedChatPage />,
</BrainProviderMock>
</SupabaseProviderMock>
</ChatProviderMock>
<QueryClientProvider client={queryClient}>
<ChatProviderMock>
<SupabaseProviderMock>
<BrainProviderMock>
<SelectedChatPage />,
</BrainProviderMock>
</SupabaseProviderMock>
</ChatProviderMock>
</QueryClientProvider>
);
expect(getByTestId("chat-page")).toBeDefined();

41
frontend/app/chat/[chatId]/hooks/useHandleStream.ts Normal file
View File

@ -0,0 +1,41 @@
import { useChatContext } from "@/lib/context";
import { ChatMessage } from "../types";
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const useHandleStream = () => {
const { updateStreamingHistory } = useChatContext();
const handleStream = async (
reader: ReadableStreamDefaultReader<Uint8Array>
): Promise<void> => {
const decoder = new TextDecoder("utf-8");
const handleStreamRecursively = async () => {
const { done, value } = await reader.read();
if (done) {
return;
}
const dataStrings = decoder
.decode(value)
.trim()
.split("data: ")
.filter(Boolean);
dataStrings.forEach((data) => {
const parsedData = JSON.parse(data) as ChatMessage;
updateStreamingHistory(parsedData);
});
await handleStreamRecursively();
};
await handleStreamRecursively();
};
return {
handleStream,
};
};

69
frontend/app/chat/[chatId]/hooks/useQuestion.ts
View File

@ -1,11 +1,10 @@
import axios from "axios";
import { useTranslation } from "react-i18next";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { useChatContext } from "@/lib/context/ChatProvider/hooks/useChatContext";
import { useFetch, useToast } from "@/lib/hooks";
import { ChatMessage, ChatQuestion } from "../types";
import { useHandleStream } from "./useHandleStream";
import { ChatQuestion } from "../types";
interface UseChatService {
addStreamQuestion: (
@ -16,43 +15,29 @@ interface UseChatService {
export const useQuestion = (): UseChatService => {
const { fetchInstance } = useFetch();
const { updateStreamingHistory } = useChatContext();
const { currentBrain } = useBrainContext();
const { t } = useTranslation(["chat"]);
const { publish } = useToast();
const { handleStream } = useHandleStream();
const handleStream = async (
reader: ReadableStreamDefaultReader<Uint8Array>
): Promise<void> => {
const decoder = new TextDecoder("utf-8");
const handleStreamRecursively = async () => {
const { done, value } = await reader.read();
if (done) {
return;
}
const dataStrings = decoder
.decode(value)
.trim()
.split("data: ")
.filter(Boolean);
dataStrings.forEach((data) => {
try {
const parsedData = JSON.parse(data) as ChatMessage;
updateStreamingHistory(parsedData);
} catch (error) {
console.error(t("errorParsingData", { ns: "chat" }), error);
}
const handleFetchError = async (response: Response) => {
if (response.status === 429) {
publish({
variant: "danger",
text: t("tooManyRequests", { ns: "chat" }),
});
await handleStreamRecursively();
};
return;
}
await handleStreamRecursively();
const errorMessage = (await response.json()) as { detail: string };
publish({
variant: "danger",
text: errorMessage.detail,
});
return;
};
const addStreamQuestion = async (
@ -64,29 +49,29 @@ export const useQuestion = (): UseChatService => {
Accept: "text/event-stream",
};
const body = JSON.stringify(chatQuestion);
console.log("Calling API...");
try {
const response = await fetchInstance.post(
`/chat/${chatId}/question/stream?brain_id=${currentBrain?.id ?? ""}`,
body,
headers
);
if (!response.ok) {
void handleFetchError(response);
return;
}
if (response.body === null) {
throw new Error(t("resposeBodyNull", { ns: "chat" }));
}
console.log(t("receivedResponse"), response);
await handleStream(response.body.getReader());
} catch (error) {
if (axios.isAxiosError(error) && error.response?.status === 429) {
publish({
variant: "danger",
text: t("tooManyRequests", { ns: "chat" }),
});
}
console.error(t("errorCallingAPI", { ns: "chat" }), error);
publish({
variant: "danger",
text: String(error),
});
}
};

18
frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts
View File

@ -4,6 +4,7 @@ import { useEffect, useState } from "react";
import { useForm } from "react-hook-form";
import { useTranslation } from "react-i18next";
import { PUBLIC_BRAINS_KEY } from "@/lib/api/brain/config";
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
import { usePromptApi } from "@/lib/api/prompt/usePromptApi";
import { USER_DATA_KEY } from "@/lib/api/user/config";
@ -14,7 +15,7 @@ import { defineMaxTokens } from "@/lib/helpers/defineMaxTokens";
import { getAccessibleModels } from "@/lib/helpers/getAccessibleModels";
import { useToast } from "@/lib/hooks";
import { BrainStatus } from "@/lib/types/brainConfig";
import { useQuery } from "@tanstack/react-query";
import { useQuery, useQueryClient } from "@tanstack/react-query";
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const useAddBrainModal = () => {
@ -30,6 +31,7 @@ export const useAddBrainModal = () => {
setIsPublicAccessConfirmationModalOpened,
] = useState(false);
const { getUser } = useUserApi();
const queryClient = useQueryClient();
const brainStatusOptions: {
label: string;
@ -61,7 +63,14 @@ export const useAddBrainModal = () => {
},
};
const { register, getValues, reset, watch, setValue } = useForm({
const {
register,
getValues,
reset,
watch,
setValue,
formState: { dirtyFields },
} = useForm({
defaultValues,
});
@ -77,7 +86,7 @@ export const useAddBrainModal = () => {
});
useEffect(() => {
if (status === "public") {
if (status === "public" && dirtyFields.status) {
setIsPublicAccessConfirmationModalOpened(true);
}
}, [status]);
@ -145,6 +154,9 @@ export const useAddBrainModal = () => {
variant: "success",
text: t("brainCreated", { ns: "brain" }),
});
void queryClient.invalidateQueries({
queryKey: [PUBLIC_BRAINS_KEY],
});
} catch (err) {
if (axios.isAxiosError(err) && err.response?.status === 429) {
publish({

7
frontend/lib/config/defaultBrainConfig.ts
View File

@ -12,4 +12,11 @@ export const defaultBrainConfig: BrainConfig = {
supabaseUrl: undefined,
prompt_id: undefined,
status: "private",
prompt: {
title: "",
content: "",
},
name: "",
description: "",
setDefault: false,
};

28
frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts
View File

@ -33,14 +33,26 @@ export const useBrainProvider = () => {
);
const currentBrain = allBrains.find((brain) => brain.id === currentBrainId);
const fetchAllBrains = useCallback(async () => {
setIsFetchingBrains(true);
try {
const brains = await getBrains();
setAllBrains(brains);
} catch (error) {
console.error(error);
} finally {
setIsFetchingBrains(false);
}
}, [getBrains]);
const createBrainHandler = useCallback(
async (brain: CreateBrainInput): Promise<UUID | undefined> => {
const createdBrain = await createBrain(brain);
try {
setAllBrains((prevBrains) => [...prevBrains, createdBrain]);
setCurrentBrainId(createdBrain.id);
void track("BRAIN_CREATED");
void fetchAllBrains();
return createdBrain.id;
} catch {
@ -50,7 +62,7 @@ export const useBrainProvider = () => {
});
}
},
[createBrain, publish, track]
[createBrain, fetchAllBrains, publish, track]
);
const deleteBrainHandler = useCallback(
@ -68,18 +80,6 @@ export const useBrainProvider = () => {
[deleteBrain, publish, track]
);
const fetchAllBrains = useCallback(async () => {
setIsFetchingBrains(true);
try {
const brains = await getBrains();
setAllBrains(brains);
} catch (error) {
console.error(error);
} finally {
setIsFetchingBrains(false);
}
}, [getBrains]);
const fetchDefaultBrain = useCallback(async () => {
const userDefaultBrain = await getDefaultBrain();
if (userDefaultBrain !== undefined) {

18
frontend/lib/hooks/useUserData.ts Normal file
View File

@ -0,0 +1,18 @@
import { useQuery } from "@tanstack/react-query";
import { USER_DATA_KEY } from "../api/user/config";
import { useUserApi } from "../api/user/useUserApi";
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const useUserData = () => {
const { getUser } = useUserApi();
const { data: userData } = useQuery({
queryKey: [USER_DATA_KEY],
queryFn: getUser,
});
return {
userData,
};
};

11
frontend/lib/types/brainConfig.ts
View File

@ -1,5 +1,3 @@
import { UUID } from "crypto";
export const brainStatuses = ["private", "public"] as const;
export type BrainStatus = (typeof brainStatuses)[number];
@ -14,8 +12,15 @@ export type BrainConfig = {
anthropicKey?: string;
supabaseUrl?: string;
supabaseKey?: string;
prompt_id?: UUID;
prompt_id?: string;
status: BrainStatus;
prompt: {
title: string;
content: string;
};
name: string;
description: string;
setDefault: boolean;
};
export type BrainConfigContextType = {

7
frontend/public/locales/en/brain.json
View File

@ -39,5 +39,10 @@
"public_brain_subscribe_button_label":"Subscribe",
"public_brain_subscription_success_message":"You have successfully subscribed to the brain",
"public_brain_last_update_label":"Last update",
"public_brain_already_subscribed_button_label":"Subscribed"
"public_brain_already_subscribed_button_label":"Subscribed",
"set_brain_status_to_private_modal_title":"Are you sure you want to set this as <span class='text-purple-800'>Private</span>?<br/><br/>",
"set_brain_status_to_private_modal_description":"Every Quivr users won't be able to use this brain anymore and they won't see it in the brain library.",
"confirm_set_brain_status_to_private":"Yes, set as private",
"cancel_set_brain_status_to_private":"No, keep it public"
}

6
frontend/public/locales/es/brain.json
View File

@ -39,5 +39,9 @@
"public_brain_subscribe_button_label": "Suscribirse",
"public_brain_subscription_success_message": "Te has suscrito con éxito al cerebro",
"public_brain_last_update_label": "Última actualización",
"public_brain_already_subscribed_button_label": "Ya suscrito"
"public_brain_already_subscribed_button_label": "Ya suscrito",
"set_brain_status_to_private_modal_title": "¿Estás seguro de que quieres establecer esto como <span class='text-purple-800'>Privado</span>?<br/><br/>",
"set_brain_status_to_private_modal_description": "Los usuarios de Quivr ya no podrán utilizar este cerebro y no lo verán en la biblioteca de cerebros.",
"confirm_set_brain_status_to_private": "Sí, establecer como privado",
"cancel_set_brain_status_to_private": "No, mantenerlo público"
}

6
frontend/public/locales/fr/brain.json
View File

@ -39,5 +39,9 @@
"public_brain_subscribe_button_label": "S'abonner",
"public_brain_subscription_success_message": "Vous vous êtes abonné avec succès au cerveau",
"public_brain_last_update_label": "Dernière mise à jour",
"public_brain_already_subscribed_button_label": "Abonné"
"public_brain_already_subscribed_button_label": "Abonné",
"set_brain_status_to_private_modal_title": "Êtes-vous sûr de vouloir définir ceci comme <span class='text-purple-800'>Privé</span>?<br/><br/>",
"set_brain_status_to_private_modal_description": "Les utilisateurs de Quivr ne pourront plus utiliser ce cerveau et ne le verront plus dans la bibliothèque des cerveaux.",
"confirm_set_brain_status_to_private": "Oui, définir comme privé",
"cancel_set_brain_status_to_private": "Non, le laisser public"
}

6
frontend/public/locales/pt-br/brain.json
View File

@ -39,5 +39,9 @@
"public_brain_subscribe_button_label": "Inscrever-se",
"public_brain_subscription_success_message": "Você se inscreveu com sucesso no cérebro",
"public_brain_last_update_label": "Última atualização",
"public_brain_already_subscribed_button_label": "Inscrito"
"public_brain_already_subscribed_button_label": "Inscrito",
"set_brain_status_to_private_modal_title": "Tem a certeza de que deseja definir isto como <span class='text-purple-800'>Privado</span>?<br/><br/>",
"set_brain_status_to_private_modal_description": "Os utilizadores do Quivr não poderão mais utilizar este cérebro e não o verão na biblioteca de cérebros.",
"confirm_set_brain_status_to_private": "Sim, definir como privado",
"cancel_set_brain_status_to_private": "Não, mantê-lo público"
}

6
frontend/public/locales/ru/brain.json
View File

@ -39,5 +39,9 @@
"public_brain_subscribe_button_label": "Подписаться",
"public_brain_subscription_success_message": "Вы успешно подписались на мозг",
"public_brain_last_update_label": "Последнее обновление",
"public_brain_already_subscribed_button_label": "Вы уже подписаны"
"public_brain_already_subscribed_button_label": "Вы уже подписаны",
"set_brain_status_to_private_modal_title": "Вы уверены, что хотите установить это как <span class='text-purple-800'>Частное</span>?<br/><br/>",
"set_brain_status_to_private_modal_description": "Пользователи Quivr больше не смогут использовать этот мозг, и они не увидят его в библиотеке мозгов.",
"confirm_set_brain_status_to_private": "Да, установить как частное",
"cancel_set_brain_status_to_private": "Нет, оставить общедоступным"
}

6
frontend/public/locales/zh-cn/brain.json
View File

@ -39,5 +39,9 @@
"public_brain_subscribe_button_label": "订阅",
"public_brain_subscription_success_message": "Вы успешно подписались на мозг",
"public_brain_last_update_label": "Последнее обновление",
"public_brain_already_subscribed_button_label": "已订阅"
"public_brain_already_subscribed_button_label": "已订阅",
"set_brain_status_to_private_modal_title": "您确定要将此设置为<span class='text-purple-800'>私有</span>吗？<br/><br/>",
"set_brain_status_to_private_modal_description": "Quivr的用户将无法再使用此大脑，并且不会在大脑库中看到它。",
"confirm_set_brain_status_to_private": "是的，设为私有",
"cancel_set_brain_status_to_private": "不，保持为公开"
}