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

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

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

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

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])

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:
         """

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

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

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]

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)

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])

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)

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)

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

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

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

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)

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"""
     <p>Brain {brain_name} has been shared with you by {inviter_email}.</p>

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)

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:

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",

backend/core/routes/authorizations/types.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class RoleEnum(str, Enum):
+    Viewer = "Viewer"
+    Editor = "Editor"
+    Owner = "Owner"

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."}

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),

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)

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,

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

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
 

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

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;

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