cmhub/apps/api/views.py

from rest_framework.exceptions import AuthenticationFailed
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework import status
from rest_framework.views import APIView

from apps.api.authentication import ApiKeyAuthentication
from apps.api.errors import api_error
from apps.api.generation import (
    ApiRequestError,
    generate_image_response,
    generate_title_response,
)
from apps.api.serializers import (
    GenerateImageRequestSerializer,
    GenerateTitleRequestSerializer,
)
from apps.billing.services import get_balance_snapshot


class ExternalApiView(APIView):
    authentication_classes = (ApiKeyAuthentication,)
    permission_classes = (IsAuthenticated,)

    def permission_denied(self, request, message=None, code=None):
        if request.authenticators and not request.successful_authenticator:
            raise AuthenticationFailed(api_error("unauthorized", "缺失或无效 API Key"))
        super().permission_denied(request, message=message, code=code)


class GenerateTitleView(ExternalApiView):
    def post(self, request):
        serializer = GenerateTitleRequestSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(
                api_error("bad_request", "参数错误"),
                status=status.HTTP_400_BAD_REQUEST,
            )
        try:
            data = generate_title_response(
                user=request.user,
                api_key=request.auth,
                request_data=serializer.validated_data,
            )
        except ApiRequestError as exc:
            return Response(exc.as_response_data(), status=exc.http_status)
        return Response(data, status=status.HTTP_200_OK)


class GenerateImageView(ExternalApiView):
    def post(self, request):
        serializer = GenerateImageRequestSerializer(data=request.data)
        if not serializer.is_valid():
            return Response(
                api_error("bad_request", "参数错误"),
                status=status.HTTP_400_BAD_REQUEST,
            )
        try:
            data = generate_image_response(
                user=request.user,
                api_key=request.auth,
                request=request,
                request_data=serializer.validated_data,
            )
        except ApiRequestError as exc:
            return Response(exc.as_response_data(), status=exc.http_status)
        return Response(data, status=status.HTTP_200_OK)


class BalanceView(ExternalApiView):
    def get(self, request):
        balance = get_balance_snapshot(request.user)
        return Response(
            {
                "user": request.user.get_username(),
                "points_balance": balance.points_balance,
            },
            status=status.HTTP_200_OK,
        )