mirror of
https://github.com/StanGirard/quivr.git
synced 2024-09-11 13:15:41 +03:00
[Brain management] Add new fields to creation modal (#755)
* refactor(ModalConfig): move defineMaxTokensto helpers * feat(AddBrain): add new properties * feat(sdk): update createBrain * feat(sdk): add setAsDefaultBrain
This commit is contained in:
parent
10b0cce992
commit
046cc3fc1d
@ -2,7 +2,7 @@
|
||||
import { motion, MotionConfig } from "framer-motion";
|
||||
import { MdChevronRight } from "react-icons/md";
|
||||
|
||||
import { AddBrainModal } from "@/lib/components/AddBrainModal";
|
||||
import { AddBrainModal } from "@/lib/components/AddBrainModal/AddBrainModal";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
import { BrainListItem } from "./BrainListItem";
|
||||
|
@ -12,6 +12,7 @@ import {
|
||||
models,
|
||||
paidModels,
|
||||
} from "@/lib/context/BrainConfigProvider/types";
|
||||
import { defineMaxTokens } from "@/lib/helpers/defineMexTokens";
|
||||
|
||||
interface ModelConfigProps {
|
||||
register: UseFormRegister<BrainConfig>;
|
||||
@ -28,21 +29,7 @@ export const ModelConfig = ({
|
||||
temperature,
|
||||
maxTokens,
|
||||
}: ModelConfigProps): JSX.Element => {
|
||||
const defineMaxTokens = (model: Model | PaidModels): number => {
|
||||
//At the moment is evaluating only models from OpenAI
|
||||
switch (model) {
|
||||
case "gpt-3.5-turbo-0613":
|
||||
return 500;
|
||||
case "gpt-3.5-turbo-16k":
|
||||
return 2000;
|
||||
case "gpt-4":
|
||||
return 1000;
|
||||
case "gpt-4-0613":
|
||||
return 100;
|
||||
default:
|
||||
return 250;
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
return (
|
||||
<>
|
||||
|
@ -2,7 +2,7 @@
|
||||
import { renderHook } from "@testing-library/react";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import { Subscription } from "../brain";
|
||||
import { CreateBrainInput, Subscription } from "../brain";
|
||||
import { SubscriptionUpdatableProperties } from "../types";
|
||||
import { useBrainApi } from "../useBrainApi";
|
||||
|
||||
@ -57,11 +57,21 @@ describe("useBrainApi", () => {
|
||||
current: { createBrain },
|
||||
},
|
||||
} = renderHook(() => useBrainApi());
|
||||
const name = "Test Brain";
|
||||
await createBrain(name);
|
||||
|
||||
const brain: CreateBrainInput = {
|
||||
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 createBrain(brain);
|
||||
|
||||
expect(axiosPostMock).toHaveBeenCalledTimes(1);
|
||||
expect(axiosPostMock).toHaveBeenCalledWith("/brains/", { name });
|
||||
expect(axiosPostMock).toHaveBeenCalledWith("/brains/", brain);
|
||||
});
|
||||
|
||||
it("should call deleteBrain with the correct parameters", async () => {
|
||||
@ -179,4 +189,16 @@ describe("useBrainApi", () => {
|
||||
{ rights: "Viewer", email }
|
||||
);
|
||||
});
|
||||
|
||||
it("should call setAsDefaultBrain with correct brainId", async () => {
|
||||
const {
|
||||
result: {
|
||||
current: { setAsDefaultBrain },
|
||||
},
|
||||
} = renderHook(() => useBrainApi());
|
||||
const brainId = "123";
|
||||
await setAsDefaultBrain(brainId);
|
||||
expect(axiosPutMock).toHaveBeenCalledTimes(1);
|
||||
expect(axiosPutMock).toHaveBeenCalledWith(`/brains/${brainId}/default`);
|
||||
});
|
||||
});
|
||||
|
@ -28,12 +28,21 @@ 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 (
|
||||
name: string,
|
||||
brain: CreateBrainInput,
|
||||
axiosInstance: AxiosInstance
|
||||
): Promise<MinimalBrainForUser> => {
|
||||
return mapBackendMinimalBrainToMinimalBrain(
|
||||
(await axiosInstance.post<BackendMinimalBrainForUser>(`/brains/`, { name }))
|
||||
(await axiosInstance.post<BackendMinimalBrainForUser>(`/brains/`, brain))
|
||||
.data
|
||||
);
|
||||
};
|
||||
@ -117,3 +126,10 @@ export const updateBrainAccess = async (
|
||||
email: userEmail,
|
||||
});
|
||||
};
|
||||
|
||||
export const setAsDefaultBrain = async (
|
||||
brainId: string,
|
||||
axiosInstance: AxiosInstance
|
||||
): Promise<void> => {
|
||||
await axiosInstance.put(`/brains/${brainId}/default`);
|
||||
};
|
||||
|
@ -3,12 +3,14 @@ import { useAxios } from "@/lib/hooks";
|
||||
import {
|
||||
addBrainSubscriptions,
|
||||
createBrain,
|
||||
CreateBrainInput,
|
||||
deleteBrain,
|
||||
getBrain,
|
||||
getBrainDocuments,
|
||||
getBrains,
|
||||
getBrainUsers,
|
||||
getDefaultBrain,
|
||||
setAsDefaultBrain,
|
||||
Subscription,
|
||||
updateBrainAccess,
|
||||
} from "./brain";
|
||||
@ -21,7 +23,8 @@ export const useBrainApi = () => {
|
||||
return {
|
||||
getBrainDocuments: async (brainId: string) =>
|
||||
getBrainDocuments(brainId, axiosInstance),
|
||||
createBrain: async (name: string) => createBrain(name, axiosInstance),
|
||||
createBrain: async (brain: CreateBrainInput) =>
|
||||
createBrain(brain, axiosInstance),
|
||||
deleteBrain: async (id: string) => deleteBrain(id, axiosInstance),
|
||||
getDefaultBrain: async () => getDefaultBrain(axiosInstance),
|
||||
getBrains: async () => getBrains(axiosInstance),
|
||||
@ -37,5 +40,7 @@ export const useBrainApi = () => {
|
||||
userEmail: string,
|
||||
subscription: SubscriptionUpdatableProperties
|
||||
) => updateBrainAccess(brainId, userEmail, subscription, axiosInstance),
|
||||
setAsDefaultBrain: async (brainId: string) =>
|
||||
setAsDefaultBrain(brainId, axiosInstance),
|
||||
};
|
||||
};
|
||||
|
@ -1,90 +0,0 @@
|
||||
import axios from "axios";
|
||||
import { FormEvent, useState } from "react";
|
||||
import { MdAdd } from "react-icons/md";
|
||||
|
||||
import Button from "@/lib/components/ui/Button";
|
||||
import Field from "@/lib/components/ui/Field";
|
||||
import { Modal } from "@/lib/components/ui/Modal";
|
||||
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
|
||||
import { useToast } from "@/lib/hooks";
|
||||
|
||||
export const AddBrainModal = (): JSX.Element => {
|
||||
const [newBrainName, setNewBrainName] = useState("");
|
||||
const [isPending, setIsPending] = useState(false);
|
||||
const { publish } = useToast();
|
||||
const { createBrain } = useBrainContext();
|
||||
const [isShareModalOpen, setIsShareModalOpen] = useState(false);
|
||||
|
||||
const handleSubmit = async (e: FormEvent) => {
|
||||
e.preventDefault();
|
||||
if (newBrainName.trim() === "" || isPending) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
setIsPending(true);
|
||||
await createBrain(newBrainName);
|
||||
setNewBrainName("");
|
||||
setIsShareModalOpen(false);
|
||||
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 {
|
||||
setIsPending(false);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Modal
|
||||
Trigger={
|
||||
<Button variant={"secondary"}>
|
||||
Add New Brain
|
||||
<MdAdd className="text-xl" />
|
||||
</Button>
|
||||
}
|
||||
title="Add Brain"
|
||||
desc="Add a new brain"
|
||||
isOpen={isShareModalOpen}
|
||||
setOpen={setIsShareModalOpen}
|
||||
CloseTrigger={<div />}
|
||||
>
|
||||
<form
|
||||
onSubmit={(e) => void handleSubmit(e)}
|
||||
className="my-10 flex items-center gap-2"
|
||||
>
|
||||
<Field
|
||||
name="brainname"
|
||||
label="Enter a brain name"
|
||||
autoFocus
|
||||
placeholder="E.g. History notes"
|
||||
autoComplete="off"
|
||||
value={newBrainName}
|
||||
onChange={(e) => setNewBrainName(e.currentTarget.value)}
|
||||
className="flex-1"
|
||||
/>
|
||||
<Button isLoading={isPending} className="self-end" type="submit">
|
||||
Create
|
||||
<MdAdd className="text-xl" />
|
||||
</Button>
|
||||
</form>
|
||||
</Modal>
|
||||
);
|
||||
};
|
135
frontend/lib/components/AddBrainModal/AddBrainModal.tsx
Normal file
135
frontend/lib/components/AddBrainModal/AddBrainModal.tsx
Normal file
@ -0,0 +1,135 @@
|
||||
/* eslint-disable max-lines */
|
||||
import { MdAdd } from "react-icons/md";
|
||||
|
||||
import Button from "@/lib/components/ui/Button";
|
||||
import Field from "@/lib/components/ui/Field";
|
||||
import { Modal } from "@/lib/components/ui/Modal";
|
||||
import { models, paidModels } from "@/lib/context/BrainConfigProvider/types";
|
||||
import { defineMaxTokens } from "@/lib/helpers/defineMexTokens";
|
||||
|
||||
import { useAddBrainModal } from "./hooks/useAddBrainModal";
|
||||
import { TextArea } from "../ui/TextField";
|
||||
|
||||
export const AddBrainModal = (): JSX.Element => {
|
||||
const {
|
||||
handleSubmit,
|
||||
isShareModalOpen,
|
||||
setIsShareModalOpen,
|
||||
register,
|
||||
openAiKey,
|
||||
temperature,
|
||||
maxTokens,
|
||||
model,
|
||||
isPending,
|
||||
} = useAddBrainModal();
|
||||
|
||||
return (
|
||||
<Modal
|
||||
Trigger={
|
||||
<Button variant={"secondary"}>
|
||||
Add New Brain
|
||||
<MdAdd className="text-xl" />
|
||||
</Button>
|
||||
}
|
||||
title="Add Brain"
|
||||
desc="Create a new brain to start aggregating content"
|
||||
isOpen={isShareModalOpen}
|
||||
setOpen={setIsShareModalOpen}
|
||||
CloseTrigger={<div />}
|
||||
>
|
||||
<form
|
||||
onSubmit={(e) => void handleSubmit(e)}
|
||||
className="my-10 flex flex-col items-center gap-2"
|
||||
>
|
||||
<Field
|
||||
label="Enter a brain name"
|
||||
autoFocus
|
||||
placeholder="E.g. History notes"
|
||||
autoComplete="off"
|
||||
className="flex-1"
|
||||
{...register("name")}
|
||||
/>
|
||||
|
||||
<TextArea
|
||||
label="Enter a brain description"
|
||||
autoFocus
|
||||
placeholder="My new brain is about..."
|
||||
autoComplete="off"
|
||||
className="flex-1 m-3"
|
||||
{...register("description")}
|
||||
/>
|
||||
|
||||
<Field
|
||||
label="OpenAI API Key"
|
||||
autoFocus
|
||||
placeholder="sk-xxx"
|
||||
autoComplete="off"
|
||||
className="flex-1"
|
||||
{...register("openAiKey")}
|
||||
/>
|
||||
|
||||
<fieldset className="w-full flex flex-col">
|
||||
<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="256"
|
||||
max={defineMaxTokens(model)}
|
||||
step="32"
|
||||
value={maxTokens}
|
||||
{...register("maxTokens")}
|
||||
/>
|
||||
</fieldset>
|
||||
<div className="flex flex-row justify-start w-full mt-4">
|
||||
<label className="flex items-center">
|
||||
<span className="mr-2 text-gray-700">Set as default brain</span>
|
||||
<input
|
||||
type="checkbox"
|
||||
{...register("setDefault")}
|
||||
className="form-checkbox h-5 w-5 text-indigo-600 rounded focus:ring-2 focus:ring-indigo-400"
|
||||
/>
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<Button isLoading={isPending} className="mt-12 self-end" type="submit">
|
||||
Create
|
||||
<MdAdd className="text-xl" />
|
||||
</Button>
|
||||
</form>
|
||||
</Modal>
|
||||
);
|
||||
};
|
116
frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts
Normal file
116
frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts
Normal file
@ -0,0 +1,116 @@
|
||||
/* eslint-disable max-lines */
|
||||
import axios from "axios";
|
||||
import { FormEvent, useState } from "react";
|
||||
import { useForm } from "react-hook-form";
|
||||
|
||||
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
|
||||
import { useBrainConfig } from "@/lib/context/BrainConfigProvider";
|
||||
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
|
||||
import { useToast } from "@/lib/hooks";
|
||||
|
||||
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
|
||||
export const useAddBrainModal = () => {
|
||||
const [isPending, setIsPending] = useState(false);
|
||||
const { publish } = useToast();
|
||||
const { createBrain } = useBrainContext();
|
||||
const { setAsDefaultBrain } = useBrainApi();
|
||||
const [isShareModalOpen, setIsShareModalOpen] = useState(false);
|
||||
const { config } = useBrainConfig();
|
||||
const defaultValues = {
|
||||
...config,
|
||||
name: "",
|
||||
description: "",
|
||||
setDefault: false,
|
||||
};
|
||||
|
||||
const { register, getValues, reset, watch } = useForm({
|
||||
defaultValues,
|
||||
});
|
||||
|
||||
const openAiKey = watch("openAiKey");
|
||||
const model = watch("model");
|
||||
const temperature = watch("temperature");
|
||||
const maxTokens = watch("maxTokens");
|
||||
|
||||
const handleSubmit = async (e: FormEvent) => {
|
||||
e.preventDefault();
|
||||
const { name, description, setDefault } = getValues();
|
||||
|
||||
console.log({
|
||||
name,
|
||||
description,
|
||||
maxTokens,
|
||||
model,
|
||||
setDefault,
|
||||
openAiKey,
|
||||
temperature,
|
||||
});
|
||||
|
||||
if (name.trim() === "" || isPending) {
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
setIsPending(true);
|
||||
const createdBrainId = await createBrain({
|
||||
name,
|
||||
description,
|
||||
max_tokens: maxTokens,
|
||||
model,
|
||||
openai_api_key: openAiKey,
|
||||
temperature,
|
||||
});
|
||||
|
||||
if (setDefault) {
|
||||
if (createdBrainId === undefined) {
|
||||
publish({
|
||||
variant: "danger",
|
||||
text: "Error occurred while creating a brain",
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
await setAsDefaultBrain(createdBrainId);
|
||||
}
|
||||
|
||||
setIsShareModalOpen(false);
|
||||
reset(defaultValues);
|
||||
publish({
|
||||
variant: "success",
|
||||
text: "Brain created successfully",
|
||||
});
|
||||
} catch (err) {
|
||||
if (axios.isAxiosError(err) && err.response?.status === 429) {
|
||||
publish({
|
||||
variant: "danger",
|
||||
text: `${JSON.stringify(
|
||||
(
|
||||
err.response as {
|
||||
data: { detail: string };
|
||||
}
|
||||
).data.detail
|
||||
)}`,
|
||||
});
|
||||
} else {
|
||||
publish({
|
||||
variant: "danger",
|
||||
text: `${JSON.stringify(err)}`,
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
setIsPending(false);
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
isShareModalOpen,
|
||||
setIsShareModalOpen,
|
||||
handleSubmit,
|
||||
register,
|
||||
openAiKey,
|
||||
model,
|
||||
temperature,
|
||||
maxTokens,
|
||||
isPending,
|
||||
};
|
||||
};
|
1
frontend/lib/components/AddBrainModal/index.ts
Normal file
1
frontend/lib/components/AddBrainModal/index.ts
Normal file
@ -0,0 +1 @@
|
||||
export * from "./AddBrainModal";
|
@ -7,7 +7,7 @@ import Popover from "@/lib/components/ui/Popover";
|
||||
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
|
||||
|
||||
import { BrainActions } from "./components/BrainActions/BrainActions";
|
||||
import { AddBrainModal } from "../../../../../AddBrainModal";
|
||||
import { AddBrainModal } from "../../../../../AddBrainModal/AddBrainModal";
|
||||
|
||||
export const BrainsDropDown = (): JSX.Element => {
|
||||
const [searchQuery, setSearchQuery] = useState("");
|
||||
|
@ -1,2 +1,2 @@
|
||||
export * from "@/lib/components/AddBrainModal";
|
||||
export * from "@/lib/components/AddBrainModal/AddBrainModal";
|
||||
export * from "./BrainActions";
|
||||
|
39
frontend/lib/components/ui/TextField.tsx
Normal file
39
frontend/lib/components/ui/TextField.tsx
Normal file
@ -0,0 +1,39 @@
|
||||
/* eslint-disable */
|
||||
import {
|
||||
DetailedHTMLProps,
|
||||
forwardRef,
|
||||
InputHTMLAttributes,
|
||||
RefObject,
|
||||
} from "react";
|
||||
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
interface FieldProps
|
||||
extends DetailedHTMLProps<
|
||||
InputHTMLAttributes<HTMLTextAreaElement>,
|
||||
HTMLTextAreaElement
|
||||
> {
|
||||
label?: string;
|
||||
name: string;
|
||||
}
|
||||
|
||||
export const TextArea = forwardRef(
|
||||
({ label, className, name, ...props }: FieldProps, forwardedRef) => {
|
||||
return (
|
||||
<fieldset className={cn("flex flex-col w-full", className)} name={name}>
|
||||
{label && (
|
||||
<label htmlFor={name} className="text-sm">
|
||||
{label}
|
||||
</label>
|
||||
)}
|
||||
<textarea
|
||||
ref={forwardedRef as RefObject<HTMLTextAreaElement>}
|
||||
className="w-full bg-gray-50 dark:bg-gray-900 px-4 py-2 border rounded-md border-black/10 dark:border-white/25"
|
||||
name={name}
|
||||
id={name}
|
||||
{...props}
|
||||
/>
|
||||
</fieldset>
|
||||
);
|
||||
}
|
||||
);
|
@ -2,6 +2,7 @@
|
||||
import { UUID } from "crypto";
|
||||
import { useCallback, useState } from "react";
|
||||
|
||||
import { CreateBrainInput } from "@/lib/api/brain/brain";
|
||||
import { useBrainApi } from "@/lib/api/brain/useBrainApi";
|
||||
import { useToast } from "@/lib/hooks";
|
||||
import { useEventTracking } from "@/services/analytics/useEventTracking";
|
||||
@ -28,9 +29,9 @@ export const useBrainProvider = () => {
|
||||
const currentBrain = allBrains.find((brain) => brain.id === currentBrainId);
|
||||
|
||||
const createBrainHandler = async (
|
||||
name: string
|
||||
brain: CreateBrainInput
|
||||
): Promise<UUID | undefined> => {
|
||||
const createdBrain = await createBrain(name);
|
||||
const createdBrain = await createBrain(brain);
|
||||
try {
|
||||
setAllBrains((prevBrains) => [...prevBrains, createdBrain]);
|
||||
saveBrainInLocalStorage(createdBrain);
|
||||
|
17
frontend/lib/helpers/defineMexTokens.ts
Normal file
17
frontend/lib/helpers/defineMexTokens.ts
Normal file
@ -0,0 +1,17 @@
|
||||
import { Model, PaidModels } from "../context/BrainConfigProvider/types";
|
||||
|
||||
export const defineMaxTokens = (model: Model | PaidModels): number => {
|
||||
//At the moment is evaluating only models from OpenAI
|
||||
switch (model) {
|
||||
case "gpt-3.5-turbo-0613":
|
||||
return 500;
|
||||
case "gpt-3.5-turbo-16k":
|
||||
return 2000;
|
||||
case "gpt-4":
|
||||
return 1000;
|
||||
case "gpt-4-0613":
|
||||
return 100;
|
||||
default:
|
||||
return 250;
|
||||
}
|
||||
};
|
Loading…
Reference in New Issue
Block a user