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, AutoModelForCausalLM, pipeline
import torch
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):
        # Inicializa o modelo de linguagem
        self.tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
        self.model = AutoModelForCausalLM.from_pretrained(
            LLM_MODEL,
            torch_dtype=torch.float16,
            device_map="auto",
            load_in_8bit=True
        )
        
        # Configura o pipeline
        pipe = pipeline(
            "text-generation",
            model=self.model,
            tokenizer=self.tokenizer,
            max_length=2048,
            temperature=0.7,
            top_p=0.95,
            repetition_penalty=1.15
        )
        
        # Configura o modelo LangChain
        self.llm = HuggingFacePipeline(pipeline=pipe)
        
        # Configura 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=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 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=1000,
                chunk_overlap=200,
                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": 3,
                        "fetch_k": 5
                    }
                ),
                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": 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 com Base de Conhecimento")
        gr.Markdown(f"""
        ### Instruções:
        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 (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("🔍 Pesquisar", 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()