from datetime import timedelta from typing import ClassVar from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin from django.core.exceptions import ValidationError from django.db import transaction from django.db.models import CASCADE, BooleanField, CharField, DateField, DateTimeField, EmailField, ForeignKey, IntegerField, ManyToManyField, Model, TextField from django.db.models.signals import m2m_changed, post_delete, post_save from django.dispatch import receiver from django.utils import timezone from django.utils.translation import gettext_lazy as _ from apps.accounts.managers import UserManager from apps.accounts.mixins import IdentifierMixin, TimeStampMixin class User(AbstractBaseUser, IdentifierMixin, TimeStampMixin, PermissionsMixin): email_address = EmailField(verbose_name = _("Email Address"), max_length = 255, unique = True) first_name = CharField(verbose_name = _("First Name"), max_length = 255) last_name = CharField(verbose_name = _("Last Name"), max_length = 255) date_of_birth = DateField(verbose_name = _("Date of Birth"), null = True, blank = True) is_active = BooleanField(verbose_name = _("Account Active"), default = True) is_staff = BooleanField(verbose_name = _("Account Admin"), default = False) is_manager = BooleanField(verbose_name = _("Organization Manager"), default = False) USERNAME_FIELD = 'email_address' EMAIL_FIELD = 'email_address' REQUIRED_FIELDS = ['first_name', 'last_name', 'date_of_birth'] objects: ClassVar[UserManager] = UserManager() def has_perm(self, perm, obj=None): return True def has_module_perms(self, app_label): return True class Meta: verbose_name = _('User') verbose_name_plural = _('Users') @property def full_name(self) -> str: return f"{self.first_name} {self.last_name}" def is_owner_of(self, organization: 'Organization') -> bool: return organization.owner.id == self.id def is_member_of(self, organization: 'Organization') -> bool: return organization.members.filter(id=self.id).exists() def __str__(self) -> str: return self.full_name class Organization(IdentifierMixin, TimeStampMixin, Model): name = CharField(verbose_name = _("Name"), max_length = 255, unique = True) description = TextField(verbose_name = _("Description"), blank = True, default = '') owner = ForeignKey(User, on_delete = CASCADE, related_name = 'owned_organizations') members = ManyToManyField(User, related_name = 'organizations') class Meta: verbose_name = _('Organization') verbose_name_plural = _('Organizations') def __str__(self) -> str: return self.name class Invite(IdentifierMixin, TimeStampMixin, Model): organization = ForeignKey(Organization, on_delete = CASCADE, related_name = "invites") created_by = ForeignKey(User, on_delete = CASCADE, related_name = "created_invites") expires_at = DateTimeField(verbose_name=_("Expires At")) uses = IntegerField(verbose_name=_("Uses"), default = 0) max_uses = IntegerField(verbose_name=_("Max Uses"), default = 1) is_active = BooleanField(verbose_name=_("Is Active"), default = True) class Meta: verbose_name = _("Invite") verbose_name_plural = _("Invites") def save(self, *args, **kwargs): if not self.expires_at: self.expires_at = timezone.now() + timedelta(days=7) super().save(*args, **kwargs) def is_valid(self): return self.is_active and self.uses < self.max_uses and timezone.now() < self.expires_at def __str__(self) -> str: return f"Invite for {self.organization.name} by {self.created_by.full_name} (expires {self.expires_at})" class Role(IdentifierMixin, TimeStampMixin, Model): name = CharField(verbose_name = _("Name"), max_length = 100) description = TextField(verbose_name = _("Description"), blank = True, default = '') organization = ForeignKey(Organization, on_delete = CASCADE, related_name = "roles") members = ManyToManyField(User, related_name = "roles") class Meta: verbose_name = _('Role') verbose_name_plural = _('Roles') unique_together = [('organization', 'name')] def __str__(self) -> str: return f"{self.name} ({self.organization.name})" @receiver(post_save, sender=Role) def create_default_agents_for_role(sender, instance: Role, created: bool, **kwargs): if created: from apps.onboarding.models import AgentConfig # L: circular import :( from apps.onboarding.consumers.prompts import OnboardingPrompts default_agents = [ { 'type': 'curriculum', 'name': f"{instance.name} Curriculum Agent", 'prompt': OnboardingPrompts.default_curriculum_prompt(instance.name), }, { 'type': 'knowledge', 'name': f"{instance.name} Knowledge Agent", 'prompt': OnboardingPrompts.default_knowledge_prompt(instance.name), }, { 'type': 'assessment', 'name': f"{instance.name} Assessment Agent", 'prompt': OnboardingPrompts.default_assessment_prompt(instance.name), }, { 'type': 'monitor', 'name': f"{instance.name} Progress Monitor", 'prompt': OnboardingPrompts.default_monitor_prompt(instance.name), }, ] with transaction.atomic(): for agent_data in default_agents: AgentConfig.objects.create( organization=instance.organization, role=instance, name=agent_data['name'], agent_type=agent_data['type'], system_prompt=agent_data['prompt'] ) @receiver(post_delete, sender=Role) def delete_role_agents_on_role_delete(sender, instance: Role, **kwargs): from apps.onboarding.models import AgentConfig AgentConfig.objects.filter(role=instance).delete() @receiver(post_save, sender=Organization) def ensure_owner_is_member(sender, instance: Organization, **kwargs): if instance.owner.id and not instance.members.filter(id=instance.owner.id).exists(): instance.members.add(instance.owner) @receiver(m2m_changed, sender=Organization.members.through) def prevent_owner_from_being_removed(sender, instance: Organization, action: str, pk_set, **kwargs): if action == 'pre_remove' and instance.owner.id in (pk_set or set()): raise ValidationError(_('Organization owner must remain a member.')) if action in ['post_add', 'post_remove', 'post_clear']: if instance.owner.id and not instance.members.filter(id=instance.owner.id).exists(): instance.members.add(instance.owner)