234 lines
7.2 KiB
Python
234 lines
7.2 KiB
Python
from django.db import models
|
|
from django.contrib.auth.models import AbstractUser
|
|
from django.utils import timezone
|
|
from autoslug import AutoSlugField
|
|
from django.core.files.storage import FileSystemStorage
|
|
|
|
# Create your models here.
|
|
|
|
FILE_STORAGE = FileSystemStorage(location="prompt_files")
|
|
|
|
|
|
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",
|
|
)
|
|
|
|
|
|
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):
|
|
return f"https://www.chat.aimloperations.com/set_password?slug={self.slug}"
|
|
|
|
|
|
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=10 * 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=FILE_STORAGE, blank=True, null=True, help_text="file for the prompt"
|
|
)
|
|
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"
|
|
)
|
|
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/')
|
|
uploaded_at = models.DateTimeField(auto_now_add=True)
|
|
processed = models.BooleanField(default=False)
|
|
active = models.BooleanField(default=False)
|