Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Dict
 import gradio as gr
 from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -31,8 +31,9 @@ class RAGSystem:
             model=self.model,
             tokenizer=self.tokenizer,
             max_length=512,
-            temperature=0.7,
-            top_p=0.95
+            temperature=0.3,  # Reduzido para respostas mais precisas
+            top_p=0.95,
+            repetition_penalty=1.2
         )
         
         self.llm = HuggingFacePipeline(pipeline=pipe)
@@ -104,6 +105,26 @@ class RAGSystem:
             print(f"Erro ao processar PDF: {str(e)}")
             return None
 
+    def format_response(self, raw_response: str, source_type: str, context_found: bool) -> str:
+        """Formata a resposta para um formato padronizado e claro"""
+        if not context_found:
+            return "🔍 Não foram encontradas informações suficientes nos documentos para responder esta pergunta."
+        
+        prefix = ""
+        if source_type == "pdf":
+            prefix = "📄 [Resposta baseada no PDF enviado]\n\n"
+        elif source_type == "base":
+            prefix = "📚 [Resposta baseada na base de documentos]\n\n"
+        elif source_type == "both":
+            prefix = "📚📄 [Resposta baseada em ambas as fontes]\n\n"
+            
+        # Limpa e formata a resposta
+        response = raw_response.strip()
+        if not response:
+            return "🔍 Não foi possível gerar uma resposta adequada com as informações disponíveis."
+            
+        return f"{prefix}{response}"
+
     def generate_response(self, file_obj, query: str, progress=gr.Progress()) -> Tuple[str, str, str]:
         """Retorna (resposta, status, tempo_decorrido)"""
         if not query.strip():
@@ -111,42 +132,68 @@ class RAGSystem:
         
         start_time = time.time()
         try:
-            progress(0, desc="Iniciando processamento...")
+            progress(0.2, desc="Processando documentos...")
+            
+            # Determina a fonte dos documentos
+            has_pdf = file_obj is not None
+            has_base = self.base_db is not None
+            source_type = "both" if has_pdf and has_base else "pdf" if has_pdf else "base" if has_base else None
+            
+            if not source_type:
+                return "Nenhuma fonte de documentos disponível.", "❌ Sem documentos", "0s"
             
             # Processa documento
-            progress(0.2, desc="Processando documento...")
-            if file_obj is not None:
+            if has_pdf:
                 db = self.process_pdf(file_obj)
                 if db is None:
                     return "Não foi possível processar o PDF.", "❌ Erro no processamento", "0s"
-            elif self.base_db is not None:
-                db = self.base_db
             else:
-                return "Nenhuma base de conhecimento disponível.", "❌ Sem documentos", "0s"
+                db = self.base_db
             
             progress(0.4, desc="Buscando informações relevantes...")
+            
+            # Configuração do RAG
             qa_chain = RetrievalQA.from_chain_type(
                 llm=self.llm,
                 chain_type="stuff",
                 retriever=db.as_retriever(
-                    search_kwargs={"k": 4, "fetch_k": 6}
+                    search_kwargs={
+                        "k": 4,
+                        "fetch_k": 6,
+                        "score_threshold": 0.5  # Filtra resultados pouco relevantes
+                    }
                 ),
                 return_source_documents=True
             )
             
             progress(0.6, desc="Gerando resposta...")
-            prompt = f"""Baseado nos documentos fornecidos, responda em português à seguinte pergunta:
-            {query}
             
-            Se a resposta vier da base de documentos permanente, indique isso no início.
-            Se a resposta vier do PDF enviado, indique isso no início.
-            Se não encontrar informações suficientes, indique isso claramente."""
+            # Prompt mais estruturado
+            prompt = f"""Instruções:
+            1. Analise cuidadosamente os documentos fornecidos
+            2. Responda à seguinte pergunta em português de forma clara e direta: {query}
+            3. Use apenas informações encontradas nos documentos
+            4. Se não houver informações suficientes, indique explicitamente
+            5. Mantenha a resposta objetiva e baseada em fatos
+            6. Cite exemplos específicos dos documentos quando relevante"""
             
+            # Gera resposta
             result = qa_chain({"query": prompt})
-            elapsed_time = f"{time.time() - start_time:.1f}s"
             
+            # Verifica se encontrou contexto relevante
+            context_found = bool(result.get("source_documents", []))
+            
+            # Formata a resposta
+            formatted_response = self.format_response(
+                result["result"],
+                source_type,
+                context_found
+            )
+            
+            elapsed_time = f"{time.time() - start_time:.1f}s"
             progress(1.0, desc="Concluído!")
-            return result["result"], "✅ Sucesso", elapsed_time
+            
+            return formatted_response, "✅ Sucesso", elapsed_time
             
         except Exception as e:
             elapsed_time = f"{time.time() - start_time:.1f}s"
@@ -162,7 +209,7 @@ def create_demo():
                 """
                 # 🤖 Assistente de Documentos Inteligente
                 
-                Este sistema usa tecnologia RAG (Retrieval-Augmented Generation) para responder perguntas sobre seus documentos.
+                Sistema de consulta avançada que responde perguntas sobre seus documentos usando RAG.
                 """
             )
             
@@ -180,14 +227,16 @@ def create_demo():
                         )
                         info = gr.Markdown(
                             f"""
-                            ℹ️ Além do upload, o sistema também consulta a pasta `{DOCS_DIR}`
+                            ℹ️ O sistema consulta:
+                            - PDFs enviados por você
+                            - Documentos na pasta `{DOCS_DIR}`
                             """
                         )
                     
                     with gr.Group():
                         gr.Markdown("### ❓ Sua Pergunta")
                         query_input = gr.Textbox(
-                            placeholder="Digite sua pergunta aqui...",
+                            placeholder="Digite sua pergunta sobre os documentos...",
                             lines=3,
                             max_lines=6,
                             show_label=False,
@@ -227,10 +276,10 @@ def create_demo():
             with gr.Accordion("📚 Exemplos de Perguntas", open=False):
                 gr.Examples(
                     examples=[
-                        [None, "Qual é o tema principal dos documentos?"],
-                        [None, "Pode resumir os pontos principais?"],
-                        [None, "Quais são as principais conclusões?"],
-                        [None, "Explique o contexto deste documento."],
+                        [None, "Quais são os principais tópicos abordados neste documento?"],
+                        [None, "Resuma as conclusões mais importantes."],
+                        [None, "O que o documento diz sobre [tema específico]?"],
+                        [None, "Quais são as recomendações apresentadas?"],
                     ],
                     inputs=[file_input, query_input],
                 )
@@ -239,11 +288,11 @@ def create_demo():
             gr.Markdown(
                 """
                 ---
-                ### 🔧 Sobre o Sistema
-                * Usa modelo T5 para geração de respostas
-                * Processamento de documentos com tecnologia RAG
-                * Suporte a múltiplos documentos PDF
-                * Respostas baseadas apenas no conteúdo dos documentos
+                ### 🔧 Informações do Sistema
+                * Respostas geradas usando tecnologia RAG (Retrieval-Augmented Generation)
+                * Processamento inteligente de documentos PDF
+                * Respostas baseadas exclusivamente no conteúdo dos documentos
+                * Suporte a múltiplos documentos e contextos
                 """
             )