chat_backend/llm_be/chat_backend/serializers.py

from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
from rest_framework import serializers
from .models import (
    CustomUser,
    Announcement,
    Company,
    Conversation,
    Prompt,
    Feedback,
    FEEDBACK_CATEGORIES,
    DocumentWorkspace,
    Document,
)


class MyTokenObtainPairSerializer(TokenObtainPairSerializer):
    @classmethod
    def get_token(cls, user):
        token = super(MyTokenObtainPairSerializer, cls).get_token(user)

        # add custom claim
        token["company"] = "something here"

        return token


class CompanySerializer(serializers.ModelSerializer):
    class Meta:
        model = Company
        fields = "__all__"


class AnnouncmentSerializer(serializers.ModelSerializer):
    class Meta:
        model = Announcement
        fields = "__all__"


class FeedbackSerializer(serializers.ModelSerializer):
    class Meta:
        model = Feedback
        fields = "__all__"


class CustomUserSerializer(serializers.ModelSerializer):
    email = serializers.EmailField(required=True)
    username = serializers.CharField()
    password = serializers.CharField(min_length=8, write_only=True)
    company = CompanySerializer()
    has_usable_password = serializers.BooleanField()

    class Meta:
        model = CustomUser
        fields = "__all__"
        extra_kwargs = {"password": {"write_only": True}}


class SelfServeRegistrationSerializer(serializers.Serializer):
    """Minimal payload for public self-serve sign-up (gated by settings)."""

    email = serializers.EmailField(required=True)
    password = serializers.CharField(min_length=8, write_only=True)
    first_name = serializers.CharField(
        required=False, allow_blank=True, max_length=150, default=""
    )
    last_name = serializers.CharField(
        required=False, allow_blank=True, max_length=150, default=""
    )
    company_name = serializers.CharField(
        required=False, allow_blank=True, max_length=256, default=""
    )

    def validate_email(self, value: str) -> str:
        email = value.strip().lower()
        if CustomUser.objects.filter(email__iexact=email).exists():
            raise serializers.ValidationError("A user with this email already exists.")
        if CustomUser.objects.filter(username__iexact=email).exists():
            raise serializers.ValidationError("A user with this email already exists.")
        return email

    def create(self, validated_data):
        email = validated_data["email"]
        password = validated_data["password"]
        first_name = (validated_data.get("first_name") or "").strip()
        last_name = (validated_data.get("last_name") or "").strip()
        company_name = (validated_data.get("company_name") or "").strip()
        if not company_name:
            company_name = f"{email}'s workspace"

        company = Company.objects.create(
            name=company_name,
            state="NA",
            zipcode="00000",
            address="N/A",
        )
        user = CustomUser.objects.create_user(
            username=email,
            email=email,
            password=password,
            first_name=first_name,
            last_name=last_name,
            company=company,
            is_company_manager=True,
        )
        return user


class ConversationSerializer(serializers.ModelSerializer):
    class Meta:
        model = Conversation
        fields = ("title", "created", "last_modified", "id")


class PromptSerializer(serializers.ModelSerializer):

    class Meta:
        model = Prompt
        fields = (
            "message",
            "user_created",
            "created",
            "id",
        )


class BasicUserSerializer(serializers.ModelSerializer):
    class Meta:
        model = CustomUser
        fields = (
            "email",
            "first_name",
            "last_name",
            "is_active",
            "has_usable_password",
            "is_company_manager",
            "has_signed_tos",
        )


# document serializers
class DocumentWorkspaceSerializer(serializers.ModelSerializer):
    class Meta:
        model = DocumentWorkspace
        fields = ["id", "name", "created"]
        read_only_fields = ["id", "created"]


class DocumentSerializer(serializers.ModelSerializer):
    class Meta:
        model = Document
        fields = [
            "id",
            "workspace",
            "file",
            "uploaded_at",
            "processed",
            "created",
            "active",
        ]
        read_only_fields = ["id", "uploaded_at", "processed", "created"]