feat: adicionando possibilidade de chamar criar um resumo com iterative_refinement

langchain_backend/main.py

@@ -1,10 +1,11 @@
 import os
-from langchain_backend.utils import create_prompt_llm_chain, create_retriever, getPDF
+from langchain_backend.utils import create_prompt_llm_chain, create_retriever, getPDF, create_llm
 from langchain_backend import utils
 from langchain.chains import create_retrieval_chain
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_chroma import Chroma
 from langchain_openai import OpenAIEmbeddings
+from langchain.chains.summarize import load_summarize_chain
 
 os.environ.get("OPENAI_API_KEY")
 
@@ -29,18 +30,27 @@ def get_llm_answer(system_prompt, user_prompt, pdf_url, model, embedding):
   retriever = create_retriever(pages, vectorstore)
   rag_chain = create_retrieval_chain(retriever, create_prompt_llm_chain(system_prompt, model))
   results = rag_chain.invoke({"input": user_prompt})
-  print('allIds ARQUIVO MAIN: ',  utils.allIds)
+  # print('allIds ARQUIVO MAIN: ',  utils.allIds)
   vectorstore.delete( utils.allIds)
   vectorstore.delete_collection()
   utils.allIds = []
-  print('utils.allIds: ', utils.allIds)
+  # print('utils.allIds: ', utils.allIds)
   return results
 
-def get_llm_answer_summary(system_prompt, user_prompt, pdf_url, model):
+def get_llm_answer_summary(system_prompt, user_prompt, pdf_url, model, isIterativeRefinement):
   print('model: ', model)
+  print('isIterativeRefinement: ', isIterativeRefinement)
+  print('\n\n\n')
   pages = getPDF(pdf_url)
-  rag_chain = create_prompt_llm_chain(system_prompt, model)
+  if not isIterativeRefinement:
+    rag_chain = create_prompt_llm_chain(system_prompt, model)
 
-  results = rag_chain.invoke({"input": user_prompt, "context": pages})
+    results = rag_chain.invoke({"input": user_prompt, "context": pages})
 
-  return results
+    return results
+  else:
+    chain = load_summarize_chain(create_llm(model), "refine", True)
+    result = chain.invoke({"input_documents": pages})
+    print('result: ', result)
+    # Obs --> Para passar informações personalizadas --> chain = load_summarize_chain(llm, "refine", True, question_prompt=initial_prompt, refine_prompt=PromptTemplate.from_template(refine_prompt))
+    # Para ver mais opções --> Acessa a origem da função load_summarize_chain , e nela acessa a origem da função _load_refine_chain --> As opções são os parâmetros que esta última função recebe

langchain_backend/utils.py

@@ -9,7 +9,10 @@ from langchain_huggingface import HuggingFaceEndpoint, HuggingFaceEmbeddings
 from setup.environment import default_model
 from uuid import uuid4
 
-
+os.environ["LANGCHAIN_TRACING_V2"]="true"
+os.environ["LANGCHAIN_ENDPOINT"]="https://api.smith.langchain.com"
+os.environ.get("LANGCHAIN_API_KEY")
+os.environ["LANGCHAIN_PROJECT"]="VELLA"
 os.environ.get("OPENAI_API_KEY")
 os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
@@ -29,8 +32,8 @@ def getPDF(file_paths):
   # loader = PyPDFLoader(file_paths, extract_images=False)
   # pages = loader.load_and_split(text_splitter)
   for page in pages:
-    print('\n')
-    print('allIds: ', allIds)
+    # print('\n')
+    # print('allIds: ', allIds)
     documentId = str(uuid4())
     allIds.append(documentId)
     page.id = documentId
@@ -50,16 +53,7 @@ def create_retriever(documents, vectorstore):
   return retriever
 
 def create_prompt_llm_chain(system_prompt, modelParam):
-  if modelParam == default_model:
-    model = ChatOpenAI(model=modelParam)
-  else:
-    model = HuggingFaceEndpoint(
-        repo_id=modelParam,
-        task="text-generation",
-        # max_new_tokens=100,
-        do_sample=False,
-        huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
-    )
+  model = create_llm(modelParam)
 
   system_prompt = system_prompt + "\n\n" + "{context}"
   prompt = ChatPromptTemplate.from_messages(
@@ -69,4 +63,16 @@ def create_prompt_llm_chain(system_prompt, modelParam):
       ]
   )
   question_answer_chain = create_stuff_documents_chain(model, prompt)
-  return question_answer_chain
+  return question_answer_chain
+
+def create_llm(modelParam):
+  if modelParam == default_model:
+    return ChatOpenAI(model=modelParam)
+  else:
+    return HuggingFaceEndpoint(
+        repo_id=modelParam,
+        task="text-generation",
+        # max_new_tokens=100,
+        do_sample=False,
+        huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+    )

resumos/serializer.py

@@ -20,5 +20,6 @@ from rest_framework import serializers
 class ResumoPDFSerializer(serializers.Serializer):
     files = serializers.ListField(child=serializers.FileField(), required=True)
     system_prompt = serializers.CharField(required=True)
-    user_message = serializers.CharField(required=False)
-    model = serializers.CharField(required=False)
+    user_message = serializers.CharField(required=False, default="")
+    model = serializers.CharField(required=False)
+    iterative_refinement = serializers.BooleanField(required=False, default=False)

resumos/views.py

@@ -18,9 +18,9 @@ class ResumoView(APIView):
     serializer = ResumoPDFSerializer(data=request.data)
     if serializer.is_valid(raise_exception=True):
         listaPDFs = []
-        data = request.data
+        data = serializer.validated_data
         model = serializer.validated_data.get("model", default_model)
-        user_message = data.get("user_message", "")
+        print('serializer.validated_data: ', serializer.validated_data)
 
         for file in serializer.validated_data['files']:
             print("file: ", file)
@@ -32,7 +32,7 @@ class ResumoView(APIView):
                 listaPDFs.append(temp_file_path)
             # print('listaPDFs: ', listaPDFs)
 
-        resposta_llm = get_llm_answer_summary(data["system_prompt"], user_message, listaPDFs, model=model)
+        resposta_llm = get_llm_answer_summary(data["system_prompt"], data["user_message"], listaPDFs, model=model, isIterativeRefinement=data["iterative_refinement"])
 
         for file in listaPDFs:
             os.remove(file)