Dynavera/apps/accounts/models.py

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'],
                    llm_config={"model_id": "meta-llama-3.1-8b-instruct"}
                )

@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)