MathApp/rag/views.py

import base64
import binascii
import logging

from drf_spectacular.types import OpenApiTypes
from drf_spectacular.utils import OpenApiParameter, extend_schema
from rest_framework.exceptions import APIException
from rest_framework.parsers import FormParser, JSONParser, MultiPartParser
from rest_framework.response import Response
from rest_framework.throttling import ScopedRateThrottle
from rest_framework.views import APIView

from payment.permissions import HasInternalServiceKey

from .models import LessonChunk
from .serializers import (
    ChunkListResponseSerializer,
    ChunkSerializer,
    IngestPdfRequestSerializer,
    IngestResponseSerializer,
    IngestTextRequestSerializer,
)
from .services import ingest as ingest_service
from .services.clients import RagNotConfigured

logger = logging.getLogger(__name__)

_INTERNAL_KEY_PARAM = OpenApiParameter(
    'X-Internal-Key', OpenApiTypes.STR, OpenApiParameter.HEADER, required=True,
    description='کلید مخفیِ فراخوانی داخلی/سرور-به-سرور (settings.INTERNAL_SERVICE_KEY).',
)


class RagProviderUnavailable(APIException):
    status_code = 503
    default_detail = 'سرویس هوش مصنوعی/بردار موقتاً در دسترس نیست. لطفاً بعداً تلاش کنید.'
    default_code = 'rag_provider_unavailable'


def _ingest_response(result: dict) -> Response:
    return Response(IngestResponseSerializer({
        'status': result['status'],
        'collection': result['collection'],
        'chunksCreated': result['chunks_created'],
        'chunkIds': result.get('chunk_ids', []),
    }).data)


class IngestTextView(APIView):
    """POST /api/rag/ingest/text — آپلودِ متنِ یک درس، chunk/خلاصه/embedding و ذخیره در DB + Qdrant."""

    permission_classes = [HasInternalServiceKey]
    throttle_scope = 'rag_ingest'
    throttle_classes = [ScopedRateThrottle]

    @extend_schema(
        request=IngestTextRequestSerializer,
        responses=IngestResponseSerializer,
        parameters=[_INTERNAL_KEY_PARAM],
        tags=['RAG (rag)'],
    )
    def post(self, request):
        serializer = IngestTextRequestSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        v = serializer.validated_data
        try:
            result = ingest_service.ingest_text(
                lesson_id=v['lessonId'],
                grade_id=v.get('gradeId'),
                subject=v['subject'],
                lesson_number=v['lessonNumber'],
                text=v['text'],
                chapter=v.get('chapter'),
                title=v.get('title') or '',
            )
        except RagNotConfigured as exc:
            logger.warning('rag ingest_text: not configured: %s', exc)
            raise RagProviderUnavailable(detail=str(exc))
        except Exception:
            logger.exception('rag ingest_text failed')
            raise RagProviderUnavailable()
        return _ingest_response(result)


class IngestPdfView(APIView):
    """POST /api/rag/ingest/pdf — آپلودِ PDF کتاب، استخراجِ صفحات با vision و ذخیره در DB + Qdrant."""

    permission_classes = [HasInternalServiceKey]
    throttle_scope = 'rag_ingest'
    throttle_classes = [ScopedRateThrottle]
    parser_classes = [MultiPartParser, FormParser, JSONParser]

    @extend_schema(
        request=IngestPdfRequestSerializer,
        responses=IngestResponseSerializer,
        parameters=[_INTERNAL_KEY_PARAM],
        tags=['RAG (rag)'],
    )
    def post(self, request):
        serializer = IngestPdfRequestSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        v = serializer.validated_data

        upload = v.get('pdf')
        if upload is not None:
            pdf_bytes = upload.read()
        else:
            try:
                pdf_bytes = base64.b64decode(v['pdfBase64'], validate=True)
            except (binascii.Error, ValueError):
                raise RagProviderUnavailable(detail='pdfBase64 معتبر نیست.')

        try:
            result = ingest_service.ingest_pdf(
                lesson_id=v['lessonId'],
                grade_id=v.get('gradeId'),
                subject=v['subject'],
                lesson_number=v['lessonNumber'],
                pdf_bytes=pdf_bytes,
                chapter=v.get('chapter'),
                title=v.get('title') or '',
                chapter_name=v.get('chapterName') or '',
            )
        except RagNotConfigured as exc:
            logger.warning('rag ingest_pdf: not available: %s', exc)
            raise RagProviderUnavailable(detail=str(exc))
        except Exception:
            logger.exception('rag ingest_pdf failed')
            raise RagProviderUnavailable()
        return _ingest_response(result)


class ChunkListView(APIView):
    """GET /api/rag/chunks — فهرستِ chunkهای ذخیره‌شده (بازبینی)، فیلترِ lessonId/subject."""

    permission_classes = [HasInternalServiceKey]

    @extend_schema(
        responses=ChunkListResponseSerializer,
        parameters=[
            _INTERNAL_KEY_PARAM,
            OpenApiParameter('lessonId', OpenApiTypes.INT, OpenApiParameter.QUERY, required=False),
            OpenApiParameter('subject', OpenApiTypes.STR, OpenApiParameter.QUERY, required=False),
        ],
        tags=['RAG (rag)'],
    )
    def get(self, request):
        qs = LessonChunk.objects.all()
        lesson_id = request.query_params.get('lessonId')
        subject = request.query_params.get('subject')
        if lesson_id:
            qs = qs.filter(lesson_id=lesson_id)
        if subject:
            qs = qs.filter(subject=subject)
        rows = list(qs.order_by('-id')[:200])
        data = {'count': len(rows), 'results': ChunkSerializer(rows, many=True).data}
        return Response(ChunkListResponseSerializer(data).data)