diff --git a/backend/core/models/brains.py b/backend/core/models/brains.py index 362cf7728..475ded951 100644 --- a/backend/core/models/brains.py +++ b/backend/core/models/brains.py @@ -2,11 +2,10 @@ from typing import Any, List, Optional from uuid import UUID from logger import get_logger -from pydantic import BaseModel -from utils.vectors import get_unique_files_from_vector_ids - from models.settings import BrainRateLimiting, CommonsDep, common_dependencies from models.users import User +from pydantic import BaseModel +from utils.vectors import get_unique_files_from_vector_ids logger = get_logger(__name__) @@ -14,10 +13,12 @@ logger = get_logger(__name__) class Brain(BaseModel): id: Optional[UUID] = None name: Optional[str] = "Default brain" - status: Optional[str] = "public" + description: Optional[str] = "This is a description" + status: Optional[str] = "private" model: Optional[str] = "gpt-3.5-turbo-0613" temperature: Optional[float] = 0.0 max_tokens: Optional[int] = 256 + openai_api_key: Optional[str] = None files: List[Any] = [] max_brain_size = BrainRateLimiting().max_brain_size @@ -150,7 +151,20 @@ class Brain(BaseModel): def create_brain(self): commons = common_dependencies() response = ( - commons["supabase"].table("brains").insert({"name": self.name}).execute() + commons["supabase"] + .table("brains") + .insert( + { + "name": self.name, + "description": self.description, + "temperature": self.temperature, + "model": self.model, + "max_tokens": self.max_tokens, + "openai_api_key": self.openai_api_key, + "status": self.status, + } + ) + .execute() ) self.id = response.data[0]["brain_id"] @@ -174,6 +188,18 @@ class Brain(BaseModel): return response.data + def set_as_default_brain_for_user(self, user: User): + old_default_brain = get_default_user_brain(user) + + if old_default_brain is not None: + self.commons["supabase"].table("brains_users").update( + {"default_brain": False} + ).match({"brain_id": old_default_brain["id"], "user_id": user.id}).execute() + + self.commons["supabase"].table("brains_users").update( + {"default_brain": True} + ).match({"brain_id": self.id, "user_id": user.id}).execute() + def create_brain_vector(self, vector_id, file_sha1): response = ( self.commons["supabase"] @@ -201,15 +227,17 @@ class Brain(BaseModel): return vectorsResponse.data def update_brain_fields(self): - self.commons["supabase"].table("brains").update({"name": self.name}).match( - {"brain_id": self.id} - ).execute() - - def update_brain_with_file(self, file_sha1: str): - # not used - vector_ids = self.get_vector_ids_from_file_sha1(file_sha1) - for vector_id in vector_ids: - self.create_brain_vector(vector_id, file_sha1) + self.commons["supabase"].table("brains").update( + { + "name": self.name, + "description": self.description, + "model": self.model, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + "openai_api_key": self.openai_api_key, + "status": self.status, + } + ).match({"brain_id": self.id}).execute() def get_unique_brain_files(self): """ diff --git a/backend/core/routes/brain_routes.py b/backend/core/routes/brain_routes.py index efe0062e7..0588a2ce2 100644 --- a/backend/core/routes/brain_routes.py +++ b/backend/core/routes/brain_routes.py @@ -8,10 +8,9 @@ from models.brains import ( get_default_user_brain, get_default_user_brain_or_create_new, ) -from models.settings import BrainRateLimiting, common_dependencies +from models.settings import BrainRateLimiting from models.users import User - -from routes.authorizations.brain_authorization import has_brain_authorization +from routes.authorizations.brain_authorization import RoleEnum, has_brain_authorization logger = get_logger(__name__) @@ -102,8 +101,6 @@ async def create_brain_endpoint( In the brains table & in the brains_users table and put the creator user as 'Owner' """ - brain = Brain(name=brain.name) # pyright: ignore reportPrivateUsage=none - user_brains = brain.get_user_brains(current_user.id) max_brain_per_user = BrainRateLimiting().max_brain_per_user @@ -142,7 +139,7 @@ async def create_brain_endpoint( Depends( AuthBearer(), ), - Depends(has_brain_authorization()), + Depends(has_brain_authorization([RoleEnum.Editor, RoleEnum.Owner])), ], tags=["Brain"], ) @@ -151,22 +148,35 @@ async def update_brain_endpoint( input_brain: Brain, ): """ - Update an existing brain with new brain parameters/files. - If the file is contained in Add file to brain : - if given a fileName/ file sha1 / -> add all the vector Ids to the brains_vectors - Modify other brain fields: - name, status, model, max_tokens, temperature - Return modified brain ? No need -> do an optimistic update + Update an existing brain with new brain configuration + """ + input_brain.id = brain_id + print("brain", input_brain) + + input_brain.update_brain_fields() + return {"message": f"Brain {brain_id} has been updated."} + + +# set as default brain +@brain_router.post( + "/brains/{brain_id}/default", + dependencies=[ + Depends( + AuthBearer(), + ), + Depends(has_brain_authorization()), + ], + tags=["Brain"], +) +async def set_as_default_brain_endpoint( + brain_id: UUID, + user: User = Depends(get_current_user), +): + """ + Set a brain as default for the current user. """ - commons = common_dependencies() brain = Brain(id=brain_id) - # Add new file to brain , il file_sha1 already exists in brains_vectors -> out (not now) - if brain.file_sha1: # pyright: ignore reportPrivateUsage=none - # add all the vector Ids to the brains_vectors with the given brain.brain_id - brain.update_brain_with_file( - file_sha1=input_brain.file_sha1 # pyright: ignore reportPrivateUsage=none - ) + brain.set_as_default_brain_for_user(user) - brain.update_brain_fields(commons, brain) # pyright: ignore reportPrivateUsage=none - return {"message": f"Brain {brain_id} has been updated."} + return {"message": f"Brain {brain_id} has been set as default brain."} diff --git a/backend/core/tests/test_brains.py b/backend/core/tests/test_brains.py index 2c3f0a0d8..66460d355 100644 --- a/backend/core/tests/test_brains.py +++ b/backend/core/tests/test_brains.py @@ -1,6 +1,8 @@ import random import string +from models.brains import get_default_user_brain + def test_retrieve_default_brain(client, api_key): # Making a GET request to the /brains/default/ endpoint @@ -161,3 +163,55 @@ def test_delete_all_brains_and_get_default_brain(client, api_key): # Assert that the response status code is 200 (HTTP OK) assert response.status_code == 200 assert response.json()["name"] == "Default brain" + + +def test_set_as_default_brain_endpoint(client, api_key): + random_brain_name = "".join( + random.choices(string.ascii_letters + string.digits, k=10) + ) + # Set up the request payload + payload = { + "name": random_brain_name, + "status": "public", + "model": "gpt-3.5-turbo-0613", + "temperature": 0, + "max_tokens": 256, + } + + # Making a POST request to the /brains/ endpoint + response = client.post( + "/brains/", + json=payload, + headers={"Authorization": "Bearer " + api_key}, + ) + + response_data = response.json() + + brain_id = response_data["id"] + + # Make a POST request to set the brain as default for the user + response = client.post( + f"/brains/{brain_id}/default", + headers={"Authorization": "Bearer " + api_key}, + ) + + # Assert that the response status code is 200 (HTTP OK) + assert response.status_code == 200 + + # Assert the response message + assert response.json() == { + "message": f"Brain {brain_id} has been set as default brain." + } + + # Check if the brain is now the default for the user + + # Send a request to get user information + response = client.get("/user", headers={"Authorization": "Bearer " + api_key}) + # Assert that the response contains the expected fields + user_info = response.json() + user_id = user_info["id"] + + default_brain = get_default_user_brain(user_id) + assert default_brain is not None + assert default_brain["id"] == brain_id + assert default_brain["default_brain"] is True diff --git a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts index dbb25aed4..9d0e08e8b 100644 --- a/frontend/lib/api/brain/__tests__/useBrainApi.test.ts +++ b/frontend/lib/api/brain/__tests__/useBrainApi.test.ts @@ -198,7 +198,7 @@ describe("useBrainApi", () => { } = renderHook(() => useBrainApi()); const brainId = "123"; await setAsDefaultBrain(brainId); - expect(axiosPutMock).toHaveBeenCalledTimes(1); - expect(axiosPutMock).toHaveBeenCalledWith(`/brains/${brainId}/default`); + expect(axiosPostMock).toHaveBeenCalledTimes(1); + expect(axiosPostMock).toHaveBeenCalledWith(`/brains/${brainId}/default`); }); }); diff --git a/frontend/lib/api/brain/brain.ts b/frontend/lib/api/brain/brain.ts index a3dc592cb..236237b9f 100644 --- a/frontend/lib/api/brain/brain.ts +++ b/frontend/lib/api/brain/brain.ts @@ -131,5 +131,5 @@ export const setAsDefaultBrain = async ( brainId: string, axiosInstance: AxiosInstance ): Promise => { - await axiosInstance.put(`/brains/${brainId}/default`); + await axiosInstance.post(`/brains/${brainId}/default`); }; diff --git a/frontend/lib/components/AddBrainModal/AddBrainModal.tsx b/frontend/lib/components/AddBrainModal/AddBrainModal.tsx index da4254295..fe7a2b19b 100644 --- a/frontend/lib/components/AddBrainModal/AddBrainModal.tsx +++ b/frontend/lib/components/AddBrainModal/AddBrainModal.tsx @@ -107,9 +107,8 @@ export const AddBrainModal = (): JSX.Element => { diff --git a/frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts b/frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts index d737dd0b6..4dd25703a 100644 --- a/frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts +++ b/frontend/lib/components/AddBrainModal/hooks/useAddBrainModal.ts @@ -1,11 +1,12 @@ /* eslint-disable max-lines */ import axios from "axios"; -import { FormEvent, useState } from "react"; +import { FormEvent, useEffect, 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 { defineMaxTokens } from "@/lib/helpers/defineMexTokens"; import { useToast } from "@/lib/hooks"; // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types @@ -23,7 +24,7 @@ export const useAddBrainModal = () => { setDefault: false, }; - const { register, getValues, reset, watch } = useForm({ + const { register, getValues, reset, watch, setValue } = useForm({ defaultValues, }); @@ -32,20 +33,14 @@ export const useAddBrainModal = () => { const temperature = watch("temperature"); const maxTokens = watch("maxTokens"); + useEffect(() => { + setValue("maxTokens", Math.min(maxTokens, defineMaxTokens(model))); + }, [maxTokens, model, setValue]); + 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; } diff --git a/scripts/202307241530031_add_fields_to_brain.sql b/scripts/202307241530031_add_fields_to_brain.sql new file mode 100644 index 000000000..9262ca7bb --- /dev/null +++ b/scripts/202307241530031_add_fields_to_brain.sql @@ -0,0 +1,31 @@ +BEGIN; + +-- Change max_tokens type to INT +ALTER TABLE brains ALTER COLUMN max_tokens TYPE INT USING max_tokens::INT; + +-- Add or rename the api_key column to openai_api_key +DO $$ +BEGIN + BEGIN + -- Check if the api_key column exists + IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'brains' AND column_name = 'api_key') THEN + -- Rename the api_key column to openai_api_key + ALTER TABLE brains RENAME COLUMN api_key TO openai_api_key; + ELSE + -- Create the openai_api_key column if it doesn't exist + ALTER TABLE brains ADD COLUMN openai_api_key TEXT; + END IF; + END; +END $$; + +-- Add description column +ALTER TABLE brains ADD COLUMN description TEXT; + +-- Update migrations table +INSERT INTO migrations (name) +SELECT '202307241530031_add_fields_to_brain' +WHERE NOT EXISTS ( + SELECT 1 FROM migrations WHERE name = '202307241530031_add_fields_to_brain' +); + +COMMIT; diff --git a/scripts/tables.sql b/scripts/tables.sql index 4ad5c43b6..2f0816baa 100644 --- a/scripts/tables.sql +++ b/scripts/tables.sql @@ -126,14 +126,15 @@ CREATE TABLE IF NOT EXISTS api_keys( is_active BOOLEAN DEFAULT true ); --- Create brains table -CREATE TABLE IF NOT EXISTS brains ( +CREATE TABLE IF NOT EXISTS brains ( brain_id UUID DEFAULT gen_random_uuid() PRIMARY KEY, - name TEXT, + name TEXT NOT NULL, status TEXT, + description TEXT, model TEXT, - max_tokens TEXT, - temperature FLOAT + max_tokens INT, + temperature FLOAT, + openai_api_key TEXT ); -- Create brains X users table @@ -193,7 +194,7 @@ CREATE TABLE IF NOT EXISTS migrations ( ); INSERT INTO migrations (name) -SELECT '20230717173000_add_get_user_id_by_user_email' +SELECT '202307241530031_add_fields_to_brain' WHERE NOT EXISTS ( - SELECT 1 FROM migrations WHERE name = '20230717173000_add_get_user_id_by_user_email' + SELECT 1 FROM migrations WHERE name = '202307241530031_add_fields_to_brain' );