app.py

import os
from typing import Optional
import gradio as gr
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
from langchain_community.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch
import tempfile

# Configurações
EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
LLM_MODEL = "google/flan-t5-large"  # Modelo aberto e sem necessidade de autenticação
DOCS_DIR = "documentos"

class RAGSystem:
    def __init__(self):
        # Inicializa o modelo de linguagem
        print("Carregando modelo de linguagem...")
        self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(
            LLM_MODEL,
            device_map="auto",
            torch_dtype=torch.float32  # T5 funciona bem com float32
        )
        
        # Configura o pipeline
        pipe = pipeline(
            "text2text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            max_length=512,
            temperature=0.7,
            top_p=0.95
        )
        
        # Configura o modelo LangChain
        self.llm = HuggingFacePipeline(pipeline=pipe)
        
        # Configura embeddings
        print("Configurando embeddings...")
        self.embeddings = HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            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=500,  # Chunks menores para o T5
                chunk_overlap=100,
                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 do usuário"""
        try:
            # Cria arquivo temporário
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                tmp_file.write(file_content)
                tmp_path = tmp_file.name

            # Carrega e processa o PDF
            loader = PyPDFLoader(tmp_path)
            documents = loader.load()
            
            # Remove arquivo temporário
            os.unlink(tmp_path)
            
            if not documents:
                return None
            
            # Divide o texto em chunks
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=500,
                chunk_overlap=100,
                length_function=len,
                separators=["\n\n", "\n", ".", " ", ""]
            )
            texts = text_splitter.split_documents(documents)
            
            # 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:
            print(f"Erro ao processar PDF: {str(e)}")
            return None

    def generate_response(self, file_obj, query: str) -> str:
        """Gera resposta para a consulta"""
        if not query.strip():
            return "Por favor, insira uma pergunta."
        
        try:
            # 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(
                llm=self.llm,
                chain_type="stuff",
                retriever=db.as_retriever(
                    search_kwargs={
                        "k": 4,  # Aumentamos o k para ter mais contexto
                        "fetch_k": 6
                    }
                ),
                return_source_documents=True
            )
            
            # Adiciona contexto sobre a fonte da 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."""
            
            # Gera resposta
            result = qa_chain({"query": prompt})
            return result["result"]
            
        except Exception as e:
            return f"Erro ao gerar resposta: {str(e)}"

# Interface Gradio
def create_demo():
    rag = RAGSystem()
    
    with gr.Blocks() as demo:
        gr.Markdown("# 📚 Sistema RAG de Consulta a Documentos")
        gr.Markdown(f"""
        ### Como usar:
        1. Os documentos da pasta `{DOCS_DIR}` são usados como base de conhecimento
        2. Você pode fazer upload de PDFs adicionais para consulta
        3. Digite sua pergunta e aguarde a resposta
        4. As respostas são baseadas no conteúdo dos documentos
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload do PDF (opcional)",
                    type="binary",
                    file_types=[".pdf"]
                )
                query_input = gr.Textbox(
                    label="Sua Pergunta",
                    placeholder="Digite sua pergunta sobre o documento...",
                    lines=3
                )
                submit_btn = gr.Button("🔍 Buscar Resposta", variant="primary")
            
            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Resposta",
                    lines=10
                )
        
        submit_btn.click(
            fn=rag.generate_response,
            inputs=[file_input, query_input],
            outputs=output
        )
        
        gr.Examples(
            examples=[
                [None, "Qual é o tema principal dos documentos?"],
                [None, "Pode fazer um resumo dos pontos principais?"],
                [None, "Quais são as principais conclusões?"]
            ],
            inputs=[file_input, query_input]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_demo()
    demo.launch()