File size: 3,203 Bytes
4e93adb
 
 
4cd3056
 
4e93adb
7dc6d22
4e93adb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4d3bceb
4e93adb
4d3bceb
4e93adb
 
 
 
 
 
 
 
 
 
 
4cd3056
e4d76bf
4e93adb
7dc6d22
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4e93adb
 
 
4dcf767
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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, 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})