[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

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

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

frontend/app/config/components/ModelConfig.tsx

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

frontend/lib/api/brain/__tests__/useBrainApi.test.ts

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

frontend/lib/api/brain/brain.ts

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

frontend/lib/api/brain/useBrainApi.ts

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

frontend/lib/components/AddBrainModal.tsx

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

frontend/lib/components/AddBrainModal/AddBrainModal.tsx

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

frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts

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

frontend/lib/components/AddBrainModal/index.ts

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

frontend/lib/components/NavBar/components/NavItems/components/BrainsDropDown/BrainsDropDown.tsx

@@ -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("");

frontend/lib/components/NavBar/components/NavItems/components/BrainsDropDown/components/index.ts

@@ -1,2 +1,2 @@
-export * from "@/lib/components/AddBrainModal";
+export * from "@/lib/components/AddBrainModal/AddBrainModal";
 export * from "./BrainActions";

frontend/lib/components/ui/TextField.tsx

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

frontend/lib/context/BrainProvider/hooks/useBrainProvider.ts

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

frontend/lib/helpers/defineMexTokens.ts

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