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})