diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 35ad78f0..131902b6 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -249,7 +249,7 @@ def get_total_number_ai_credits(self): return number_ai_credits def get_total_monthly_ai_credits(self): - """Get total monthly AI credits including parent and groups""" + """Get total monthly AI credits remaining including parent and groups""" monthly_ai_credits = self.monthly_ai_credits if self.parent and self.parent.share_resources: monthly_ai_credits += self.parent.monthly_ai_credits @@ -257,6 +257,23 @@ def get_total_monthly_ai_credits(self): monthly_ai_credits += group.monthly_ai_credits return monthly_ai_credits + def get_total_monthly_ai_credits_allowance(self): + """ + Get the total monthly AI credits allowance, including parent and shared groups. + This is the amount that monthly_credits will reset to each month. + """ + total = self.ai_credits_per_month + + # Include parent if it shares resources + if self.parent and self.parent.share_resources: + total += self.parent.ai_credits_per_month + + # Include groups that share resources + for group in self.groups.filter(share_resources=True): + total += group.ai_credits_per_month + + return total + class AICreditLog(models.Model): """Log usage of AI Credits""" diff --git a/documentcloud/organizations/serializers.py b/documentcloud/organizations/serializers.py index e075073c..d8995fe4 100644 --- a/documentcloud/organizations/serializers.py +++ b/documentcloud/organizations/serializers.py @@ -19,21 +19,23 @@ class OrganizationSerializer(serializers.ModelSerializer): "Only viewable by organization members." ), ) - monthly_credits = serializers.IntegerField( - source="monthly_ai_credits", + monthly_credits = serializers.SerializerMethodField( + label=_("Monthly Credits"), read_only=True, help_text=( "Number of monthly premium credits this organization has left. " "This will reset to monthly_credit_allowance on credit_reset_date. " + "This includes shared credits from parents and groups. " "Only viewable be organization members." ), ) - purchased_credits = serializers.IntegerField( - source="number_ai_credits", + purchased_credits = serializers.SerializerMethodField( + label=_("Purchased Credits"), read_only=True, help_text=( "Number of purchased premium credits. " "These do not reset or expire. " + "This includes shared credits from parents and groups. " "Only viewable by organization members." ), ) @@ -45,8 +47,7 @@ class OrganizationSerializer(serializers.ModelSerializer): "Only viewable by organization members." ), ) - monthly_credit_allowance = serializers.IntegerField( - source="ai_credits_per_month", + monthly_credit_allowance = serializers.SerializerMethodField( read_only=True, help_text=( "The amount of credits that monthly_credits will reset to. " @@ -102,6 +103,15 @@ def get_plan(self, obj): else: return "Free" + def get_monthly_credits(self, obj): + return obj.get_total_monthly_ai_credits() + + def get_purchased_credits(self, obj): + return obj.get_total_number_ai_credits() + + def get_monthly_credit_allowance(self, obj): + return obj.get_total_monthly_ai_credits_allowance() + class AICreditSerializer(serializers.Serializer): """Serializer for the AI credit endpoint"""