Merge pull request #357 from MuckRock/organization-collective-fixes

Some fixes for organization-collectives

Author: duckduckgrayduck

documentcloud/organizations/models.py

@@ -249,14 +249,31 @@ 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
         for group in self.groups.filter(share_resources=True):
             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"""

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