feat

endpoint_teste/views.py

@@ -3,6 +3,7 @@ from rest_framework import viewsets, filters
 from django_filters.rest_framework import DjangoFilterBackend
 from endpoint_teste.models import EndpointTesteModel
 from endpoint_teste.serializer import EndpointTesteSerializer, PDFUploadSerializer
+from ..environment import default_model
 
 from rest_framework.decorators import api_view
 from rest_framework.response import Response
@@ -87,7 +88,7 @@ def getPDF(request):
             try:
                 resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=serializer.validated_data['model'])
             except:
-                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path)
+                resposta_llm = get_llm_answer(data["system_prompt"], data["user_message"], temp_file_path, model=default_model)
 
             os.remove(temp_file_path)
 

environment.py

@@ -0,0 +1 @@
+default_model = "gpt-4o-mini"

langchain_backend/main.py

@@ -12,10 +12,11 @@ def get_llm_answer(system_prompt, user_prompt, pdf_url, model):
   else:
     pages = getPDF()
   retriever = create_retriever(pages)
-  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))
+  # rag_chain = None
+  rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
+  # 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

@@ -7,6 +7,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.chains.combine_documents import create_stuff_documents_chain
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
+from ..environment import default_model
 
 os.environ.get("OPENAI_API_KEY")
 os.environ.get("HUGGINGFACEHUB_API_TOKEN")
@@ -30,14 +31,10 @@ def create_retriever(documents):
   
   return retriever
 
-def create_prompt_llm_chain(system_prompt, modelParam="gpt-4o-mini"):
-  if modelParam == "gpt-4o-mini":
+def create_prompt_llm_chain(system_prompt, modelParam):
+  if modelParam == default_model:
     model = ChatOpenAI(model=modelParam)
   else:
-    
-    print('\n\n\n')
-    print(os.environ.get("HUGGINGFACEHUB_API_TOKEN"))
-
     model = HuggingFaceEndpoint(
         repo_id=modelParam,
         task="text-generation",