feat(billing): handle limit aggregation for llm requests DEV-1031 (#6291)

jamesrkiger · web-flow · commit f87d0fa3978b · 2025-11-10T13:41:22.000-05:00
### 📣 Summary
Enables aggregation of automated QA billing limits and the inclusion of
these limits in service usage data returned by the API.
diff --git a/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py b/kobo/apps/openrosa/apps/api/tests/viewsets/test_xform_submission_api.py
@@ -79,7 +79,7 @@ def test_query_counts(self):
             # so should find a way to keep that out of this count
             if settings.STRIPE_ENABLED:
                 # But because of cache, sometimes goes down to 62
-                expected_queries = FuzzyInt(62, 84)
+                expected_queries = FuzzyInt(62, 90)
             with self.assertNumQueries(expected_queries):
                 self.view(request)
 
diff --git a/kobo/apps/organizations/constants.py b/kobo/apps/organizations/constants.py
@@ -1,16 +1,6 @@
 from django.db import models
 
 
-# Temporary const for code that iterates over usage types
-# TODO: Remove when Stripe code has been updated to
-# handle LLM usage type
-class SupportedUsageType(models.TextChoices):
-    SUBMISSION = 'submission'
-    STORAGE_BYTES = 'storage_bytes'
-    MT_CHARACTERS = 'mt_characters'
-    ASR_SECONDS = 'asr_seconds'
-
-
 class UsageType(models.TextChoices):
     SUBMISSION = 'submission'
     STORAGE_BYTES = 'storage_bytes'
diff --git a/kobo/apps/organizations/tests/test_organizations_service_usage_api.py b/kobo/apps/organizations/tests/test_organizations_service_usage_api.py
@@ -84,8 +84,8 @@ def test_check_api_response_without_stripe(self):
         assert response.data['total_nlp_usage']['asr_seconds_all_time'] == 4728
         assert response.data['total_nlp_usage']['mt_characters_current_period'] == 5473
         assert response.data['total_nlp_usage']['mt_characters_all_time'] == 6726
-        assert response.data['total_nlp_usage']['llm_requests_current_period'] == 20
-        assert response.data['total_nlp_usage']['llm_requests_all_time'] == 70
+        assert response.data['total_nlp_usage']['llm_requests_current_period'] == 30
+        assert response.data['total_nlp_usage']['llm_requests_all_time'] == 80
         assert response.data['total_storage_bytes'] == self.expected_file_size()
 
         # Without stripe, there are no usage limits and
diff --git a/kobo/apps/stripe/models.py b/kobo/apps/stripe/models.py
@@ -9,7 +9,7 @@
 from djstripe.models import Charge, Price, Subscription
 
 from kobo.apps.kobo_auth.shortcuts import User
-from kobo.apps.organizations.constants import SupportedUsageType, UsageType
+from kobo.apps.organizations.constants import UsageType
 from kobo.apps.organizations.models import Organization
 from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES
 from kobo.apps.stripe.utils.subscription_limits import get_default_add_on_limits
@@ -176,6 +176,26 @@ def get_organizations_totals(organizations: list['Organization'] = None):
                     ),
                     0,
                 ),
+                llm_requests_remaining=Coalesce(
+                    Sum(
+                        Cast(
+                            PlanAddOn.get_limits_remaining_field(
+                                UsageType.LLM_REQUESTS
+                            ),
+                            output_field=IntegerField(),
+                        )
+                    ),
+                    0,
+                ),
+                total_llm_requests_limit=Coalesce(
+                    Sum(
+                        Cast(
+                            PlanAddOn.get_usage_limits_field(UsageType.LLM_REQUESTS),
+                            output_field=IntegerField(),
+                        )
+                    ),
+                    0,
+                ),
                 submission_remaining=Coalesce(
                     Sum(
                         Cast(
@@ -311,6 +331,6 @@ class ExceededLimitCounter(models.Model):
         on_delete=models.CASCADE,
     )
     days = models.PositiveSmallIntegerField(default=0)
-    limit_type = models.CharField(choices=SupportedUsageType.choices, max_length=20)
+    limit_type = models.CharField(choices=UsageType.choices, max_length=20)
     date_created = models.DateTimeField(auto_now_add=True)
     date_modified = models.DateTimeField(auto_now=True)
diff --git a/kobo/apps/stripe/tests/test_stripe_utils.py b/kobo/apps/stripe/tests/test_stripe_utils.py
@@ -14,7 +14,7 @@
 from model_bakery import baker
 
 from kobo.apps.kobo_auth.shortcuts import User
-from kobo.apps.organizations.constants import SupportedUsageType, UsageType
+from kobo.apps.organizations.constants import UsageType
 from kobo.apps.organizations.models import Organization
 from kobo.apps.organizations.utils import get_billing_dates
 from kobo.apps.stripe.models import ExceededLimitCounter
@@ -67,75 +67,57 @@ def setUpTestData(cls):
 
     def test_get_organization_subscription_limits(self):
         free_plan = generate_free_plan()
+        first_paid_plan_limits = {
+            f'{UsageType.MT_CHARACTERS}_limit': '100',
+            f'{UsageType.ASR_SECONDS}_limit': '200',
+            f'{UsageType.LLM_REQUESTS}_limit': '300',
+            f'{UsageType.SUBMISSION}_limit': '400',
+            f'{UsageType.STORAGE_BYTES}_limit': '500',
+        }
+        # second plan for org2 where we append (not add) 1 to all the limits
+        second_paid_plan_limits = {
+            key: f'{value}1' for key, value in first_paid_plan_limits.items()
+        }
         product_metadata = {
-            f'{UsageType.MT_CHARACTERS}_limit': '1234',
-            f'{UsageType.ASR_SECONDS}_limit': '5678',
-            f'{UsageType.SUBMISSION}_limit': '91011',
-            f'{UsageType.STORAGE_BYTES}_limit': '121314',
             'product_type': 'plan',
             'plan_type': 'enterprise',
         }
-        # create a second plan for org2 where we append (not add) 1 to all the limits
+        first_product_metadata = {
+            **first_paid_plan_limits,
+            **product_metadata,
+        }
         second_product_metadata = {
-            key: f'{value}1' if key.endswith('limit') else value
-            for key, value in product_metadata.items()
+            **second_paid_plan_limits,
+            **product_metadata,
         }
-        generate_plan_subscription(self.organization, metadata=product_metadata)
+
+        generate_plan_subscription(self.organization, metadata=first_product_metadata)
         generate_plan_subscription(
             self.second_organization, metadata=second_product_metadata
         )
         all_limits = get_organizations_subscription_limits()
-        assert (
-            all_limits[self.organization.id][f'{UsageType.MT_CHARACTERS}_limit'] == 1234
-        )
-        assert (
-            all_limits[self.second_organization.id][f'{UsageType.MT_CHARACTERS}_limit']
-            == 12341
-        )
-        assert (
-            all_limits[self.organization.id][f'{UsageType.ASR_SECONDS}_limit'] == 5678
-        )
-        assert (
-            all_limits[self.second_organization.id][f'{UsageType.ASR_SECONDS}_limit']
-            == 56781
-        )
-        assert (
-            all_limits[self.organization.id][f'{UsageType.SUBMISSION}_limit'] == 91011
-        )
-        assert (
-            all_limits[self.second_organization.id][f'{UsageType.SUBMISSION}_limit']
-            == 910111
-        )
-        assert (
-            all_limits[self.organization.id][f'{UsageType.STORAGE_BYTES}_limit']
-            == 121314
-        )
-        assert (
-            all_limits[self.second_organization.id][f'{UsageType.STORAGE_BYTES}_limit']
-            == 1213141
-        )
 
         other_orgs = Organization.objects.exclude(
             id__in=[self.organization.id, self.second_organization.id]
         )
-        for org in other_orgs:
-            assert all_limits[org.id][f'{UsageType.MT_CHARACTERS}_limit'] == int(
-                free_plan.metadata[f'{UsageType.MT_CHARACTERS}_limit']
-            )
-            assert all_limits[org.id][f'{UsageType.ASR_SECONDS}_limit'] == int(
-                free_plan.metadata[f'{UsageType.ASR_SECONDS}_limit']
-            )
-            assert all_limits[org.id][f'{UsageType.SUBMISSION}_limit'] == int(
-                free_plan.metadata[f'{UsageType.SUBMISSION}_limit']
-            )
-            assert all_limits[org.id][f'{UsageType.STORAGE_BYTES}_limit'] == int(
-                free_plan.metadata[f'{UsageType.STORAGE_BYTES}_limit']
+        for usage_type, _ in UsageType.choices:
+            assert all_limits[self.organization.id][f'{usage_type}_limit'] == float(
+                first_paid_plan_limits[f'{usage_type}_limit']
             )
+            assert all_limits[self.second_organization.id][
+                f'{usage_type}_limit'
+            ] == float(second_paid_plan_limits[f'{usage_type}_limit'])
+
+            for org in other_orgs:
+                assert all_limits[org.id][f'{usage_type}_limit'] == float(
+                    free_plan.metadata[f'{usage_type}_limit']
+                )
 
     def test__prioritizes_price_metadata(self):
         product_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '1',
             f'{UsageType.ASR_SECONDS}_limit': '1',
+            f'{UsageType.LLM_REQUESTS}_limit': '1',
             f'{UsageType.SUBMISSION}_limit': '1',
             f'{UsageType.STORAGE_BYTES}_limit': '1',
             'product_type': 'plan',
@@ -144,14 +126,15 @@ def test__prioritizes_price_metadata(self):
         price_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '2',
             f'{UsageType.ASR_SECONDS}_limit': '2',
+            f'{UsageType.LLM_REQUESTS}_limit': '2',
             f'{UsageType.SUBMISSION}_limit': '2',
             f'{UsageType.STORAGE_BYTES}_limit': '2',
         }
         generate_plan_subscription(
             self.organization, metadata=product_metadata, price_metadata=price_metadata
         )
         limits = get_paid_subscription_limits([self.organization.id]).first()
-        for usage_type, _ in SupportedUsageType.choices:
+        for usage_type, _ in UsageType.choices:
             assert limits[f'{usage_type}_limit'] == '2'
 
     def test_get_subscription_limits_takes_most_recent_active_subscriptions(self):
@@ -477,6 +460,7 @@ def test_get_org_effective_limits(self, include_onetime_addons):
         plan_product_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '1',
             f'{UsageType.ASR_SECONDS}_limit': '1',
+            f'{UsageType.LLM_REQUESTS}_limit': '1',
             f'{UsageType.SUBMISSION}_limit': '1',
             f'{UsageType.STORAGE_BYTES}_limit': '1',
             'product_type': 'plan',
@@ -485,6 +469,7 @@ def test_get_org_effective_limits(self, include_onetime_addons):
         product_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '2',
             f'{UsageType.ASR_SECONDS}_limit': '2',
+            f'{UsageType.LLM_REQUESTS}_limit': '2',
             f'{UsageType.SUBMISSION}_limit': '2',
             f'{UsageType.STORAGE_BYTES}_limit': '2',
             'product_type': 'plan',
@@ -497,6 +482,7 @@ def test_get_org_effective_limits(self, include_onetime_addons):
         one_time_nlp_addon_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '15',
             f'{UsageType.ASR_SECONDS}_limit': '20',
+            f'{UsageType.LLM_REQUESTS}_limit': '20',
         }
         nlp_addon = _create_one_time_addon_product(one_time_nlp_addon_metadata)
         customer = baker.make(Customer, subscriber=self.organization)
@@ -526,11 +512,17 @@ def test_get_org_effective_limits(self, include_onetime_addons):
         assert (
             results[self.second_organization.id][f'{UsageType.ASR_SECONDS}_limit'] == 2
         )
+        assert (
+            results[self.second_organization.id][f'{UsageType.LLM_REQUESTS}_limit'] == 2
+        )
         if include_onetime_addons:
             assert (
                 results[self.organization.id][f'{UsageType.MT_CHARACTERS}_limit'] == 16
             )
             assert results[self.organization.id][f'{UsageType.ASR_SECONDS}_limit'] == 21
+            assert (
+                results[self.organization.id][f'{UsageType.LLM_REQUESTS}_limit'] == 21
+            )
             assert (
                 results[self.second_organization.id][f'{UsageType.SUBMISSION}_limit']
                 == 12
@@ -540,6 +532,7 @@ def test_get_org_effective_limits(self, include_onetime_addons):
                 results[self.organization.id][f'{UsageType.MT_CHARACTERS}_limit'] == 1
             )
             assert results[self.organization.id][f'{UsageType.ASR_SECONDS}_limit'] == 1
+            assert results[self.organization.id][f'{UsageType.LLM_REQUESTS}_limit'] == 1
             assert (
                 results[self.second_organization.id][f'{UsageType.SUBMISSION}_limit']
                 == 2
@@ -603,6 +596,7 @@ def setUp(self):
         product_metadata = {
             f'{UsageType.MT_CHARACTERS}_limit': '1',
             f'{UsageType.ASR_SECONDS}_limit': '1',
+            f'{UsageType.LLM_REQUESTS}_limit': '1',
             f'{UsageType.SUBMISSION}_limit': '1',
             f'{UsageType.STORAGE_BYTES}_limit': '1',
             'product_type': 'plan',
@@ -627,7 +621,7 @@ def test_check_exceeded_limit_adds_counters(self):
         ):
             self.add_submissions(count=2, asset=self.asset, username='someuser')
         self.add_nlp_trackers()
-        for usage_type, _ in SupportedUsageType.choices:
+        for usage_type, _ in UsageType.choices:
             check_exceeded_limit(self.someuser, usage_type)
             assert (
                 ExceededLimitCounter.objects.filter(
@@ -638,7 +632,7 @@ def test_check_exceeded_limit_adds_counters(self):
 
     def test_check_exceeded_limit_updates_counters(self):
         today = timezone.now()
-        for usage_type, _ in SupportedUsageType.choices:
+        for usage_type, _ in UsageType.choices:
             baker.make(
                 ExceededLimitCounter,
                 user=self.anotheruser,
@@ -655,7 +649,7 @@ def test_check_exceeded_limit_updates_counters(self):
         ):
             self.add_submissions(count=2, asset=self.asset, username='someuser')
         self.add_nlp_trackers()
-        for usage_type, _ in SupportedUsageType.choices:
+        for usage_type, _ in UsageType.choices:
             check_exceeded_limit(self.someuser, usage_type)
             counter = ExceededLimitCounter.objects.get(
                 user_id=self.anotheruser.id, limit_type=usage_type
diff --git a/kobo/apps/stripe/tests/utils.py b/kobo/apps/stripe/tests/utils.py
@@ -23,6 +23,7 @@ def generate_free_plan():
         'product_type': 'plan',
         f'{UsageType.SUBMISSION}_limit': '5000',
         f'{UsageType.ASR_SECONDS}_limit': '600',
+        f'{UsageType.LLM_REQUESTS}_limit': '20',
         f'{UsageType.MT_CHARACTERS}_limit': '6000',
         f'{UsageType.STORAGE_BYTES}_limit': '1000',
         'default_free_plan': 'true',
diff --git a/kobo/apps/stripe/utils/subscription_limits.py b/kobo/apps/stripe/utils/subscription_limits.py
@@ -6,7 +6,7 @@
 from django.db.models import F, Max, Q, QuerySet, Window
 from django.db.models.functions import Coalesce
 
-from kobo.apps.organizations.constants import SupportedUsageType, UsageType
+from kobo.apps.organizations.constants import UsageType
 from kobo.apps.organizations.models import Organization, OrganizationUser
 from kobo.apps.organizations.types import UsageLimits
 from kobo.apps.stripe.constants import ACTIVE_STRIPE_STATUSES
@@ -34,6 +34,7 @@ def get_default_add_on_limits():
         f'{UsageType.SUBMISSION}_limit': 0,
         f'{UsageType.ASR_SECONDS}_limit': 0,
         f'{UsageType.MT_CHARACTERS}_limit': 0,
+        f'{UsageType.LLM_REQUESTS}_limit': 0,
     }
 
 
@@ -114,7 +115,7 @@ def get_organizations_subscription_limits(
         if row['product_type'] == 'plan':
             row_limits = {
                 f'{usage_type}_limit': row[f'{usage_type}_limit']
-                for usage_type, _ in SupportedUsageType.choices
+                for usage_type, _ in UsageType.choices
             }
         elif row['product_type'] == 'addon':
             row_limits['addon_storage_limit'] = row[f'{UsageType.STORAGE_BYTES}_limit']
@@ -124,6 +125,7 @@ def get_organizations_subscription_limits(
     submission_limit = _get_limit_key(UsageType.SUBMISSION)
     characters_limit = _get_limit_key(UsageType.MT_CHARACTERS)
     seconds_limit = _get_limit_key(UsageType.ASR_SECONDS)
+    requests_limit = _get_limit_key(UsageType.LLM_REQUESTS)
     # Anyone who does not have a subscription is on the free tier plan by default
     default_plan = (
         Product.objects.filter(metadata__default_free_plan='true')
@@ -132,11 +134,12 @@ def get_organizations_subscription_limits(
             submission_limit=F(f'metadata__{submission_limit}'),
             mt_characters_limit=F(f'metadata__{characters_limit}'),
             asr_seconds_limit=F(f'metadata__{seconds_limit}'),
+            llm_requests_limit=F(f'metadata__{requests_limit}'),
         )
         .first()
     ) or {}
     default_plan_limits = {}
-    for usage_type, _ in SupportedUsageType.choices:
+    for usage_type, _ in UsageType.choices:
         limit_key = f'{usage_type}_limit'
         default_limit = default_plan.get(limit_key)
         if default_limit is None:
@@ -147,7 +150,7 @@ def get_organizations_subscription_limits(
     results = {}
     for org_id in all_org_ids:
         all_org_limits = {}
-        for usage_type, _ in SupportedUsageType.choices:
+        for usage_type, _ in UsageType.choices:
             plan_limit = subscription_limits_by_org_id.get(org_id, {}).get(
                 f'{usage_type}_limit'
             )
@@ -205,7 +208,7 @@ def get_organizations_effective_limits(
         PlanAddOn = apps.get_model('stripe', 'PlanAddOn')  # noqa
         addon_limits = PlanAddOn.get_organizations_totals(organizations=organizations)
         for org_id, limits in effective_limits.items():
-            for usage_type, _ in SupportedUsageType.choices:
+            for usage_type, _ in UsageType.choices:
                 addon = addon_limits.get(org_id, {}).get(f'total_{usage_type}_limit', 0)
                 limits[f'{usage_type}_limit'] += addon
     return effective_limits
@@ -235,6 +238,9 @@ def get_paid_subscription_limits(organization_ids: list[str], **kwargs) -> Query
     price_seconds_key, product_seconds_key = (
         _get_subscription_metadata_fields_for_usage_type(UsageType.ASR_SECONDS)
     )
+    price_requests_key, product_requests_key = (
+        _get_subscription_metadata_fields_for_usage_type(UsageType.LLM_REQUESTS)
+    )
 
     # Get organizations we care about (either those in the 'organizations' param or all)
     org_filter = Q(customer__subscriber_id__in=[org_id for org_id in organization_ids])
@@ -256,6 +262,7 @@ def get_paid_subscription_limits(organization_ids: list[str], **kwargs) -> Query
             mt_characters_limit=Coalesce(
                 F(price_characters_key), F(product_characters_key)
             ),
+            llm_requests_limit=Coalesce(F(price_requests_key), F(product_requests_key)),
             sub_start_date=F('start_date'),
             product_type=F('items__price__product__metadata__product_type'),
         )
diff --git a/kpi/serializers/v2/service_usage.py b/kpi/serializers/v2/service_usage.py
@@ -66,6 +66,7 @@ def _get_nlp_tracking_data(self, asset, start_date=None):
         if not asset.has_deployment:
             return {
                 'total_nlp_asr_seconds': 0,
+                'total_nlp_llm_requests': 0,
                 'total_nlp_mt_characters': 0,
             }
         return OpenRosaDeploymentBackend.nlp_tracking_data(
diff --git a/kpi/tests/test_usage_calculator.py b/kpi/tests/test_usage_calculator.py
diff --git a/kpi/utils/usage_calculator.py b/kpi/utils/usage_calculator.py

Original file line number	Diff line number	Diff line change
`@@ -66,6 +66,7 @@ def _get_nlp_tracking_data(self, asset, start_date=None):`
`66`	`66`	`if not asset.has_deployment:`
`67`	`67`	`return {`
`68`	`68`	`'total_nlp_asr_seconds': 0,`
	`69`	`+ 'total_nlp_llm_requests': 0,`
`69`	`70`	`'total_nlp_mt_characters': 0,`
`70`	`71`	`}`
`71`	`72`	`return OpenRosaDeploymentBackend.nlp_tracking_data(`