resumos/views.py

from rest_framework.views import APIView
import tempfile, os
from rest_framework.response import Response

from langchain_backend.utils import DEFAULT_SYSTEM_PROMPT
from .serializer import ResumoPDFSerializer
from langchain_backend.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, "Por favor faça com que o tamanho do conteúdo total gerado seja no formato JSON e que tenha pelo menos 20000 caracteres", 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})