From 00b4e4188c957f63f3ad688810dfc99ebdb5e975 Mon Sep 17 00:00:00 2001 From: Mitchell Kotler Date: Tue, 23 Dec 2025 13:17:03 -0500 Subject: [PATCH] Add resource sharing for organization collectives --- .../0018_alter_organization_merged.py | 20 ++ ...on_members_organization_parent_and_more.py | 30 +++ documentcloud/organizations/models.py | 71 ++++++- .../organizations/tests/test_models.py | 197 +++++++++++++++++- .../migrations/0013_alter_project_options.py | 17 ++ requirements/base.txt | 2 +- requirements/local.txt | 2 +- requirements/production.txt | 2 +- 8 files changed, 333 insertions(+), 8 deletions(-) create mode 100644 documentcloud/organizations/migrations/0018_alter_organization_merged.py create mode 100644 documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py create mode 100644 documentcloud/projects/migrations/0013_alter_project_options.py diff --git a/documentcloud/organizations/migrations/0018_alter_organization_merged.py b/documentcloud/organizations/migrations/0018_alter_organization_merged.py new file mode 100644 index 00000000..62733267 --- /dev/null +++ b/documentcloud/organizations/migrations/0018_alter_organization_merged.py @@ -0,0 +1,20 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:06 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('organizations', '0017_organization_merged'), + ] + + operations = [ + migrations.AlterField( + model_name='organization', + name='merged', + field=models.ForeignKey(blank=True, help_text='The organization this organization was merged in to', null=True, on_delete=django.db.models.deletion.PROTECT, related_name='+', to=settings.SQUARELET_ORGANIZATION_MODEL), + ), + ] diff --git a/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py b/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py new file mode 100644 index 00000000..2cd7cb47 --- /dev/null +++ b/documentcloud/organizations/migrations/0019_organization_members_organization_parent_and_more.py @@ -0,0 +1,30 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:32 + +from django.conf import settings +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('organizations', '0018_alter_organization_merged'), + ] + + operations = [ + migrations.AddField( + model_name='organization', + name='members', + field=models.ManyToManyField(blank=True, help_text='Organizations which are members of this organization (useful for trade associations or other member groups)', related_name='groups', to=settings.SQUARELET_ORGANIZATION_MODEL), + ), + migrations.AddField( + model_name='organization', + name='parent', + field=models.ForeignKey(blank=True, help_text='The parent organization', null=True, on_delete=django.db.models.deletion.PROTECT, related_name='children', to=settings.SQUARELET_ORGANIZATION_MODEL, verbose_name='parent'), + ), + migrations.AddField( + model_name='organization', + name='share_resources', + field=models.BooleanField(default=True, help_text='Share resources (subscriptions, credits) with all children and member organizations. Global toggle that applies to all relationships.', verbose_name='share resources'), + ), + ] diff --git a/documentcloud/organizations/models.py b/documentcloud/organizations/models.py index 483de406..89543219 100644 --- a/documentcloud/organizations/models.py +++ b/documentcloud/organizations/models.py @@ -130,6 +130,19 @@ def merge(self, uuid): self.addons.update(organization=other) self.visual_addons.update(organization=other) + # transfer children to the other organization + self.children.update(parent=other) + + # transfer group memberships + groups = self.groups.all() + other.groups.add(*groups) + self.groups.clear() + + # transfer members + members = self.members.all() + other.members.add(*members) + self.members.clear() + self.merged = other def calc_ai_credits_per_month(self, users): @@ -147,20 +160,54 @@ def use_ai_credits(self, amount, user_id, note): initial_amount = amount ai_credit_count = {"monthly": 0, "regular": 0} organization = Organization.objects.select_for_update().get(pk=self.pk) + if organization.parent and organization.parent.share_resources: + parent = Organization.objects.select_for_update().get( + pk=organization.parent_id + ) + else: + parent = None + groups = organization.groups.filter(share_resources=True).select_for_update() + # Deduct from own resources first ai_credit_count["monthly"] = min(amount, organization.monthly_ai_credits) amount -= ai_credit_count["monthly"] ai_credit_count["regular"] = min(amount, organization.number_ai_credits) amount -= ai_credit_count["regular"] - if amount > 0: - raise InsufficientAICreditsError(amount) - organization.monthly_ai_credits -= ai_credit_count["monthly"] organization.number_ai_credits -= ai_credit_count["regular"] organization.save() + # Then deduct from parent resources + if parent: + parent_monthly = min(amount, parent.monthly_ai_credits) + ai_credit_count["monthly"] += parent_monthly + amount -= parent_monthly + parent.monthly_ai_credits -= parent_monthly + + parent_regular = min(amount, parent.number_ai_credits) + ai_credit_count["regular"] += parent_regular + amount -= parent_regular + parent.number_ai_credits -= parent_regular + parent.save() + + # Then deduct from group resources + for group in groups: + group_monthly = min(amount, group.monthly_ai_credits) + ai_credit_count["monthly"] += group_monthly + amount -= group_monthly + group.monthly_ai_credits -= group_monthly + + group_regular = min(amount, group.number_ai_credits) + ai_credit_count["regular"] += group_regular + amount -= group_regular + group.number_ai_credits -= group_regular + group.save() + + if amount > 0: + raise InsufficientAICreditsError(amount) + organization.ai_credit_logs.create( user_id=user_id, organization=organization, @@ -170,6 +217,24 @@ def use_ai_credits(self, amount, user_id, note): return ai_credit_count + def get_total_number_ai_credits(self): + """Get total number AI credits including parent and groups""" + number_ai_credits = self.number_ai_credits + if self.parent and self.parent.share_resources: + number_ai_credits += self.parent.number_ai_credits + for group in self.groups.filter(share_resources=True): + number_ai_credits += group.number_ai_credits + return number_ai_credits + + def get_total_monthly_ai_credits(self): + """Get total monthly AI credits 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 + class AICreditLog(models.Model): """Log usage of AI Credits""" diff --git a/documentcloud/organizations/tests/test_models.py b/documentcloud/organizations/tests/test_models.py index 30407323..d61a4090 100644 --- a/documentcloud/organizations/tests/test_models.py +++ b/documentcloud/organizations/tests/test_models.py @@ -2,6 +2,7 @@ import pytest # DocumentCloud +from documentcloud.organizations.exceptions import InsufficientAICreditsError from documentcloud.organizations.models import Organization from documentcloud.organizations.tests.factories import OrganizationFactory from documentcloud.users.models import User @@ -50,7 +51,7 @@ def test_merge_fks(self): if f.is_relation and f.auto_created ] ) - == 6 + == 8 ) # Many to many relations defined on the Organization model assert ( @@ -61,5 +62,197 @@ def test_merge_fks(self): if f.many_to_many and not f.auto_created ] ) - == 1 + == 2 ) + + +class TestOrganizationCollective: + """Tests for Organization collective resource sharing""" + + @pytest.mark.django_db() + def test_use_ai_credits_with_parent(self): + """Test using AI credits with parent's resources when own resources exhausted""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, parent=parent_org + ) + + # Use 15 credits - 10 from child monthly, 5 from child regular + result = child_org.use_ai_credits(15, user.pk, "Test") + + child_org.refresh_from_db() + parent_org.refresh_from_db() + + assert result == {"monthly": 10, "regular": 5} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert parent_org.monthly_ai_credits == 50 + assert parent_org.number_ai_credits == 25 + + @pytest.mark.django_db() + def test_use_ai_credits_parent_no_sharing(self): + """Test that resources are not shared when parent.share_resources=False""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=False + ) + child_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, parent=parent_org + ) + + # Try to use 15 credits - should fail after child's 15 credits + with pytest.raises(InsufficientAICreditsError): + child_org.use_ai_credits(20, user.pk, "Test") + + @pytest.mark.django_db() + def test_use_ai_credits_with_groups(self): + """Test using AI credits with group's resources""" + user = UserFactory() + group_org = OrganizationFactory( + monthly_ai_credits=50, number_ai_credits=25, share_resources=True + ) + child_org = OrganizationFactory(monthly_ai_credits=10, number_ai_credits=5) + child_org.groups.add(group_org) + + # Use 20 credits - should use 10 from child monthly, 5 from child regular, + # 5 from group monthly + result = child_org.use_ai_credits(20, user.pk, "Test") + + child_org.refresh_from_db() + group_org.refresh_from_db() + + assert result == {"monthly": 15, "regular": 5} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert group_org.monthly_ai_credits == 45 + + @pytest.mark.django_db() + def test_use_ai_credits_with_multiple_groups(self): + """Test using AI credits from multiple groups""" + user = UserFactory() + group1 = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + group2 = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + child_org = OrganizationFactory(monthly_ai_credits=10, number_ai_credits=0) + child_org.groups.add(group1, group2) + + # Use 40 credits - should use 5 from child, then from groups + result = child_org.use_ai_credits(40, user.pk, "Test") + + child_org.refresh_from_db() + group1.refresh_from_db() + group2.refresh_from_db() + + assert result == {"monthly": 30, "regular": 10} + assert child_org.monthly_ai_credits == 0 + # Groups are consumed in arbitrary order + assert group1.monthly_ai_credits + group2.monthly_ai_credits == 20 + assert group1.number_ai_credits + group2.number_ai_credits == 10 + + @pytest.mark.django_db() + def test_use_ai_credits_parent_and_groups(self): + """Test using AI credits with both parent and groups""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=20, number_ai_credits=10, share_resources=True + ) + group_org = OrganizationFactory( + monthly_ai_credits=30, number_ai_credits=15, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=5, number_ai_credits=0, parent=parent_org + ) + child_org.groups.add(group_org) + + # Use 60 credits: 5 child monthly, 20 parent monthly, 10 parent regular, + # 25 group monthly + result = child_org.use_ai_credits(60, user.pk, "Test") + + child_org.refresh_from_db() + parent_org.refresh_from_db() + group_org.refresh_from_db() + + assert result == {"monthly": 50, "regular": 10} + assert child_org.monthly_ai_credits == 0 + assert child_org.number_ai_credits == 0 + assert parent_org.monthly_ai_credits == 0 + assert parent_org.number_ai_credits == 0 + assert group_org.monthly_ai_credits == 5 + assert group_org.number_ai_credits == 15 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_own_only(self): + """Test get_total_number_ai_credits with no parent or groups""" + org = OrganizationFactory(number_ai_credits=100) + assert org.get_total_number_ai_credits() == 100 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_with_parent(self): + """Test get_total_number_ai_credits including parent""" + parent_org = OrganizationFactory(number_ai_credits=50, share_resources=True) + child_org = OrganizationFactory(number_ai_credits=25, parent=parent_org) + + assert child_org.get_total_number_ai_credits() == 75 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_parent_no_sharing(self): + """Test get_total_number_ai_credits when parent doesn't share""" + parent_org = OrganizationFactory(number_ai_credits=50, share_resources=False) + child_org = OrganizationFactory(number_ai_credits=25, parent=parent_org) + + assert child_org.get_total_number_ai_credits() == 25 + + @pytest.mark.django_db() + def test_get_total_number_ai_credits_with_groups(self): + """Test get_total_number_ai_credits including groups""" + group1 = OrganizationFactory(number_ai_credits=30, share_resources=True) + group2 = OrganizationFactory(number_ai_credits=20, share_resources=True) + org = OrganizationFactory(number_ai_credits=10) + org.groups.add(group1, group2) + + assert org.get_total_number_ai_credits() == 60 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_own_only(self): + """Test get_total_monthly_ai_credits with no parent or groups""" + org = OrganizationFactory(monthly_ai_credits=50) + assert org.get_total_monthly_ai_credits() == 50 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_with_parent(self): + """Test get_total_monthly_ai_credits including parent""" + parent_org = OrganizationFactory(monthly_ai_credits=100, share_resources=True) + child_org = OrganizationFactory(monthly_ai_credits=25, parent=parent_org) + + assert child_org.get_total_monthly_ai_credits() == 125 + + @pytest.mark.django_db() + def test_get_total_monthly_ai_credits_with_groups(self): + """Test get_total_monthly_ai_credits including groups""" + group1 = OrganizationFactory(monthly_ai_credits=40, share_resources=True) + group2 = OrganizationFactory(monthly_ai_credits=30, share_resources=True) + org = OrganizationFactory(monthly_ai_credits=15) + org.groups.add(group1, group2) + + assert org.get_total_monthly_ai_credits() == 85 + + @pytest.mark.django_db() + def test_insufficient_ai_credits_with_parent(self): + """Test InsufficientAICreditsError even with parent resources""" + user = UserFactory() + parent_org = OrganizationFactory( + monthly_ai_credits=10, number_ai_credits=5, share_resources=True + ) + child_org = OrganizationFactory( + monthly_ai_credits=5, number_ai_credits=2, parent=parent_org + ) + + # Try to use more credits than available (total is 22, trying to use 25) + with pytest.raises(InsufficientAICreditsError): + child_org.use_ai_credits(25, user.pk, "Test") diff --git a/documentcloud/projects/migrations/0013_alter_project_options.py b/documentcloud/projects/migrations/0013_alter_project_options.py new file mode 100644 index 00000000..283ddcd7 --- /dev/null +++ b/documentcloud/projects/migrations/0013_alter_project_options.py @@ -0,0 +1,17 @@ +# Generated by Django 4.2.2 on 2025-12-22 21:06 + +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('projects', '0012_auto_20210407_1801'), + ] + + operations = [ + migrations.AlterModelOptions( + name='project', + options={'ordering': ('slug',), 'permissions': (('add_remove_project', 'Can add & remove documents from a project'), ('change_project_all', 'Can edit all fields on a project (not just pinned)'))}, + ), + ] diff --git a/requirements/base.txt b/requirements/base.txt index e6cf0655..202e8920 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -429,7 +429,7 @@ sqlparse==0.4.4 # via # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.11 # via -r requirements/base.in stack-data==0.3.0 # via ipython diff --git a/requirements/local.txt b/requirements/local.txt index c0b7be23..41434c43 100644 --- a/requirements/local.txt +++ b/requirements/local.txt @@ -736,7 +736,7 @@ sqlparse==0.4.4 # -r requirements/./base.txt # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.11 # via -r requirements/./base.txt stack-data==0.3.0 # via diff --git a/requirements/production.txt b/requirements/production.txt index 513ec384..fdabbbb8 100644 --- a/requirements/production.txt +++ b/requirements/production.txt @@ -608,7 +608,7 @@ sqlparse==0.4.4 # -r requirements/./base.txt # django # django-debug-toolbar -squarelet-auth==0.1.10 +squarelet-auth==0.1.11 # via -r requirements/./base.txt stack-data==0.3.0 # via