resumos/views.py

from rest_framework.views import APIView
import tempfile, os
from rest_framework.response import Response
from _utils.resumo_simples_cursor import get_llm_summary_answer_by_cursor
from _utils.utils import DEFAULT_SYSTEM_PROMPT
from .serializer import (
    ResumoPDFSerializer,
    ResumoCursorSerializer,
)
from _utils.main import get_llm_answer_summary, get_llm_answer_summary_with_embedding
from setup.environment import default_model
from rest_framework.parsers import MultiPartParser
from drf_spectacular.utils import extend_schema


class ResumoView(APIView):
    parser_classes = [MultiPartParser]

    @extend_schema(
        request=ResumoPDFSerializer,
    )
    def post(self, request):
        serializer = ResumoPDFSerializer(data=request.data)
        if serializer.is_valid(raise_exception=True):
            listaPDFs = []
            data = serializer.validated_data
            model = serializer.validated_data.get("model", default_model)
            print("serializer.validated_data: ", serializer.validated_data)

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

            system_prompt = data.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
            resposta_llm = get_llm_answer_summary(
                system_prompt,
                data["user_message"],
                listaPDFs,
                model=model,
                isIterativeRefinement=data["iterative_refinement"],
            )

            for file in listaPDFs:
                os.remove(file)

            return Response({"resposta": resposta_llm})


class ResumoEmbeddingView(APIView):
    parser_classes = [MultiPartParser]

    @extend_schema(
        request=ResumoPDFSerializer,
    )
    def post(self, request):
        serializer = ResumoPDFSerializer(data=request.data)
        if serializer.is_valid(raise_exception=True):
            listaPDFs = []
            data = serializer.validated_data
            model = serializer.validated_data.get("model", default_model)
            print("serializer.validated_data: ", serializer.validated_data)

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

            system_prompt = data.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
            resposta_llm = get_llm_answer_summary_with_embedding(
                system_prompt,
                data["user_message"],
                listaPDFs,
                model=model,
                isIterativeRefinement=data["iterative_refinement"],
            )

            for file in listaPDFs:
                os.remove(file)

            return Response({"resposta": resposta_llm})


class ResumoSimplesCursorView(APIView):
    parser_classes = [MultiPartParser]

    @extend_schema(
        request=ResumoCursorSerializer,
    )
    def post(self, request):
        serializer = ResumoCursorSerializer(data=request.data)
        if serializer.is_valid(raise_exception=True):
            listaPDFs = []
            data = serializer.validated_data
            print("\nserializer.validated_data: ", serializer.validated_data)

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

            resposta_llm = get_llm_summary_answer_by_cursor(data, listaPDFs)

            for file in listaPDFs:
                os.remove(file)

            return Response({"resposta": resposta_llm})