Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 from typing import Optional
 import gradio as gr
-from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
@@ -14,6 +14,7 @@ import tempfile
 # Configurações
 EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
 LLM_MODEL = "mistralai/Mistral-7B-v0.1"
+DOCS_DIR = "documents"  # Pasta com os documentos base
 
 class RAGSystem:
     def __init__(self):
@@ -23,7 +24,7 @@ class RAGSystem:
             LLM_MODEL,
             torch_dtype=torch.float16,
             device_map="auto",
-            load_in_8bit=True  # Usa quantização 8-bit para reduzir uso de memória
+            load_in_8bit=True
         )
         
         # Configura o pipeline
@@ -46,8 +47,49 @@ class RAGSystem:
             model_kwargs={'device': 'cpu'}
         )
 
+        # Carrega a base de conhecimento permanente
+        self.base_db = self.load_base_knowledge()
+
+    def load_base_knowledge(self) -> Optional[FAISS]:
+        """Carrega a base de conhecimento permanente da pasta de documentos"""
+        try:
+            if not os.path.exists(DOCS_DIR):
+                print(f"Pasta {DOCS_DIR} não encontrada. Criando...")
+                os.makedirs(DOCS_DIR)
+                return None
+
+            # Carrega todos os PDFs da pasta
+            loader = DirectoryLoader(
+                DOCS_DIR,
+                glob="**/*.pdf",
+                loader_cls=PyPDFLoader
+            )
+            documents = loader.load()
+            
+            if not documents:
+                print("Nenhum documento encontrado na pasta base.")
+                return None
+            
+            # Divide o texto em chunks
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=1000,
+                chunk_overlap=200,
+                length_function=len,
+                separators=["\n\n", "\n", ".", " ", ""]
+            )
+            texts = text_splitter.split_documents(documents)
+            
+            # Cria base de conhecimento
+            print(f"Criando base de conhecimento com {len(texts)} chunks...")
+            db = FAISS.from_documents(texts, self.embeddings)
+            return db
+            
+        except Exception as e:
+            print(f"Erro ao carregar base de conhecimento: {str(e)}")
+            return None
+
     def process_pdf(self, file_content: bytes) -> Optional[FAISS]:
-        """Processa o PDF e cria a base de conhecimento"""
+        """Processa o PDF do usuário"""
         try:
             # Cria arquivo temporário
             with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
@@ -75,6 +117,11 @@ class RAGSystem:
             
             # Cria base de conhecimento
             db = FAISS.from_documents(texts, self.embeddings)
+            
+            # Se existir uma base permanente, mescla com ela
+            if self.base_db is not None:
+                db.merge_from(self.base_db)
+            
             return db
             
         except Exception as e:
@@ -83,17 +130,20 @@ class RAGSystem:
 
     def generate_response(self, file_obj, query: str) -> str:
         """Gera resposta para a consulta"""
-        if file_obj is None:
-            return "Por favor, faça upload de um arquivo PDF."
-        
         if not query.strip():
             return "Por favor, insira uma pergunta."
         
         try:
-            # Processa o PDF
-            db = self.process_pdf(file_obj)
-            if db is None:
-                return "Não foi possível processar o PDF."
+            # Se tiver arquivo do usuário, processa e mescla com a base
+            if file_obj is not None:
+                db = self.process_pdf(file_obj)
+                if db is None:
+                    return "Não foi possível processar o PDF."
+            # Se não tiver arquivo do usuário, usa só a base permanente
+            elif self.base_db is not None:
+                db = self.base_db
+            else:
+                return "Nenhuma base de conhecimento disponível. Por favor, faça upload de um PDF ou adicione documentos à pasta base."
             
             # Configura o chain RAG
             qa_chain = RetrievalQA.from_chain_type(
@@ -108,8 +158,16 @@ class RAGSystem:
                 return_source_documents=True
             )
             
+            # Adiciona contexto sobre a fonte da resposta
+            prompt = f"""Com base nos documentos fornecidos, responda à seguinte pergunta:
+            {query}
+            
+            Se a resposta vier da base de documentos permanente, indique isso no início da resposta.
+            Se a resposta vier do PDF enviado pelo usuário, indique isso no início da resposta.
+            Se não encontrar informações suficientes, indique isso claramente."""
+            
             # Gera resposta
-            result = qa_chain({"query": query})
+            result = qa_chain({"query": prompt})
             return result["result"]
             
         except Exception as e:
@@ -120,18 +178,18 @@ def create_demo():
     rag = RAGSystem()
     
     with gr.Blocks() as demo:
-        gr.Markdown("# 📚 Sistema RAG com Mistral-7B")
-        gr.Markdown("""
+        gr.Markdown("# 📚 Sistema RAG com Base de Conhecimento")
+        gr.Markdown(f"""
         ### Instruções:
-        1. Faça upload de um arquivo PDF
-        2. Digite sua pergunta sobre o conteúdo
-        3. Aguarde a resposta gerada pelo modelo
+        1. Os documentos da pasta `{DOCS_DIR}` são usados como base de conhecimento permanente
+        2. Você pode fazer upload de PDFs adicionais para consulta
+        3. As respostas serão baseadas em ambas as fontes quando disponíveis
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
                 file_input = gr.File(
-                    label="Upload do PDF",
+                    label="Upload do PDF (opcional)",
                     type="binary",
                     file_types=[".pdf"]
                 )
@@ -156,7 +214,7 @@ def create_demo():
         
         gr.Examples(
             examples=[
-                [None, "Qual é o tema principal deste documento?"],
+                [None, "Qual é o tema principal dos documentos?"],
                 [None, "Pode fazer um resumo dos pontos principais?"],
                 [None, "Quais são as principais conclusões?"]
             ],