Dynavera/apps/accounts/models.py

from typing import ClassVar
from uuid import uuid4
from datetime import timedelta

from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin
from django.db import transaction
from django.db.models import BooleanField, CASCADE, CharField, DateField, DateTimeField, EmailField, ForeignKey, IntegerField, ManyToManyField, Model, TextField, UUIDField
from django.db.models.signals import post_save
from django.dispatch import receiver
from django.utils.translation import gettext_lazy as _
from django.utils import timezone

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 __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, unique = True)
    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')

    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 :(
        
        default_agents = [
            {
                'type': 'curriculum',
                'name': f"{instance.name} Curriculum Agent",
                'prompt': f"You are a curriculum specialist. Design a learning path for someone in the role of {instance.name}..."
            },
            {
                'type': 'knowledge',
                'name': f"{instance.name} Knowledge Agent",
                'prompt': f"You are a knowledge assistant. Search the provided docs and help with questions about {instance.name}..."
            },
            {
                'type': 'assessment',
                'name': f"{instance.name} Assessment Agent",
                'prompt': f"You are an evaluator. Create questions based on the knowledge of {instance.name}..."
            },
            {
                'type': 'monitor',
                'name': f"{instance.name} Progress Monitor",
                'prompt': f"You are a supervisor tracking the progress of someone in the role of {instance.name}..."
            }
        ]

        with transaction.atomic():
            for agent_data in default_agents:
                AgentConfig.objects.create(
                    organization=instance.organization,
                    name=agent_data['name'],
                    agent_type=agent_data['type'],
                    system_prompt=agent_data['prompt'],
                    llm_config={"model_id": "meta-llama-3.1-8b-instruct"}
                )