From 4ca6c667da3d5daf0339c65f077c8956c7ef42e8 Mon Sep 17 00:00:00 2001 From: Mamadou DICKO <63923024+mamadoudicko@users.noreply.github.com> Date: Thu, 3 Aug 2023 10:37:13 +0200 Subject: [PATCH] feat(backend): implement brain-prompt link (#831) * feat: add prompt_id field to brain * feat(Prompt controller): update prompt routes * feat: remove unused private prompts * refactor: add BrainEntity and repo and service * tests: partially type main Repository * feat: add PromptStatusEnum enum * feat: change delete prompt repository return type --- backend/core/models/brain_entity.py | 23 ++++ backend/core/models/brains.py | 71 +--------- backend/core/models/databases/repository.py | 12 +- .../core/models/databases/supabase/brains.py | 124 ++++++++++++++---- .../core/models/databases/supabase/prompts.py | 25 +++- backend/core/models/prompt.py | 8 +- backend/core/repository/brain/create_brain.py | 9 ++ .../repository/brain/create_brain_user.py | 16 +++ .../core/repository/brain/get_brain_by_id.py | 10 ++ .../repository/brain/get_brain_details.py | 19 +++ .../repository/brain/get_brain_for_user.py | 9 ++ .../brain/get_default_user_brain.py | 22 ++++ .../get_default_user_brain_or_create_new.py | 17 +++ .../core/repository/brain/get_user_brains.py | 11 ++ .../brain/set_as_default_brain_for_user.py | 19 +++ backend/core/repository/brain/update_brain.py | 12 ++ .../resend_invitation_email.py | 7 +- .../core/repository/prompt/create_prompt.py | 1 - .../repository/prompt/delete_prompt_py_id.py | 4 +- .../authorizations/brain_authorization.py | 14 +- backend/core/routes/authorizations/types.py | 7 + backend/core/routes/brain_routes.py | 85 ++++++++---- backend/core/routes/chat_routes.py | 28 ++-- backend/core/routes/prompt_routes.py | 42 ++++++ backend/core/routes/subscription_routes.py | 24 ++-- backend/core/routes/upload_routes.py | 5 +- backend/core/routes/user_routes.py | 7 +- backend/core/tests/test_brains.py | 7 +- .../20230802120700_add_prompt_id_to_brain.sql | 19 +++ scripts/tables.sql | 8 +- 30 files changed, 473 insertions(+), 192 deletions(-) create mode 100644 backend/core/models/brain_entity.py create mode 100644 backend/core/repository/brain/create_brain.py create mode 100644 backend/core/repository/brain/create_brain_user.py create mode 100644 backend/core/repository/brain/get_brain_by_id.py create mode 100644 backend/core/repository/brain/get_brain_details.py create mode 100644 backend/core/repository/brain/get_brain_for_user.py create mode 100644 backend/core/repository/brain/get_default_user_brain.py create mode 100644 backend/core/repository/brain/get_default_user_brain_or_create_new.py create mode 100644 backend/core/repository/brain/get_user_brains.py create mode 100644 backend/core/repository/brain/set_as_default_brain_for_user.py create mode 100644 backend/core/repository/brain/update_brain.py create mode 100644 backend/core/routes/authorizations/types.py create mode 100644 scripts/20230802120700_add_prompt_id_to_brain.sql diff --git a/backend/core/models/brain_entity.py b/backend/core/models/brain_entity.py new file mode 100644 index 000000000..54b51fa92 --- /dev/null +++ b/backend/core/models/brain_entity.py @@ -0,0 +1,23 @@ +from typing import Optional +from uuid import UUID + +from pydantic import BaseModel +from routes.authorizations.types import RoleEnum + + +class BrainEntity(BaseModel): + brain_id: UUID + name: str + description: Optional[str] + temperature: Optional[float] + model: Optional[str] + max_tokens: Optional[int] + openai_api_key: Optional[str] + status: Optional[str] + prompt_id: Optional[UUID] + + +class MinimalBrainEntity(BaseModel): + id: UUID + name: str + rights: RoleEnum diff --git a/backend/core/models/brains.py b/backend/core/models/brains.py index 6b5165cdd..c2d606537 100644 --- a/backend/core/models/brains.py +++ b/backend/core/models/brains.py @@ -6,7 +6,6 @@ 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 logger = get_logger(__name__) @@ -22,6 +21,7 @@ class Brain(BaseModel): openai_api_key: Optional[str] = None files: List[Any] = [] max_brain_size = BrainRateLimiting().max_brain_size + prompt_id: Optional[UUID] = None class Config: arbitrary_types_allowed = True @@ -77,24 +77,6 @@ class Brain(BaseModel): {"brain_id": self.id, "user_id": user_id} ).execute() - def get_user_brains(self, user_id): - return self.commons["db"].get_user_brains(user_id) - - def get_brain_for_user(self, user_id): - return self.commons["db"].get_brain_for_user(user_id, self.id) - - def get_brain_details(self): - response = ( - self.commons["supabase"] - .from_("brains") - .select("id:brain_id, name, *") - .filter("brain_id", "eq", self.id) - .execute() - ) - if response.data == []: - return None - return response.data[0] - def delete_brain(self, user_id): results = self.commons["db"].delete_brain_user_by_id(user_id, self.id) @@ -105,37 +87,12 @@ class Brain(BaseModel): self.commons["db"].delete_brain_user(self.id) self.commons["db"].delete_brain(self.id) - def create_brain(self): - response = self.commons["db"].create_brain(self.name) - self.id = response.data[0]["brain_id"] - return response.data - - def create_brain_user(self, user_id: UUID, rights, default_brain): - response = self.commons["db"].create_brain_user(user_id=user_id, brain_id=self.id, rights=rights, default_brain=default_brain) - self.id = response.data[0]["brain_id"] - 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): return self.commons["db"].create_brain_vector(self.id, vector_id, file_sha1) def get_vector_ids_from_file_sha1(self, file_sha1: str): return self.commons["db"].get_vector_ids_from_file_sha1(file_sha1) - def update_brain_fields(self): - return self.commons["db"].update_brain_fields(brain_id=self.id, brain_name=self.name) - def update_brain_with_file(self, file_sha1: str): # not used vector_ids = self.get_vector_ids_from_file_sha1(file_sha1) @@ -154,29 +111,3 @@ class Brain(BaseModel): def delete_file_from_brain(self, file_name: str): return self.commons["db"].delete_file_from_brain(self.id, file_name) - - -def get_default_user_brain(user: User): - commons = common_dependencies() - response = commons["db"].get_default_user_brain_id(user.id) - - logger.info("Default brain response:", response.data) - default_brain_id = response.data[0]["brain_id"] if response.data else None - - logger.info(f"Default brain id: {default_brain_id}") - - if default_brain_id: - brain_response = commons["db"].get_brain_by_id(default_brain_id) - return brain_response.data[0] if brain_response.data else None - - -def get_default_user_brain_or_create_new(user: User) -> Brain: - default_brain = get_default_user_brain(user) - - if default_brain: - return Brain.create(**default_brain) - else: - brain = Brain.create() - brain.create_brain() - brain.create_brain_user(user.id, "Owner", True) - return brain diff --git a/backend/core/models/databases/repository.py b/backend/core/models/databases/repository.py index 889a01a8c..3973c9e8c 100644 --- a/backend/core/models/databases/repository.py +++ b/backend/core/models/databases/repository.py @@ -2,10 +2,12 @@ from abc import ABC, abstractmethod from datetime import datetime from uuid import UUID +from models.brain_entity import BrainEntity + class Repository(ABC): @abstractmethod - def get_user_brains(self, user_id: str): + def get_user_brains(self, user_id: str) -> list[BrainEntity]: pass @abstractmethod @@ -29,7 +31,7 @@ class Repository(ABC): pass @abstractmethod - def create_brain(self, name: str): + def create_brain(self, brain: str): pass @abstractmethod @@ -46,10 +48,6 @@ class Repository(ABC): def get_vector_ids_from_file_sha1(self, file_sha1: str): pass - @abstractmethod - def update_brain_fields(self, brain_id: UUID, brain_name: str): - pass - @abstractmethod def get_brain_vector_ids(self, brain_id: UUID): pass @@ -59,7 +57,7 @@ class Repository(ABC): pass @abstractmethod - def get_default_user_brain_id(self, user_id: UUID): + def get_default_user_brain_id(self, user_id: UUID) -> UUID: pass @abstractmethod diff --git a/backend/core/models/databases/supabase/brains.py b/backend/core/models/databases/supabase/brains.py index d8a24e276..677a32330 100644 --- a/backend/core/models/databases/supabase/brains.py +++ b/backend/core/models/databases/supabase/brains.py @@ -1,30 +1,77 @@ +from typing import Optional from uuid import UUID -from models.databases.repository import Repository from logger import get_logger +from models.brain_entity import BrainEntity, MinimalBrainEntity +from models.databases.repository import Repository +from pydantic import BaseModel logger = get_logger(__name__) +class CreateBrainProperties(BaseModel): + name: Optional[str] = "Default brain" + 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 + prompt_id: Optional[UUID] = None + + def dict(self, *args, **kwargs): + brain_dict = super().dict(*args, **kwargs) + if brain_dict.get("prompt_id"): + brain_dict["prompt_id"] = str(brain_dict.get("prompt_id")) + return brain_dict + + +class BrainUpdatableProperties(BaseModel): + name: Optional[str] + description: Optional[str] + temperature: Optional[float] + model: Optional[str] + max_tokens: Optional[int] + openai_api_key: Optional[str] + status: Optional[str] + prompt_id: Optional[UUID] + + def dict(self, *args, **kwargs): + brain_dict = super().dict(*args, **kwargs) + if brain_dict.get("prompt_id"): + brain_dict["prompt_id"] = str(brain_dict.get("prompt_id")) + return brain_dict + + class Brain(Repository): def __init__(self, supabase_client): self.db = supabase_client - def get_user_brains(self, user_id): + def create_brain(self, brain: CreateBrainProperties): + return BrainEntity( + **((self.db.table("brains").insert(brain)).execute()).data[0] + ) + + def get_user_brains(self, user_id) -> list[MinimalBrainEntity]: response = ( - self.db - .from_("brains_users") + self.db.from_("brains_users") .select("id:brain_id, rights, brains (id: brain_id, name)") .filter("user_id", "eq", user_id) .execute() ) - user_brains = [] + user_brains: list[MinimalBrainEntity] = [] for item in response.data: - user_brains.append(item["brains"]) - user_brains[-1]["rights"] = item["rights"] + user_brains.append( + MinimalBrainEntity( + id=item["brains"]["id"], + name=item["brains"]["name"], + rights=item["rights"], + ) + ) + user_brains[-1].rights = item["rights"] return user_brains - def get_brain_for_user(self, user_id, brain_id): + def get_brain_for_user(self, user_id, brain_id) -> MinimalBrainEntity | None: response = ( self.db.from_("brains_users") .select("id:brain_id, rights, brains (id: brain_id, name)") @@ -34,7 +81,13 @@ class Brain(Repository): ) if len(response.data) == 0: return None - return response.data[0] + brain_data = response.data[0] + + return MinimalBrainEntity( + id=brain_data["brains"]["id"], + name=brain_data["brains"]["name"], + rights=brain_data["rights"], + ) def get_brain_details(self, brain_id): response = ( @@ -81,10 +134,7 @@ class Brain(Repository): return results - def create_brain(self, name): - return self.db.table("brains").insert({"name": name}).execute() - - def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain): + def create_brain_user(self, user_id: UUID, brain_id, rights, default_brain: bool): response = ( self.db.table("brains_users") .insert( @@ -124,10 +174,20 @@ class Brain(Repository): ) return vectorsResponse.data - def update_brain_fields(self, brain_id, brain_name): - self.db.table("brains").update({"name": brain_name}).match( - {"brain_id": brain_id} - ).execute() + def update_brain_by_id( + self, brain_id: UUID, brain: BrainUpdatableProperties + ) -> BrainEntity | None: + update_brain_response = ( + self.db.table("brains") + .update(brain.dict(exclude_unset=True)) + .match({"brain_id": brain_id}) + .execute() + ).data + + if len(update_brain_response) == 0: + return None + + return BrainEntity(**update_brain_response[0]) def get_brain_vector_ids(self, brain_id): """ @@ -183,23 +243,29 @@ class Brain(Repository): return {"message": f"File {file_name} in brain {brain_id} has been deleted."} - def get_default_user_brain_id(self, user_id: UUID): + def get_default_user_brain_id(self, user_id: UUID) -> UUID | None: response = ( - self.db.from_("brains_users") - .select("brain_id") - .filter("user_id", "eq", user_id) - .filter("default_brain", "eq", True) - .execute() - ) + ( + self.db.from_("brains_users") + .select("brain_id") + .filter("user_id", "eq", user_id) + .filter("default_brain", "eq", True) + .execute() + ) + ).data + if len(response) == 0: + return None + return UUID(response[0].get("brain_id")) - return response - - def get_brain_by_id(self, brain_id: UUID): + def get_brain_by_id(self, brain_id: UUID) -> BrainEntity | None: response = ( self.db.from_("brains") .select("id:brain_id, name, *") .filter("brain_id", "eq", brain_id) .execute() - ) + ).data - return response + if len(response) == 0: + return None + + return BrainEntity(**response[0]) diff --git a/backend/core/models/databases/supabase/prompts.py b/backend/core/models/databases/supabase/prompts.py index d4f6802bf..0ee29c87c 100644 --- a/backend/core/models/databases/supabase/prompts.py +++ b/backend/core/models/databases/supabase/prompts.py @@ -3,7 +3,7 @@ from uuid import UUID from fastapi import HTTPException from models.databases.repository import Repository -from models.prompt import Prompt +from models.prompt import Prompt, PromptStatusEnum from pydantic import BaseModel @@ -12,7 +12,7 @@ class CreatePromptProperties(BaseModel): title: str content: str - status: str = "private" + status: PromptStatusEnum = PromptStatusEnum.private class PromptUpdatableProperties(BaseModel): @@ -20,7 +20,14 @@ class PromptUpdatableProperties(BaseModel): title: Optional[str] content: Optional[str] - status: Optional[str] + status: Optional[PromptStatusEnum] + + +class DeletePromptResponse(BaseModel): + """Response when deleting a prompt""" + + status: str = "delete" + prompt_id: UUID class Prompts(Repository): @@ -28,20 +35,22 @@ class Prompts(Repository): self.db = supabase_client def create_prompt(self, prompt: CreatePromptProperties) -> Prompt: - """Create a prompt by id""" + """ + Create a prompt + """ response = (self.db.from_("prompts").insert(prompt.dict()).execute()).data return Prompt(**response[0]) - def delete_prompt_by_id(self, prompt_id: UUID) -> Prompt | None: + def delete_prompt_by_id(self, prompt_id: UUID) -> DeletePromptResponse: """ Delete a prompt by id Args: prompt_id (UUID): The id of the prompt Returns: - Prompt: The prompt + A dictionary containing the status of the delete and prompt_id of the deleted prompt """ response = ( self.db.from_("prompts") @@ -50,9 +59,11 @@ class Prompts(Repository): .execute() .data ) + if response == []: raise HTTPException(404, "Prompt not found") - return Prompt(**response[0]) + + return DeletePromptResponse(status="deleted", prompt_id=prompt_id) def get_prompt_by_id(self, prompt_id: UUID) -> Prompt | None: """ diff --git a/backend/core/models/prompt.py b/backend/core/models/prompt.py index e5882d380..91112b643 100644 --- a/backend/core/models/prompt.py +++ b/backend/core/models/prompt.py @@ -1,10 +1,16 @@ +from enum import Enum from uuid import UUID from pydantic import BaseModel +class PromptStatusEnum(str, Enum): + private = "private" + public = "public" + + class Prompt(BaseModel): title: str content: str - status: str = "private" + status: PromptStatusEnum = PromptStatusEnum.private id: UUID diff --git a/backend/core/repository/brain/create_brain.py b/backend/core/repository/brain/create_brain.py new file mode 100644 index 000000000..6271ababf --- /dev/null +++ b/backend/core/repository/brain/create_brain.py @@ -0,0 +1,9 @@ +from models.brain_entity import BrainEntity +from models.databases.supabase.brains import CreateBrainProperties +from models.settings import common_dependencies + + +def create_brain(brain: CreateBrainProperties) -> BrainEntity: + commons = common_dependencies() + + return commons["db"].create_brain(brain.dict(exclude_unset=True)) diff --git a/backend/core/repository/brain/create_brain_user.py b/backend/core/repository/brain/create_brain_user.py new file mode 100644 index 000000000..48f2764a3 --- /dev/null +++ b/backend/core/repository/brain/create_brain_user.py @@ -0,0 +1,16 @@ +from uuid import UUID + +from models.settings import common_dependencies +from routes.authorizations.types import RoleEnum + + +def create_brain_user( + user_id: UUID, brain_id: UUID, rights: RoleEnum, is_default_brain: bool +) -> None: + commons = common_dependencies() + commons["db"].create_brain_user( + user_id=user_id, + brain_id=brain_id, + rights=rights, + default_brain=is_default_brain, + ).data[0] diff --git a/backend/core/repository/brain/get_brain_by_id.py b/backend/core/repository/brain/get_brain_by_id.py new file mode 100644 index 000000000..7308c2b6d --- /dev/null +++ b/backend/core/repository/brain/get_brain_by_id.py @@ -0,0 +1,10 @@ +from uuid import UUID + +from models.brain_entity import BrainEntity +from models.settings import common_dependencies + + +def get_brain_by_id(brain_id: UUID) -> BrainEntity | None: + commons = common_dependencies() + + return commons["db"].get_brain_by_id(brain_id) diff --git a/backend/core/repository/brain/get_brain_details.py b/backend/core/repository/brain/get_brain_details.py new file mode 100644 index 000000000..fec4b47d3 --- /dev/null +++ b/backend/core/repository/brain/get_brain_details.py @@ -0,0 +1,19 @@ +from uuid import UUID + +from models.brain_entity import BrainEntity +from models.settings import common_dependencies + + +def get_brain_details(brain_id: UUID) -> BrainEntity | None: + commons = common_dependencies() + response = ( + commons["supabase"] + .from_("brains") + .select("*") + .filter("brain_id", "eq", brain_id) + .execute() + ) + if response.data == []: + return None + + return BrainEntity(**response.data[0]) diff --git a/backend/core/repository/brain/get_brain_for_user.py b/backend/core/repository/brain/get_brain_for_user.py new file mode 100644 index 000000000..769f6bfd3 --- /dev/null +++ b/backend/core/repository/brain/get_brain_for_user.py @@ -0,0 +1,9 @@ +from uuid import UUID + +from models.brain_entity import MinimalBrainEntity +from models.settings import common_dependencies + + +def get_brain_for_user(user_id: UUID, brain_id: UUID) -> MinimalBrainEntity: + commons = common_dependencies() + return commons["db"].get_brain_for_user(user_id, brain_id) diff --git a/backend/core/repository/brain/get_default_user_brain.py b/backend/core/repository/brain/get_default_user_brain.py new file mode 100644 index 000000000..5152b3646 --- /dev/null +++ b/backend/core/repository/brain/get_default_user_brain.py @@ -0,0 +1,22 @@ +from uuid import UUID + +from logger import get_logger +from models.brain_entity import BrainEntity +from models.settings import common_dependencies +from repository.brain.get_brain_by_id import get_brain_by_id + +logger = get_logger(__name__) + + +def get_user_default_brain(user_id: UUID) -> BrainEntity | None: + commons = common_dependencies() + brain_id = commons["db"].get_default_user_brain_id(user_id) + + logger.info("Default brain response:", brain_id) + + if brain_id is None: + return None + + logger.info(f"Default brain id: {brain_id}") + + return get_brain_by_id(brain_id) diff --git a/backend/core/repository/brain/get_default_user_brain_or_create_new.py b/backend/core/repository/brain/get_default_user_brain_or_create_new.py new file mode 100644 index 000000000..45b7fdb2b --- /dev/null +++ b/backend/core/repository/brain/get_default_user_brain_or_create_new.py @@ -0,0 +1,17 @@ +from models.brain_entity import BrainEntity +from models.databases.supabase.brains import CreateBrainProperties +from models.users import User +from repository.brain.create_brain import create_brain +from repository.brain.create_brain_user import create_brain_user +from repository.brain.get_default_user_brain import get_user_default_brain +from routes.authorizations.types import RoleEnum + + +def get_default_user_brain_or_create_new(user: User) -> BrainEntity: + default_brain = get_user_default_brain(user.id) + + if not default_brain: + default_brain = create_brain(CreateBrainProperties()) + create_brain_user(user.id, default_brain.brain_id, RoleEnum.Owner, True) + + return default_brain diff --git a/backend/core/repository/brain/get_user_brains.py b/backend/core/repository/brain/get_user_brains.py new file mode 100644 index 000000000..b1ec19300 --- /dev/null +++ b/backend/core/repository/brain/get_user_brains.py @@ -0,0 +1,11 @@ +from uuid import UUID + +from models.brain_entity import BrainEntity +from models.settings import common_dependencies + + +def get_user_brains(user_id: UUID) -> list[BrainEntity]: + commons = common_dependencies() + results = commons["db"].get_user_brains(user_id) + + return results diff --git a/backend/core/repository/brain/set_as_default_brain_for_user.py b/backend/core/repository/brain/set_as_default_brain_for_user.py new file mode 100644 index 000000000..678f613b1 --- /dev/null +++ b/backend/core/repository/brain/set_as_default_brain_for_user.py @@ -0,0 +1,19 @@ +from uuid import UUID + +from models.settings import common_dependencies +from repository.brain.get_default_user_brain import get_user_default_brain + + +def set_as_default_brain_for_user(user_id: UUID, brain_id: UUID): + commons = common_dependencies() + + old_default_brain = get_user_default_brain(user_id) + + if old_default_brain is not None: + commons["supabase"].table("brains_users").update( + {"default_brain": False} + ).match({"brain_id": old_default_brain.brain_id, "user_id": user_id}).execute() + + commons["supabase"].table("brains_users").update({"default_brain": True}).match( + {"brain_id": brain_id, "user_id": user_id} + ).execute() diff --git a/backend/core/repository/brain/update_brain.py b/backend/core/repository/brain/update_brain.py new file mode 100644 index 000000000..a8d64bab1 --- /dev/null +++ b/backend/core/repository/brain/update_brain.py @@ -0,0 +1,12 @@ +from uuid import UUID + +from models.brain_entity import BrainEntity +from models.databases.supabase.brains import BrainUpdatableProperties +from models.settings import common_dependencies + + +def update_brain_by_id(brain_id: UUID, brain: BrainUpdatableProperties) -> BrainEntity: + """Update a prompt by id""" + commons = common_dependencies() + + return commons["db"].update_brain_by_id(brain_id, brain) diff --git a/backend/core/repository/brain_subscription/resend_invitation_email.py b/backend/core/repository/brain_subscription/resend_invitation_email.py index 050aba5bd..9efe2e0c7 100644 --- a/backend/core/repository/brain_subscription/resend_invitation_email.py +++ b/backend/core/repository/brain_subscription/resend_invitation_email.py @@ -1,9 +1,9 @@ import resend from logger import get_logger -from models.brains import Brain from models.brains_subscription_invitations import BrainSubscription from models.settings import BrainSettings +from repository.brain.get_brain_details import get_brain_details from repository.brain_subscription.get_brain_url import get_brain_url logger = get_logger(__name__) @@ -19,11 +19,10 @@ def resend_invitation_email( brain_url = get_brain_url(origin, brain_subscription.brain_id) - invitation_brain_client = Brain(id=brain_subscription.brain_id) - invitation_brain = invitation_brain_client.get_brain_details() + invitation_brain = get_brain_details(brain_subscription.brain_id) if invitation_brain is None: raise Exception("Brain not found") - brain_name = invitation_brain["name"] + brain_name = invitation_brain.name html_body = f"""

Brain {brain_name} has been shared with you by {inviter_email}.

diff --git a/backend/core/repository/prompt/create_prompt.py b/backend/core/repository/prompt/create_prompt.py index 0518e2fb1..3e41f4682 100644 --- a/backend/core/repository/prompt/create_prompt.py +++ b/backend/core/repository/prompt/create_prompt.py @@ -4,7 +4,6 @@ from models.settings import common_dependencies def create_prompt(prompt: CreatePromptProperties) -> Prompt: - """Create a prompt by id""" commons = common_dependencies() return commons["db"].create_prompt(prompt) diff --git a/backend/core/repository/prompt/delete_prompt_py_id.py b/backend/core/repository/prompt/delete_prompt_py_id.py index c60447c81..bbd6a199d 100644 --- a/backend/core/repository/prompt/delete_prompt_py_id.py +++ b/backend/core/repository/prompt/delete_prompt_py_id.py @@ -1,10 +1,10 @@ from uuid import UUID -from models.prompt import Prompt +from models.databases.supabase.prompts import DeletePromptResponse from models.settings import common_dependencies -def delete_prompt_by_id(prompt_id: UUID) -> Prompt | None: +def delete_prompt_by_id(prompt_id: UUID) -> DeletePromptResponse: """ Delete a prompt by id Args: diff --git a/backend/core/routes/authorizations/brain_authorization.py b/backend/core/routes/authorizations/brain_authorization.py index e37ab1371..b50c15966 100644 --- a/backend/core/routes/authorizations/brain_authorization.py +++ b/backend/core/routes/authorizations/brain_authorization.py @@ -1,17 +1,12 @@ -from enum import Enum from typing import List, Optional, Union from uuid import UUID from auth.auth_bearer import get_current_user from fastapi import Depends, HTTPException, status -from models.brains import Brain from models.users import User +from repository.brain.get_brain_for_user import get_brain_for_user - -class RoleEnum(str, Enum): - Viewer = "Viewer" - Editor = "Editor" - Owner = "Owner" +from routes.authorizations.types import RoleEnum def has_brain_authorization( @@ -53,8 +48,7 @@ def validate_brain_authorization( detail="Missing required role", ) - brain = Brain(id=brain_id) - user_brain = brain.get_brain_for_user(user_id) + user_brain = get_brain_for_user(user_id, brain_id) if user_brain is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -66,7 +60,7 @@ def validate_brain_authorization( required_roles = [required_roles] # Check if the user has at least one of the required roles - if user_brain.get("rights") not in required_roles: + if user_brain.rights not in required_roles: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="You don't have the required role(s) for this brain", diff --git a/backend/core/routes/authorizations/types.py b/backend/core/routes/authorizations/types.py new file mode 100644 index 000000000..69ef9fd4c --- /dev/null +++ b/backend/core/routes/authorizations/types.py @@ -0,0 +1,7 @@ +from enum import Enum + + +class RoleEnum(str, Enum): + Viewer = "Viewer" + Editor = "Editor" + Owner = "Owner" diff --git a/backend/core/routes/brain_routes.py b/backend/core/routes/brain_routes.py index 220613e30..82426d928 100644 --- a/backend/core/routes/brain_routes.py +++ b/backend/core/routes/brain_routes.py @@ -3,15 +3,31 @@ from uuid import UUID from auth import AuthBearer, get_current_user from fastapi import APIRouter, Depends, HTTPException from logger import get_logger -from models.brains import ( - Brain, - get_default_user_brain, - get_default_user_brain_or_create_new, +from models.databases.supabase.brains import ( + BrainUpdatableProperties, + CreateBrainProperties, ) from models.settings import BrainRateLimiting from models.users import User +from repository.brain.create_brain import create_brain +from repository.brain.create_brain_user import create_brain_user +from repository.brain.get_brain_details import get_brain_details +from repository.brain.get_default_user_brain import get_user_default_brain +from repository.brain.get_default_user_brain_or_create_new import ( + get_default_user_brain_or_create_new, +) +from repository.brain.get_user_brains import get_user_brains +from repository.brain.set_as_default_brain_for_user import ( + set_as_default_brain_for_user, +) +from repository.brain.update_brain import update_brain_by_id +from repository.prompt.delete_prompt_py_id import delete_prompt_by_id +from repository.prompt.get_prompt_by_id import get_prompt_by_id -from routes.authorizations.brain_authorization import RoleEnum, has_brain_authorization +from routes.authorizations.brain_authorization import ( + has_brain_authorization, +) +from routes.authorizations.types import RoleEnum logger = get_logger(__name__) @@ -30,8 +46,7 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)): This endpoint retrieves all the brains associated with the current authenticated user. It returns a list of brains objects containing the brain ID and brain name for each brain. """ - brain = Brain() - brains = brain.get_user_brains(current_user.id) + brains = get_user_brains(current_user.id) return {"brains": brains} @@ -51,10 +66,9 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us """ brain = get_default_user_brain_or_create_new(current_user) - return {"id": brain.id, "name": brain.name, "rights": "Owner"} + return {"id": brain.brain_id, "name": brain.name, "rights": "Owner"} -# get one brain - Currently not used in FE @brain_router.get( "/brains/{brain_id}/", dependencies=[Depends(AuthBearer()), Depends(has_brain_authorization())], @@ -72,9 +86,8 @@ async def get_brain_endpoint( This endpoint retrieves the details of a specific brain identified by the provided brain ID. It returns the brain ID and its history, which includes the brain messages exchanged in the brain. """ - brain = Brain(id=brain_id) - brain_details = brain.get_brain_details() + brain_details = get_brain_details(brain_id) if brain_details is None: raise HTTPException( status_code=404, @@ -87,7 +100,7 @@ async def get_brain_endpoint( # create new brain @brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"]) async def create_brain_endpoint( - brain: Brain, + brain: CreateBrainProperties, current_user: User = Depends(get_current_user), ): """ @@ -100,7 +113,7 @@ async def create_brain_endpoint( In the brains table & in the brains_users table and put the creator user as 'Owner' """ - user_brains = brain.get_user_brains(current_user.id) + user_brains = get_user_brains(current_user.id) max_brain_per_user = BrainRateLimiting().max_brain_per_user if len(user_brains) >= max_brain_per_user: @@ -109,23 +122,31 @@ async def create_brain_endpoint( detail=f"Maximum number of brains reached ({max_brain_per_user}).", ) - brain.create_brain() # pyright: ignore reportPrivateUsage=none - default_brain = get_default_user_brain(current_user) + new_brain = create_brain( + brain, + ) + default_brain = get_user_default_brain(current_user.id) if default_brain: logger.info(f"Default brain already exists for user {current_user.id}") - brain.create_brain_user( # pyright: ignore reportPrivateUsage=none - user_id=current_user.id, rights="Owner", default_brain=False + create_brain_user( + user_id=current_user.id, + brain_id=new_brain.brain_id, + rights=RoleEnum.Owner, + is_default_brain=False, ) else: logger.info( f"Default brain does not exist for user {current_user.id}. It will be created." ) - brain.create_brain_user( # pyright: ignore reportPrivateUsage=none - user_id=current_user.id, rights="Owner", default_brain=True + create_brain_user( + user_id=current_user.id, + brain_id=new_brain.brain_id, + rights=RoleEnum.Owner, + is_default_brain=True, ) return { - "id": brain.id, # pyright: ignore reportPrivateUsage=none + "id": new_brain.brain_id, "name": brain.name, "rights": "Owner", } @@ -144,15 +165,28 @@ async def create_brain_endpoint( ) async def update_brain_endpoint( brain_id: UUID, - input_brain: Brain, + input_brain: BrainUpdatableProperties, ): """ Update an existing brain with new brain configuration """ - input_brain.id = brain_id - print("brain", input_brain) - input_brain.update_brain_fields() + # 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", + ) + 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) + return {"message": f"Brain {brain_id} has been updated."} @@ -174,8 +208,7 @@ async def set_as_default_brain_endpoint( """ Set a brain as default for the current user. """ - brain = Brain(id=brain_id) - brain.set_as_default_brain_for_user(user) + set_as_default_brain_for_user(user.id, brain_id) return {"message": f"Brain {brain_id} has been set as default brain."} diff --git a/backend/core/routes/chat_routes.py b/backend/core/routes/chat_routes.py index 54cfe9239..8f2d143ad 100644 --- a/backend/core/routes/chat_routes.py +++ b/backend/core/routes/chat_routes.py @@ -9,11 +9,15 @@ from auth import AuthBearer, get_current_user from fastapi import APIRouter, Depends, Query, Request from fastapi.responses import StreamingResponse from llm.openai import OpenAIBrainPicking -from models.brains import Brain, get_default_user_brain_or_create_new +from models.brains import Brain from models.chat import Chat, ChatHistory from models.chats import ChatQuestion from models.settings import LLMSettings, common_dependencies from models.users import User +from repository.brain.get_brain_details import get_brain_details +from repository.brain.get_default_user_brain_or_create_new import ( + get_default_user_brain_or_create_new, +) from repository.chat.create_chat import CreateChatProperties, create_chat from repository.chat.get_chat_by_id import get_chat_by_id from repository.chat.get_chat_history import get_chat_history @@ -57,10 +61,7 @@ def delete_chat_from_db(commons, chat_id): def fetch_user_stats(commons, user, date): - response = ( - commons["db"] - .get_user_stats(user.email, date) - ) + response = commons["db"].get_user_stats(user.email, date) userItem = next(iter(response.data or []), {"requests_count": 0}) return userItem @@ -173,9 +174,10 @@ async def create_question_handler( brain = Brain(id=brain_id) if not current_user.user_openai_api_key: - brain_details = brain.get_brain_details() - if brain_details: - current_user.user_openai_api_key = brain_details["openai_api_key"] + if brain_id: + brain_details = get_brain_details(brain_id) + if brain_details: + current_user.user_openai_api_key = brain_details.openai_api_key if not current_user.user_openai_api_key: user_identity = get_user_identity(current_user.id) @@ -199,7 +201,7 @@ async def create_question_handler( LLMSettings() if not brain_id: - brain_id = get_default_user_brain_or_create_new(current_user).id + brain_id = get_default_user_brain_or_create_new(current_user).brain_id gpt_answer_generator = OpenAIBrainPicking( chat_id=str(chat_id), @@ -244,10 +246,10 @@ async def create_stream_question_handler( current_user.user_openai_api_key = request.headers.get("Openai-Api-Key") brain = Brain(id=brain_id) - if not current_user.user_openai_api_key: - brain_details = brain.get_brain_details() + if not current_user.user_openai_api_key and brain_id: + brain_details = get_brain_details(brain_id) if brain_details: - current_user.user_openai_api_key = brain_details["openai_api_key"] + current_user.user_openai_api_key = brain_details.openai_api_key if not current_user.user_openai_api_key: user_identity = get_user_identity(current_user.id) @@ -270,7 +272,7 @@ async def create_stream_question_handler( logger.info(f"Streaming request for {chat_question.model}") check_user_limit(current_user) if not brain_id: - brain_id = get_default_user_brain_or_create_new(current_user).id + brain_id = get_default_user_brain_or_create_new(current_user).brain_id gpt_answer_generator = OpenAIBrainPicking( chat_id=str(chat_id), diff --git a/backend/core/routes/prompt_routes.py b/backend/core/routes/prompt_routes.py index f0825b7a3..bc323f726 100644 --- a/backend/core/routes/prompt_routes.py +++ b/backend/core/routes/prompt_routes.py @@ -1,7 +1,16 @@ +from uuid import UUID + from auth import AuthBearer from fastapi import APIRouter, Depends +from models.databases.supabase.prompts import ( + CreatePromptProperties, + PromptUpdatableProperties, +) from models.prompt import Prompt +from repository.prompt.create_prompt import create_prompt +from repository.prompt.get_prompt_by_id import get_prompt_by_id from repository.prompt.get_public_prompts import get_public_prompts +from repository.prompt.update_prompt_by_id import update_prompt_by_id prompt_router = APIRouter() @@ -13,3 +22,36 @@ async def get_prompts() -> list[Prompt]: """ return get_public_prompts() + + +@prompt_router.get( + "/prompts/{prompt_id}", dependencies=[Depends(AuthBearer())], tags=["Prompt"] +) +async def get_prompt(prompt_id: UUID) -> Prompt | None: + """ + Retrieve a prompt by its id + """ + + return get_prompt_by_id(prompt_id) + + +@prompt_router.put( + "/prompts/{prompt_id}", dependencies=[Depends(AuthBearer())], tags=["Prompt"] +) +async def update_prompt( + prompt_id: UUID, prompt: PromptUpdatableProperties +) -> Prompt | None: + """ + Update a prompt by its id + """ + + return update_prompt_by_id(prompt_id, prompt) + + +@prompt_router.post("/prompts", dependencies=[Depends(AuthBearer())], tags=["Prompt"]) +async def create_prompt_route(prompt: CreatePromptProperties) -> Prompt | None: + """ + Create a prompt by its id + """ + + return create_prompt(prompt) diff --git a/backend/core/routes/subscription_routes.py b/backend/core/routes/subscription_routes.py index 16efb5c63..4763c5430 100644 --- a/backend/core/routes/subscription_routes.py +++ b/backend/core/routes/subscription_routes.py @@ -7,6 +7,9 @@ from models.brains import Brain from models.brains_subscription_invitations import BrainSubscription from models.users import User from pydantic import BaseModel +from repository.brain.create_brain_user import create_brain_user +from repository.brain.get_brain_details import get_brain_details +from repository.brain.get_brain_for_user import get_brain_for_user from repository.brain.update_user_rights import update_brain_user_rights from repository.brain_subscription.resend_invitation_email import ( resend_invitation_email, @@ -121,14 +124,14 @@ async def remove_user_subscription( brain = Brain( id=brain_id, ) - user_brain = brain.get_brain_for_user(current_user.id) + user_brain = get_brain_for_user(current_user.id, brain_id) if user_brain is None: raise HTTPException( status_code=403, detail="You don't have permission for this brain", ) - if user_brain.get("rights") != "Owner": + if user_brain.rights != "Owner": brain.delete_user_from_brain(current_user.id) else: brain_users = brain.get_brain_users() @@ -170,8 +173,7 @@ def get_user_invitation(brain_id: UUID, current_user: User = Depends(get_current detail="You have not been invited to this brain", ) - brain = Brain(id=brain_id) - brain_details = brain.get_brain_details() + brain_details = get_brain_details(brain_id) if brain_details is None: raise HTTPException( @@ -179,7 +181,7 @@ def get_user_invitation(brain_id: UUID, current_user: User = Depends(get_current detail="Brain not found while trying to get invitation", ) - return {"name": brain_details["name"], "rights": invitation["rights"]} + return {"name": brain_details.name, "rights": invitation["rights"]} @subscription_router.post( @@ -208,9 +210,11 @@ async def accept_invitation( raise HTTPException(status_code=404, detail="Invitation not found") try: - brain = Brain(id=brain_id) - brain.create_brain_user( - user_id=current_user.id, rights=invitation["rights"], default_brain=False + create_brain_user( + user_id=current_user.id, + brain_id=brain_id, + rights=invitation["rights"], + is_default_brain=False, ) except Exception as e: raise HTTPException(status_code=400, detail=f"Error adding user to brain: {e}") @@ -299,8 +303,8 @@ def update_brain_subscription( ) # check if user is not an editor trying to update an owner right which is not allowed - current_invitation = brain.get_brain_for_user(user_id) - if current_invitation is not None and current_invitation.get("rights") == "Owner": + current_invitation = get_brain_for_user(user_id, brain_id) + if current_invitation is not None and current_invitation.rights == "Owner": try: validate_brain_authorization( brain_id, diff --git a/backend/core/routes/upload_routes.py b/backend/core/routes/upload_routes.py index 42d087126..b2a6ab595 100644 --- a/backend/core/routes/upload_routes.py +++ b/backend/core/routes/upload_routes.py @@ -7,6 +7,7 @@ from models.brains import Brain from models.files import File from models.settings import common_dependencies from models.users import User +from repository.brain.get_brain_details import get_brain_details from repository.user_identity.get_user_identity import get_user_identity from utils.file import convert_bytes, get_file_size from utils.processors import filter_file @@ -62,9 +63,9 @@ async def upload_file( else: openai_api_key = request.headers.get("Openai-Api-Key", None) if openai_api_key is None: - brain_details = brain.get_brain_details() + brain_details = get_brain_details(brain_id) if brain_details: - openai_api_key = brain_details["openai_api_key"] + openai_api_key = brain_details.openai_api_key if openai_api_key is None: openai_api_key = get_user_identity(current_user.id).openai_api_key diff --git a/backend/core/routes/user_routes.py b/backend/core/routes/user_routes.py index c54fc455f..3c13a05af 100644 --- a/backend/core/routes/user_routes.py +++ b/backend/core/routes/user_routes.py @@ -3,10 +3,11 @@ import time from auth import AuthBearer, get_current_user from fastapi import APIRouter, Depends, Request -from models.brains import Brain, get_default_user_brain +from models.brains import Brain from models.settings import BrainRateLimiting from models.user_identity import UserIdentity from models.users import User +from repository.brain.get_default_user_brain import get_user_default_brain from repository.user_identity.get_user_identity import get_user_identity from repository.user_identity.update_user_identity import ( UserIdentityUpdatableProperties, @@ -47,10 +48,10 @@ async def get_user_endpoint( date = time.strftime("%Y%m%d") max_requests_number = os.getenv("MAX_REQUESTS_NUMBER") requests_stats = current_user.get_user_request_stats() - default_brain = get_default_user_brain(current_user) + default_brain = get_user_default_brain(current_user.id) if default_brain: - defaul_brain_size = Brain(id=default_brain["id"]).brain_size + defaul_brain_size = Brain(id=default_brain.brain_id).brain_size else: defaul_brain_size = 0 diff --git a/backend/core/tests/test_brains.py b/backend/core/tests/test_brains.py index 66460d355..d365f252e 100644 --- a/backend/core/tests/test_brains.py +++ b/backend/core/tests/test_brains.py @@ -1,7 +1,7 @@ import random import string -from models.brains import get_default_user_brain +from repository.brain.get_default_user_brain import get_user_default_brain def test_retrieve_default_brain(client, api_key): @@ -211,7 +211,6 @@ def test_set_as_default_brain_endpoint(client, api_key): user_info = response.json() user_id = user_info["id"] - default_brain = get_default_user_brain(user_id) + default_brain = get_user_default_brain(user_id) assert default_brain is not None - assert default_brain["id"] == brain_id - assert default_brain["default_brain"] is True + assert default_brain.brain_id == brain_id diff --git a/scripts/20230802120700_add_prompt_id_to_brain.sql b/scripts/20230802120700_add_prompt_id_to_brain.sql new file mode 100644 index 000000000..f46bbc7cc --- /dev/null +++ b/scripts/20230802120700_add_prompt_id_to_brain.sql @@ -0,0 +1,19 @@ +BEGIN; + +-- Check if prompt_id column exists +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'brains' AND column_name = 'prompt_id') THEN + -- Add prompt_id column and reference the table prompts' id column + ALTER TABLE brains ADD COLUMN prompt_id UUID REFERENCES prompts(id); + END IF; +END $$; + +-- Update migrations table +INSERT INTO migrations (name) +SELECT '20230802120700_add_prompt_id_to_brain' +WHERE NOT EXISTS ( + SELECT 1 FROM migrations WHERE name = '20230802120700_add_prompt_id_to_brain' +); + +COMMIT; diff --git a/scripts/tables.sql b/scripts/tables.sql index de5342e4f..f8ad0f137 100644 --- a/scripts/tables.sql +++ b/scripts/tables.sql @@ -134,9 +134,11 @@ CREATE TABLE IF NOT EXISTS brains ( model TEXT, max_tokens INT, temperature FLOAT, - openai_api_key TEXT + openai_api_key TEXT, + prompt_id UUID REFERENCES prompts(id) ); + -- Create brains X users table CREATE TABLE IF NOT EXISTS brains_users ( brain_id UUID, @@ -210,7 +212,7 @@ CREATE TABLE IF NOT EXISTS migrations ( ); INSERT INTO migrations (name) -SELECT '20230701180101_add_prompts_table' +SELECT '20230802120700_add_prompt_id_to_brain' WHERE NOT EXISTS ( - SELECT 1 FROM migrations WHERE name = '20230701180101_add_prompts_table' + SELECT 1 FROM migrations WHERE name = '20230802120700_add_prompt_id_to_brain' );