feat: remove api brain secrets and schemas on delete (#1621)

Issue: #1573

backend/models/brain_entity.py

@@ -49,3 +49,10 @@ class PublicBrain(BaseModel):
     description: Optional[str]
     number_of_subscribers: int = 0
     last_update: str
+
+
+class BrainUser(BaseModel):
+    id: UUID
+    user_id: UUID
+    rights: RoleEnum
+    default_brain: bool = False

backend/models/brains.py

@@ -49,40 +49,6 @@ def create(cls, *args, **kwargs):
             commons=commons, *args, **kwargs  # pyright: ignore reportPrivateUsage=none
         )  # pyright: ignore reportPrivateUsage=none
 
-    # TODO: move this to a brand new BrainService
-    def get_brain_users(self):
-        response = (
-            self.supabase_client.table("brains_users")
-            .select("id:brain_id, *")
-            .filter("brain_id", "eq", self.id)  # type: ignore
-            .execute()
-        )
-        return response.data
-
-    # TODO: move this to a brand new BrainService
-    def delete_user_from_brain(self, user_id):
-        results = (
-            self.supabase_client.table("brains_users")
-            .select("*")
-            .match({"brain_id": self.id, "user_id": user_id})
-            .execute()
-        )
-
-        if len(results.data) != 0:
-            self.supabase_client.table("brains_users").delete().match(
-                {"brain_id": self.id, "user_id": user_id}
-            ).execute()
-
-    def delete_brain(self, user_id):
-        results = self.supabase_db.delete_brain_user_by_id(user_id, self.id)  # type: ignore
-
-        if len(results) == 0:
-            return {"message": "You are not the owner of this brain."}
-        else:
-            self.supabase_db.delete_brain_vector(self.id)  # type: ignore
-            self.supabase_db.delete_brain_users(self.id)  # type: ignore
-            self.supabase_db.delete_brain(self.id)  # type: ignore
-
     def create_brain_vector(self, vector_id, file_sha1):
         return self.supabase_db.create_brain_vector(self.id, vector_id, file_sha1)  # type: ignore
 

backend/models/databases/supabase/brains.py

@@ -2,7 +2,13 @@
 from uuid import UUID
 
 from logger import get_logger
-from models.brain_entity import BrainEntity, BrainType, MinimalBrainEntity, PublicBrain
+from models.brain_entity import (
+    BrainEntity,
+    BrainType,
+    BrainUser,
+    MinimalBrainEntity,
+    PublicBrain,
+)
 from models.databases.repository import Repository
 from models.databases.supabase.api_brain_definition import (
     CreateApiBrainDefinition,
@@ -337,3 +343,13 @@ def get_brain_subscribers_count(self, brain_id: UUID) -> int:
         if len(response) == 0:
             raise ValueError(f"Brain with id {brain_id} does not exist.")
         return response[0]["count"]
+
+    def get_brain_users(self, brain_id: UUID) -> list[BrainUser]:
+        response = (
+            self.db.table("brains_users")
+            .select("id:brain_id, *")
+            .filter("brain_id", "eq", str(brain_id))
+            .execute()
+        )
+
+        return [BrainUser(**item) for item in response.data]

backend/repository/brain/delete_brain.py

@@ -0,0 +1,36 @@
+from uuid import UUID
+
+from fastapi import HTTPException
+from models.brain_entity import BrainType
+from models.settings import get_supabase_db
+
+from repository.api_brain_definition.delete_api_brain_definition import (
+    delete_api_brain_definition,
+)
+from repository.brain import get_brain_by_id
+from repository.brain.delete_brain_secrets import delete_brain_secrets
+from repository.knowledge.remove_brain_all_knowledge import (
+    remove_brain_all_knowledge,
+)
+
+
+def delete_brain(brain_id: UUID) -> dict[str, str]:
+    supabase_db = get_supabase_db()
+
+    brain_to_delete = get_brain_by_id(brain_id=brain_id)
+    if brain_to_delete is None:
+        raise HTTPException(status_code=404, detail="Brain not found.")
+
+    if brain_to_delete.brain_type == BrainType.API:
+        delete_brain_secrets(
+            brain_id=brain_id,
+        )
+        delete_api_brain_definition(brain_id=brain_id)
+    else:
+        remove_brain_all_knowledge(brain_id)
+
+    supabase_db.delete_brain_vector(str(brain_id))
+    supabase_db.delete_brain_users(str(brain_id))
+    supabase_db.delete_brain(str(brain_id))
+
+    return {"message": "Brain deleted."}

backend/repository/brain/delete_brain_secrets.py

@@ -0,0 +1,30 @@
+from uuid import UUID
+
+from fastapi import HTTPException
+from models.settings import get_supabase_db
+
+from repository.api_brain_definition.get_api_brain_definition import (
+    get_api_brain_definition,
+)
+from repository.external_api_secret import delete_secret
+
+
+def delete_brain_secrets(brain_id: UUID) -> None:
+    supabase_db = get_supabase_db()
+
+    brain_definition = get_api_brain_definition(brain_id=brain_id)
+
+    if brain_definition is None:
+        raise HTTPException(status_code=404, detail="Brain definition not found.")
+
+    secrets = brain_definition.secrets
+
+    if len(secrets) > 0:
+        brain_users = supabase_db.get_brain_users(brain_id=brain_id)
+        for user in brain_users:
+            for secret in secrets:
+                delete_secret(
+                    user_id=user.user_id,
+                    brain_id=brain_id,
+                    secret_name=secret.name,
+                )

backend/repository/brain/delete_brain_user.py

@@ -1,10 +1,34 @@
 from uuid import UUID
 
+from fastapi import HTTPException
+from models.brain_entity import BrainType
 from models.settings import get_supabase_db
 
+from repository.api_brain_definition.get_api_brain_definition import (
+    get_api_brain_definition,
+)
+from repository.brain.get_brain_by_id import get_brain_by_id
+from repository.external_api_secret.delete_secret import delete_secret
+
 
 def delete_brain_user(user_id: UUID, brain_id: UUID) -> None:
     supabase_db = get_supabase_db()
+    brain_to_delete_user_from = get_brain_by_id(brain_id=brain_id)
+    if brain_to_delete_user_from is None:
+        raise HTTPException(status_code=404, detail="Brain not found.")
+
+    if brain_to_delete_user_from.brain_type == BrainType.API:
+        brain_definition = get_api_brain_definition(brain_id=brain_id)
+        if brain_definition is None:
+            raise HTTPException(status_code=404, detail="Brain definition not found.")
+        secrets = brain_definition.secrets
+        for secret in secrets:
+            delete_secret(
+                user_id=user_id,
+                brain_id=brain_id,
+                secret_name=secret.name,
+            )
+
     supabase_db.delete_brain_user_by_id(
         user_id=user_id,
         brain_id=brain_id,

backend/repository/brain/get_brain_users.py

@@ -0,0 +1,10 @@
+from uuid import UUID
+
+from models.brain_entity import BrainUser
+from models.settings import get_supabase_db
+
+
+def get_brain_users(brain_id: UUID) -> list[BrainUser]:
+    supabase_db = get_supabase_db()
+
+    return supabase_db.get_brain_users(brain_id)

backend/repository/external_api_secret/delete_secret.py

@@ -10,7 +10,7 @@ def delete_secret(user_id: UUID, brain_id: UUID, secret_name: str) -> bool:
     response = supabase_client.rpc(
         "delete_secret",
         {
-            "name": build_secret_unique_name(
+            "secret_name": build_secret_unique_name(
                 user_id=user_id, brain_id=brain_id, secret_name=secret_name
             ),
         },

backend/routes/subscription_routes.py

@@ -3,7 +3,7 @@
 
 from auth.auth_bearer import AuthBearer, get_current_user
 from fastapi import APIRouter, Depends, HTTPException
-from models import Brain, BrainSubscription, PromptStatusEnum, UserIdentity
+from models import BrainSubscription, PromptStatusEnum, UserIdentity
 from pydantic import BaseModel
 from repository.brain import (
     create_brain_user,
@@ -12,16 +12,15 @@
     get_brain_for_user,
     update_brain_user_rights,
 )
+from repository.brain.delete_brain import delete_brain
 from repository.brain.delete_brain_user import delete_brain_user
+from repository.brain.get_brain_users import get_brain_users
 from repository.brain_subscription import (
     SubscriptionInvitationService,
     resend_invitation_email,
 )
-from repository.knowledge.remove_brain_all_knowledge import (
-    remove_brain_all_knowledge,
-)
 from repository.prompt import delete_prompt_by_id, get_prompt_by_id
-from repository.user import get_user_email_by_user_id, get_user_id_by_user_email
+from repository.user import get_user_id_by_user_email
 
 from routes.authorizations.brain_authorization import (
     RoleEnum,
@@ -99,29 +98,6 @@ def invite_users_to_brain(
         Depends(has_brain_authorization([RoleEnum.Owner, RoleEnum.Editor])),
     ],
 )
-def get_brain_users(
-    brain_id: UUID,
-):
-    """
-    Get all users for a brain
-    """
-    brain = Brain(
-        id=brain_id,
-    )
-    brain_users = brain.get_brain_users()
-
-    brain_access_list = []
-
-    for brain_user in brain_users:
-        brain_access = {}
-        # TODO: find a way to fetch user email concurrently
-        brain_access["email"] = get_user_email_by_user_id(brain_user["user_id"])
-        brain_access["rights"] = brain_user["rights"]
-        brain_access_list.append(brain_access)
-
-    return brain_access_list
-
-
 @subscription_router.delete(
     "/brains/{brain_id}/subscription",
 )
@@ -139,9 +115,6 @@ async def remove_user_subscription(
             detail="Brain not found while trying to delete",
         )
 
-    brain = Brain(
-        id=brain_id,
-    )
     user_brain = get_brain_for_user(current_user.id, brain_id)
     if user_brain is None:
         raise HTTPException(
@@ -150,21 +123,21 @@ async def remove_user_subscription(
         )
 
     if user_brain.rights != "Owner":
-        brain.delete_user_from_brain(current_user.id)
+        delete_brain_user(current_user.id, brain_id)
     else:
-        brain_users = brain.get_brain_users()
+        brain_users = get_brain_users(
+            brain_id=brain_id,
+        )
         brain_other_owners = [
             brain
             for brain in brain_users
-            if brain["rights"] == "Owner"
-            and str(brain["user_id"]) != str(current_user.id)
+            if brain.rights == "Owner" and str(brain.user_id) != str(current_user.id)
         ]
 
         if len(brain_other_owners) == 0:
-            # Delete its prompt if it's private
-
-            remove_brain_all_knowledge(brain_id)
-            brain.delete_brain(current_user.id)
+            delete_brain(
+                brain_id=brain_id,
+            )
             if targeted_brain.prompt_id:
                 brain_to_delete_prompt = get_prompt_by_id(targeted_brain.prompt_id)
                 if brain_to_delete_prompt is not None and (
@@ -173,7 +146,10 @@ async def remove_user_subscription(
                     delete_prompt_by_id(targeted_brain.prompt_id)
 
         else:
-            brain.delete_user_from_brain(current_user.id)
+            delete_brain_user(
+                current_user.id,
+                brain_id,
+            )
 
     return {"message": f"Subscription removed successfully from brain {brain_id}"}
 
@@ -321,10 +297,6 @@ def update_brain_subscription(
             detail="User not found",
         )
 
-    brain = Brain(
-        id=brain_id,
-    )
-
     # check if user is an editor but trying to give high level permissions
     if subscription.rights == "Owner":
         try:
@@ -363,7 +335,7 @@ def update_brain_subscription(
                 current_user.id,
                 RoleEnum.Owner,
             )
-            brain.delete_user_from_brain(user_id)
+            delete_brain_user(user_id, brain_id)
         except HTTPException:
             raise HTTPException(
                 status_code=403,