chat_backend/llm_be/chat_backend/serializers.py

from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
from rest_framework import serializers
from django.db.models import Count, Q, Sum

from .models import (
    CustomUser,
    Announcement,
    Company,
    Conversation,
    Prompt,
    PromptMetric,
    Feedback,
    FEEDBACK_CATEGORIES,
    DocumentWorkspace,
    Document,
    DriveConnection,
)


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

    def validate(self, attrs):
        data = super().validate(attrs)
        if getattr(self.user, "deleted", False):
            from rest_framework_simplejwt.exceptions import AuthenticationFailed

            raise AuthenticationFailed(
                "No active account found with the given credentials",
                code="no_active_account",
            )
        return data


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()
    subscription = serializers.SerializerMethodField()

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

    def get_subscription(self, obj):
        from finance.services.plans import needs_checkout, plan_to_dict
        from finance.models import UserSubscription

        try:
            sub = obj.subscription
        except UserSubscription.DoesNotExist:
            return {
                "plan": None,
                "status": UserSubscription.Status.NONE,
                "source": UserSubscription.Source.NONE,
                "needs_checkout": True,
                "cancel_at_period_end": False,
                "current_period_end": None,
            }
        return {
            "plan": plan_to_dict(sub.plan) if sub.plan_id else None,
            "status": sub.status,
            "source": sub.source,
            "needs_checkout": needs_checkout(obj),
            "cancel_at_period_end": bool(sub.cancel_at_period_end),
            "current_period_end": (
                sub.current_period_end.isoformat() if sub.current_period_end else None
            ),
        }


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):
        from finance.services.plans import try_redeem_backer_email

        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,
        )
        try_redeem_backer_email(user)
        return user


def _conversation_token_totals(conversation_id: int):
    """
    Sum PromptMetric tokens for a conversation.

    Returns nulls when the provider never reported usage (never fabricate 0).
    """
    agg = PromptMetric.objects.filter(conversation_id=conversation_id).aggregate(
        tin=Sum("tokens_in"),
        tout=Sum("tokens_out"),
        with_in=Count("id", filter=Q(tokens_in__isnull=False)),
        with_out=Count("id", filter=Q(tokens_out__isnull=False)),
    )
    return (
        agg["tin"] if agg["with_in"] else None,
        agg["tout"] if agg["with_out"] else None,
    )


def _prompt_token_pair(prompt_id: int):
    metric = (
        PromptMetric.objects.filter(prompt_id=prompt_id)
        .order_by("-created")
        .only("tokens_in", "tokens_out")
        .first()
    )
    if metric is None:
        return None, None
    return metric.tokens_in, metric.tokens_out


class ConversationSerializer(serializers.ModelSerializer):
    tokens_in = serializers.SerializerMethodField()
    tokens_out = serializers.SerializerMethodField()

    class Meta:
        model = Conversation
        fields = (
            "title",
            "created",
            "last_modified",
            "id",
            "tokens_in",
            "tokens_out",
        )

    def _token_pair(self, obj):
        cache = self.context.setdefault("_conversation_token_cache", {})
        if obj.id not in cache:
            cache[obj.id] = _conversation_token_totals(obj.id)
        return cache[obj.id]

    def get_tokens_in(self, obj):
        tin, _ = self._token_pair(obj)
        return tin

    def get_tokens_out(self, obj):
        _, tout = self._token_pair(obj)
        return tout


class PromptSerializer(serializers.ModelSerializer):
    tokens_in = serializers.SerializerMethodField()
    tokens_out = serializers.SerializerMethodField()

    class Meta:
        model = Prompt
        fields = (
            "message",
            "user_created",
            "created",
            "id",
            "tokens_in",
            "tokens_out",
            "citations",
        )
        read_only_fields = ("citations",)

    def _token_pair(self, obj):
        cache = self.context.setdefault("_prompt_token_cache", {})
        if obj.id not in cache:
            cache[obj.id] = _prompt_token_pair(obj.id)
        return cache[obj.id]

    def get_tokens_in(self, obj):
        tin, _ = self._token_pair(obj)
        return tin

    def get_tokens_out(self, obj):
        _, tout = self._token_pair(obj)
        return tout

    def validate_message(self, value: str) -> str:
        if value is None or not str(value).strip():
            raise serializers.ValidationError("Message text cannot be empty.")
        return str(value).strip()


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"]


# drive connection serializers (#47-#52)
class DriveConnectionSerializer(serializers.ModelSerializer):
    """Never exposes access_token/refresh_token to the client."""

    class Meta:
        model = DriveConnection
        fields = [
            "id",
            "provider",
            "kind",
            "external_account_email",
            "selected_resource_ids",
            "selected_resource_labels",
            "last_sync_at",
            "last_sync_status",
            "last_sync_error",
            "sync_total",
            "sync_processed",
            "sync_added",
            "sync_updated",
            "sync_failed",
            "is_active",
            "created",
        ]
        read_only_fields = fields


class DriveConnectionResourcesSerializer(serializers.Serializer):
    resource_ids = serializers.ListField(
        child=serializers.CharField(max_length=512), allow_empty=True
    )
    resource_labels = serializers.ListField(
        child=serializers.CharField(max_length=512, allow_blank=True),
        allow_empty=True,
        required=False,
        default=list,
    )