Dynavera/apps/orgs/models.py

from datetime import timedelta
from uuid import uuid4
from django.utils import timezone
from django.utils.translation import gettext_lazy as _
from django.db.models import BigAutoField, BooleanField, CASCADE, CharField, DateTimeField, ForeignKey, ManyToManyField, Model, TextField, UUIDField, IntegerField, FileField
from django.db.models.signals import post_delete, post_save
from django.db import transaction
from django.dispatch import receiver
from apps.users.mixins import TimeStampMixin
from apps.users.models import User

class Organization(TimeStampMixin, Model):

    id   = BigAutoField(primary_key = True)
    uuid = UUIDField(default = uuid4, unique = True, editable = False)
    name = CharField(max_length = 255, unique = True)

    description = TextField(blank = True, default = '')

    owner   = ForeignKey(User, on_delete = CASCADE, related_name = 'owned_organizations')
    members = ManyToManyField(User, through = 'OrganizationMembership', related_name = 'organizations')

    class Meta:
        verbose_name        = _('Organization')
        verbose_name_plural = _('Organizations')

    def __str__(self) -> str:
        return self.name

class OrganizationMembership(TimeStampMixin, Model):
    
    id           = BigAutoField(primary_key = True)
    user         = ForeignKey(User, on_delete = CASCADE, related_name = 'organization_memberships')
    organization = ForeignKey(Organization, on_delete = CASCADE, related_name = 'memberships')

    class Meta:
        verbose_name = _('Organization Membership')
        verbose_name_plural = _('Organization Memberships')
        unique_together = [['user', 'organization']]

    def __str__(self) -> str:
        return f'{self.user.full_name} - {self.organization.name}'

class OrganizationInvitation(TimeStampMixin, Model):
      
    id           = BigAutoField(primary_key = True)
    token        = UUIDField(default = uuid4, unique = True, editable = False)

    organization = ForeignKey(Organization, on_delete = CASCADE, related_name = "invite_tokens")
    created_by   = ForeignKey(User, on_delete = CASCADE, related_name = "created_invites")

    expires_at   = DateTimeField()

    uses         = IntegerField(default = 0)
    max_uses     = IntegerField(default = 1)

    is_active    = BooleanField(default = True)

    class Meta:
        verbose_name = _("Invite Token")
        verbose_name_plural = _("Invite Tokens")

    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(TimeStampMixin, Model):

    id   = BigAutoField(primary_key = True)
    name = CharField(max_length = 100, unique = True)
    uuid = UUIDField(default = uuid4, editable = False, unique = True)
    description  = TextField(blank = True, default = '')
    organization = ForeignKey(Organization, on_delete = CASCADE, related_name = "roles")
    members = ManyToManyField(User, through = "RoleMembership", related_name = "roles")

    class Meta:
        verbose_name        = _('Role')
        verbose_name_plural = _('Roles')

    def __str__(self) -> str:
        return self.name
    
class RoleMembership(TimeStampMixin, Model):

    id   = BigAutoField(primary_key = True)
    user = ForeignKey(User, on_delete = CASCADE, related_name = "role_memberships")
    role = ForeignKey(Role, on_delete = CASCADE, related_name = "memberships")

    class Meta:
        verbose_name = _("Role Membership")
        verbose_name_plural = _("Role Memberships")
        unique_together = [["user", "role"]]

    def __str__(self) -> str:
        return f"{self.user.full_name} - {self.role.name}"

class TrainingFile(TimeStampMixin, Model):
    
    ALLOWED_EXTENSIONS = ('txt', 'pdf', 'md', 'csv', 'json', 'docx', 'doc')
    
    STATUS_CHOICES = [
        ('ingesting', 'Ingesting'),
        ('chunked', 'Chunked'),
        ('embedded', 'Embedded'),
        ('failed', 'Failed'),
    ]

    id           = BigAutoField(primary_key = True)
    uuid         = UUIDField(default = uuid4, unique = True, editable = False)
    role         = ForeignKey(Role, on_delete = CASCADE, related_name = "training_files")
    uploaded_by  = ForeignKey(User, on_delete = CASCADE, related_name = "uploaded_training_files")
    
    file         = FileField(upload_to = 'training_files/%Y/%m/%d/')
    file_name    = CharField(max_length = 255)
    file_size    = IntegerField()
    file_type    = CharField(max_length = 50)
    
    description  = TextField(blank = True, default = '')
    status       = CharField(max_length = 20, choices = STATUS_CHOICES, default = 'ingesting')
    is_processed = BooleanField(default = False)
    
    class Meta:
        verbose_name = _("Training File")
        verbose_name_plural = _("Training Files")
        ordering = ['-created_at']
    
    def __str__(self) -> str:
        return f"{self.file_name} - {self.role.name}"


@receiver(post_delete, sender=TrainingFile)
def delete_training_file_on_delete(sender, instance, **kwargs):
    if instance.file:
        try:
            import os
            if os.path.isfile(instance.file.path):
                os.remove(instance.file.path)
        except Exception:
            pass


@receiver(post_save, sender=TrainingFile)
def enqueue_training_file_ingestion(sender, instance, created, **kwargs):
    if not created:
        return

    def _enqueue():
        from apps.mlstore.tasks import ingest_training_file_task
        ingest_training_file_task.delay(str(instance.uuid))

    transaction.on_commit(_enqueue)