diff --git a/backend/models/brains.py b/backend/models/brains.py index 2378f3dc3..e0e1f341d 100644 --- a/backend/models/brains.py +++ b/backend/models/brains.py @@ -3,11 +3,12 @@ from typing import Any, List, Optional from uuid import UUID from logger import get_logger -from models.settings import CommonsDep, common_dependencies -from models.users import User from pydantic import BaseModel from utils.vectors import get_unique_files_from_vector_ids +from models.settings import CommonsDep, common_dependencies +from models.users import User + logger = get_logger(__name__) @@ -49,6 +50,32 @@ class Brain(BaseModel): 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.commons["supabase"] + .table("brains_users") + .select("id:brain_id, *") + .filter("brain_id", "eq", self.id) + .execute() + ) + return response.data + + # TODO: move this to a brand new BrainService + def delete_user_from_brain(self, user_id): + results = ( + self.commons["supabase"] + .table("brains_users") + .select("*") + .match({"brain_id": self.id, "user_id": user_id}) + .execute() + ) + + if len(results.data) != 0: + self.commons["supabase"].table("brains_users").delete().match( + {"brain_id": self.id, "user_id": user_id} + ).execute() + def get_user_brains(self, user_id): response = ( self.commons["supabase"] diff --git a/backend/routes/subscription_routes.py b/backend/routes/subscription_routes.py index 78c8a4eec..5c21ac510 100644 --- a/backend/routes/subscription_routes.py +++ b/backend/routes/subscription_routes.py @@ -3,6 +3,7 @@ from uuid import UUID from auth.auth_bearer import get_current_user from fastapi import APIRouter, Depends, HTTPException +from models.brains import Brain from models.brains_subscription_invitations import BrainSubscription from models.users import User @@ -10,12 +11,19 @@ subscription_router = APIRouter() @subscription_router.post("/brain/{brain_id}/subscription") -async def invite_user_to_brain(brain_id: UUID, users: List[dict], current_user: User = Depends(get_current_user)): +async def invite_user_to_brain( + brain_id: UUID, users: List[dict], current_user: User = Depends(get_current_user) +): # TODO: Ensure the current user has permissions to invite users to this brain - + for user in users: - subscription = BrainSubscription(brain_id=brain_id, email=user['email'], rights=user['rights'], inviter_email=current_user.email or "Quivr") - + subscription = BrainSubscription( + brain_id=brain_id, + email=user["email"], + rights=user["rights"], + inviter_email=current_user.email or "Quivr", + ) + try: subscription.create_or_update_subscription_invitation() subscription.resend_invitation_email() @@ -23,3 +31,41 @@ async def invite_user_to_brain(brain_id: UUID, users: List[dict], current_user: raise HTTPException(status_code=400, detail=f"Error inviting user: {e}") return {"message": "Invitations sent successfully"} + + +@subscription_router.delete( + "/brain/{brain_id}/subscription", +) +async def remove_user_subscription( + brain_id: UUID, current_user: User = Depends(get_current_user) +): + """ + Remove a user's subscription to a brain + """ + brain = Brain( + id=brain_id, + ) + user_brain = brain.get_brain_for_user(current_user.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": + brain.delete_user_from_brain(current_user.id) + else: + brain_other_users = brain.get_brain_users() + brain_other_owners = [ + brain + for brain in brain_other_users + if brain["rights"] == "Owner" + and str(brain["user_id"]) != str(current_user.id) + ] + + if len(brain_other_owners) == 0: + brain.delete_brain(current_user.id) + else: + brain.delete_user_from_brain(current_user.id) + + return {"message": f"Subscription removed successfully from brain {brain_id}"}