diff --git a/django/aiservice/authapp/auth.py b/django/aiservice/authapp/auth.py index 9b9ef44c7..1dcc0cc4d 100644 --- a/django/aiservice/authapp/auth.py +++ b/django/aiservice/authapp/auth.py @@ -6,7 +6,7 @@ from ninja.errors import HttpError from ninja.security import HttpBearer from authapp.auth_utils import hash_api_key, instance_for_api_key -from authapp.models import CFAPIKeys, Subscriptions +from authapp.models import CFAPIKeys, Organizations, Subscriptions class AuthenticatedRequest(Protocol): @@ -21,15 +21,16 @@ class AuthenticatedRequest(Protocol): should_log_features: bool # whether to log optimization features -async def check_subscription_status(user_id, tier) -> bool: +async def check_subscription_status(user_id, tier, organization_id=None) -> bool: """Check if a user has a premium subscription that doesn't require feature logging. Args: user_id: The ID of the user to check tier: The user's tier if already available + organization_id: The ID of the user's organization if available Returns: - bool: False if features should not be logged (premium user), True otherwise + bool: False if features should not be logged (premium user or paid org), True otherwise """ # If tier is already set, no need to check subscription @@ -37,6 +38,15 @@ async def check_subscription_status(user_id, tier) -> bool: return False try: + # Check if user belongs to a paid organization + if organization_id: + org = await Organizations.objects.filter(id=organization_id).afirst() + if org and org.name == "codeflash-ai": + return True + if org and org.subscription: + # Paid organization - don't log features + return False + subscription = await Subscriptions.objects.filter(user_id=user_id).afirst() if subscription and subscription.plan_type.lower() in ["pro", "enterprise"]: # Premium users for CF- don't log features @@ -65,7 +75,9 @@ class AuthBearer(HttpBearer): request.tier = api_key_instance.tier request.api_key_id = api_key_instance.id request.organization_id = api_key_instance.organization_id - request.should_log_features = await check_subscription_status(user_id=request.user, tier=request.tier) + request.should_log_features = await check_subscription_status( + user_id=request.user, tier=request.tier, organization_id=request.organization_id + ) return token print("THIS SHOULD NOT HAPPEN! More than one users found in the db with the same api key!")