feat: adicionando possibilidade de adicionar models do hugging face, mas deixando o chatGPT como padrão

endpoint_teste/serializer.py

@@ -16,4 +16,5 @@ class TesteSerializer(serializers.Serializer):
 class PDFUploadSerializer(serializers.Serializer):
     file = serializers.FileField()
     system_prompt = serializers.CharField(required=True)
-    user_message = serializers.CharField(required=True)
+    user_message = serializers.CharField(required=True)
+    model = serializers.CharField(required=False)

endpoint_teste/views.py

@@ -9,6 +9,7 @@ from rest_framework.response import Response
 
 from langchain_backend.main import get_llm_answer
 from .serializer import TesteSerializer
+from langchain_huggingface import HuggingFaceEndpoint
 
 class EndpointTesteViewSet(viewsets.ModelViewSet):
     """Mostrará todas as tarefas"""
@@ -31,19 +32,27 @@ def getTeste(request):
                 "Resposta": resposta_llm
             })
     if request.method == "GET":
-        hugging_face_token = os.environ.get("hugging_face_token")
-        API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B"
-        headers = {"Authorization": "Bearer " + hugging_face_token}
-        def query(payload):
-            response = requests.post(API_URL, headers=headers, json=payload)
-            return response.json()
+        # hugging_face_token = os.environ.get("hugging_face_token")
+        # API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B"
+        # headers = {"Authorization": "Bearer " + hugging_face_token}
+        # def query(payload):
+        #     response = requests.post(API_URL, headers=headers, json=payload)
+        #     return response.json()
             
-        output = query({
-            "inputs": "Can you please let us know more details about your something I don't know",
-        })
-        print('output: ', output)
-        print('output: ', dir(output))
-        return Response(output)
+        # output = query({
+        #     "inputs": "Can you please let us know more details about your something I don't know",
+        # })
+        # print('output: ', output)
+        # print('output: ', dir(output))
+        llm = HuggingFaceEndpoint(
+        repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+        task="text-generation",
+        max_new_tokens=100,
+        do_sample=False,
+        )
+        result = llm.invoke("Hugging Face is")
+        print('result: ', result)
+        return Response(result)
 
 @api_view(["POST"])
 def getPDF(request):
@@ -72,7 +81,12 @@ def getPDF(request):
                     temp_file.write(chunk)
                 temp_file_path = temp_file.name  # Get the path of the temporary file
             print('temp_file_path: ', temp_file_path)
-            resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path)
+
+            resposta_llm = None
+            if serializer.validated_data['model']:
+                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=serializer.validated_data['model'])
+            else:
+                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path)
 
             os.remove(temp_file_path)
 

langchain_backend/main.py

@@ -4,13 +4,17 @@ from langchain.chains import create_retrieval_chain
 
 os.environ.get("OPENAI_API_KEY")
 
-def get_llm_answer(system_prompt, user_prompt, pdf_url):
+def get_llm_answer(system_prompt, user_prompt, pdf_url, model):
   pages = None
   if pdf_url:
     pages = getPDF(pdf_url)
   else:
     pages = getPDF()
   retriever = create_retriever(pages)
-  rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt))
+  rag_chain = None
+  if model:
+    rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
+  else:
+    rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt))
   results = rag_chain.invoke({"input": user_prompt})
   return results

langchain_backend/utils.py

@@ -28,8 +28,8 @@ def create_retriever(documents):
   
   return retriever
 
-def create_prompt_llm_chain(system_prompt):
-  model = ChatOpenAI(model="gpt-4o-mini")
+def create_prompt_llm_chain(system_prompt, model="gpt-4o-mini"):
+  model = ChatOpenAI(model=model)
 
   system_prompt = system_prompt + "\n\n" + "{context}"
   prompt = ChatPromptTemplate.from_messages(