pdfs/views.py

import tempfile, os
from pdfs.serializer import PDFUploadSerializer
from setup.environment import default_model
from drf_spectacular.utils import extend_schema

from rest_framework.decorators import api_view, parser_classes
from rest_framework.parsers import MultiPartParser
from rest_framework.response import Response

from _utils.main import get_llm_answer

@extend_schema(
    request=PDFUploadSerializer,
)
@api_view(["POST"])
@parser_classes([MultiPartParser])
def getPDF(request):
    if request.method == "POST":
        serializer = PDFUploadSerializer(data=request.data)
        if serializer.is_valid(raise_exception=True):
            listaPDFs = []
            print('\n\n')
            data = request.data
            print('data: ', data)
            embedding = serializer.validated_data.get("embedding", "gpt")
            model = serializer.validated_data.get("model", default_model)

            # pdf_file = serializer.validated_data['file']
            for file in serializer.validated_data['files']:
                print("file: ", file)
                file.seek(0)
                # Create a temporary file to save the uploaded PDF
                with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
                    # Write the uploaded file content to the temporary file
                    for chunk in file.chunks():
                        temp_file.write(chunk)
                    temp_file_path = temp_file.name  # Get the path of the temporary file
                    listaPDFs.append(temp_file_path)
                # print('temp_file_path: ', temp_file_path)
                print('listaPDFs: ', listaPDFs)

            resposta_llm = None
            # resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=model, embedding=embedding)
            resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], listaPDFs, model=model, embedding=embedding)

            for file in listaPDFs:
                os.remove(file)
            # os.remove(temp_file_path)

            return Response({
                "Resposta": resposta_llm
            })