chat_backend/llm_be/chat_backend/models.py

from django.db import models
from django.contrib.auth.models import AbstractUser
from django.utils import timezone
from autoslug import AutoSlugField
from chat_backend.storage import DatabaseStorage
import uuid

# Create your models here.

DB_FILE_STORAGE = DatabaseStorage()


class TimeInfoBase(models.Model):

    created = models.DateTimeField(default=timezone.now)
    last_modified = models.DateTimeField(default=timezone.now)

    class Meta:
        abstract = True

    def save(self, *args, **kwargs):
        if not kwargs.pop("skip_last_modified", False) and not hasattr(
            self, "skip_last_modified"
        ):
            self.last_modified = timezone.now()
            if kwargs.get("update_fields") is not None:
                kwargs["update_fields"] = list(
                    {*kwargs["update_fields"], "last_modified"}
                )

        super().save(*args, **kwargs)


class LLMModels(TimeInfoBase):
    name = models.CharField(max_length=254, default="")
    port = models.IntegerField(help_text="This specifies the port that the LLM runs on")
    description = models.CharField(
        max_length=512,
        default="",
        help_text="A description for the LLM. Limit is 512 characters",
    )


class Company(TimeInfoBase):
    name = models.TextField(max_length=256)
    state = models.TextField(max_length=2)
    zipcode = models.TextField(max_length=5)
    address = models.TextField(max_length=256)
    available_llms = models.ForeignKey(
        "LLMModels",
        on_delete=models.CASCADE,
        blank=True,
        null=True,
        help_text="A list of LLMs that company can use",
    )

    def __str__(self):
        return self.name


class CustomUser(AbstractUser):
    company = models.ForeignKey(
        Company, on_delete=models.CASCADE, blank=True, null=True
    )
    is_company_manager = models.BooleanField(
        help_text="Allows the edit/add/remove of users for a company", default=False
    )
    deleted = models.BooleanField(help_text="This is to hid accounts", default=False)
    has_signed_tos = models.BooleanField(
        default=False, help_text="If the user has signed the TOS"
    )
    slug = AutoSlugField(populate_from="email")
    conversation_order = models.BooleanField(
        default=True, help_text="How the conversations should display"
    )

    def get_set_password_url(self):
        from django.conf import settings

        base = settings.FRONTEND_BASE_URL.rstrip("/")
        return f"{base}/set_password/?slug={self.slug}"


class UserAuthEvent(models.Model):
    """Audit trail for password reset / set actions, shown on user admin."""

    class EventType(models.TextChoices):
        PASSWORD_RESET_REQUESTED = (
            "password_reset_requested",
            "Password reset requested",
        )
        PASSWORD_SET = ("password_set", "Password set")
        INVITE_SENT = ("invite_sent", "Invite sent")

    user = models.ForeignKey(
        CustomUser,
        on_delete=models.CASCADE,
        related_name="auth_events",
    )
    event_type = models.CharField(max_length=64, choices=EventType.choices)
    created = models.DateTimeField(default=timezone.now, db_index=True)
    detail = models.CharField(max_length=512, blank=True, default="")
    ip_address = models.GenericIPAddressField(null=True, blank=True)

    class Meta:
        ordering = ["-created"]

    def __str__(self):
        return f"{self.get_event_type_display()} @ {self.created.isoformat()}"

    @classmethod
    def log(cls, user, event_type, *, detail="", ip_address=None):
        return cls.objects.create(
            user=user,
            event_type=event_type,
            detail=detail or "",
            ip_address=ip_address,
        )


class OutboundEmail(models.Model):
    """Record of emails queued/sent by the app (visible in admin)."""

    class Kind(models.TextChoices):
        PASSWORD_RESET = "password_reset", "Password reset"
        INVITE = "invite", "Invite"
        FEEDBACK = "feedback", "Feedback"

    class Status(models.TextChoices):
        QUEUED = "queued", "Queued"
        SENDING = "sending", "Sending"
        SENT = "sent", "Sent to SMTP"
        FAILED = "failed", "Failed"

    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    kind = models.CharField(max_length=32, choices=Kind.choices)
    status = models.CharField(
        max_length=16, choices=Status.choices, default=Status.QUEUED, db_index=True
    )
    to_email = models.EmailField()
    from_email = models.EmailField()
    subject = models.CharField(max_length=255)
    html_template = models.CharField(max_length=255)
    text_template = models.CharField(max_length=255)
    context = models.JSONField(default=dict, blank=True)
    error_message = models.TextField(blank=True, default="")
    user = models.ForeignKey(
        CustomUser,
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
        related_name="outbound_emails",
    )
    created = models.DateTimeField(default=timezone.now, db_index=True)
    updated = models.DateTimeField(auto_now=True)
    sent_at = models.DateTimeField(null=True, blank=True)

    class Meta:
        ordering = ["-created"]

    def __str__(self):
        return f"{self.subject} → {self.to_email} ({self.status})"


FEEDBACK_CHOICE = (
    ("SUBMITTED", "Submitted"),
    ("RESOLVED", "Resolved"),
    ("DEFFERED", "Deffered"),
    ("CLOSED", "Closed"),
)

FEEDBACK_CATEGORIES = (
    ("NOT_DEFINED", "Not defined"),
    ("BUG", "Bug"),
    ("ENHANCEMENT", "Enhancement"),
    ("OTHER", "Other"),
    ("MAX_CATEGORIES", "Max Categories"),
)


class Feedback(TimeInfoBase):
    title = models.TextField(max_length=64, default="")
    user = models.ForeignKey(
        CustomUser, on_delete=models.CASCADE, blank=True, null=True
    )
    text = models.TextField(max_length=512)
    status = models.CharField(
        max_length=24, choices=FEEDBACK_CHOICE, default="SUBMITTED"
    )
    category = models.CharField(
        max_length=24, choices=FEEDBACK_CATEGORIES, default="NOT_DEFINED"
    )

    def get_user_email(self):
        if self.user:
            return self.user.email
        else:
            return ""


MONTH_CHOICES = (
    ("JANUARY", "January"),
    ("FEBRUARY", "February"),
    ("MARCH", "March"),
    # ....
    ("DECEMBER", "December"),
)

month = models.CharField(max_length=9, choices=MONTH_CHOICES, default="JANUARY")


class Announcement(TimeInfoBase):
    class Status(models.TextChoices):
        default = "DEFAULT", "default"
        warning = "WARNING", "warning"
        info = "INFO", "info"
        danger = "DANGER", "danger"

    status = models.CharField(
        max_length=7, choices=Status.choices, default=Status.default
    )
    message = models.TextField(max_length=256)
    start_date_time = models.DateTimeField(auto_now=True)
    end_date_time = models.DateTimeField(auto_now=True)


class Conversation(TimeInfoBase):
    user = models.ForeignKey(
        CustomUser, on_delete=models.CASCADE, blank=True, null=True
    )
    title = models.CharField(
        max_length=64, help_text="The title for the conversation", default=""
    )
    deleted = models.BooleanField(
        help_text="This is to hide conversations", default=False
    )

    def get_user_email(self):
        if self.user:
            return self.user.email
        else:
            return ""

    def __str__(self):
        return self.title


class Prompt(TimeInfoBase):
    message = models.CharField(max_length=100 * 1024, help_text="The text for a prompt")
    user_created = models.BooleanField(
        help_text="True if was created by the user. False if it was generate by the LLM"
    )
    conversation = models.ForeignKey(
        "Conversation", on_delete=models.CASCADE, blank=True, null=True
    )
    file = models.FileField(
        upload_to="prompt_files/",
        storage=DB_FILE_STORAGE,
        blank=True,
        null=True,
        help_text="file for the prompt (stored in database)",
    )
    file_type = models.CharField(
        max_length=16,
        blank=True,
        null=True,
        help_text="file type of the file for the prompt",
    )

    def get_conversation_title(self):
        if self.conversation:
            return self.conversation.title
        else:
            return ""

    def file_exists(self):
        return self.file != None and self.file.storage.exists(self.file.name)


class PromptMetric(TimeInfoBase):
    PROMPT_METRIC_CHOICES = (
        ("CREATED", "Created"),
        ("SUBMITTED", "Submitted"),
        ("PROCESSED", "Processed"),
        ("FINISHED", "Finished"),
        ("MAX_PROMPT_METRIC_CHOICES", "Max Prompt Metric Choices"),
    )
    prompt_id = models.IntegerField(help_text="The id of the prompt this matches to")
    conversation_id = models.IntegerField(
        help_text="The id of the conversation this matches to"
    )
    event = models.CharField(
        max_length=26, choices=PROMPT_METRIC_CHOICES, default="CREATED"
    )
    model_name = models.CharField(max_length=215, help_text="The name of the model")
    start_time = models.DateTimeField()
    end_time = models.DateTimeField(blank=True, null=True)
    prompt_length = models.IntegerField(
        help_text="How many characters are in the prompt"
    )
    reponse_length = models.IntegerField(
        blank=True, null=True, help_text="How many characters are in the response"
    )
    tokens_in = models.IntegerField(
        blank=True,
        null=True,
        help_text=(
            "Prompt/input tokens reported by the LLM provider usage payload. "
            "Null when the provider did not report usage (never estimated)."
        ),
    )
    tokens_out = models.IntegerField(
        blank=True,
        null=True,
        help_text=(
            "Completion/output tokens reported by the LLM provider usage payload. "
            "Null when the provider did not report usage (never estimated)."
        ),
    )
    has_file = models.BooleanField(help_text="Is there a file")
    file_type = models.CharField(
        max_length=16, help_text="The file type, if any", blank=True, null=True
    )

    def get_duration(self):
        if self.start_time and self.end_time:
            difference = self.end_time - self.start_time
            return difference.seconds
        return 0


# Document Models
class DocumentWorkspace(TimeInfoBase):
    name = models.CharField(max_length=255)
    company = models.ForeignKey(Company, on_delete=models.CASCADE)


class Document(TimeInfoBase):
    workspace = models.ForeignKey(DocumentWorkspace, on_delete=models.CASCADE)
    file = models.FileField(
        upload_to="documents/",
        storage=DB_FILE_STORAGE,
        help_text="uploaded document bytes (stored in database)",
    )
    uploaded_at = models.DateTimeField(auto_now_add=True)
    processed = models.BooleanField(default=False)
    active = models.BooleanField(default=False)


class StoredFile(TimeInfoBase):
    """Blob store for DatabaseStorage — prompt attachments and documents."""

    name = models.CharField(max_length=512, unique=True, db_index=True)
    content = models.BinaryField()
    size = models.PositiveBigIntegerField(default=0)
    content_type = models.CharField(max_length=255, blank=True, default="")

    def __str__(self):
        return self.name