app.py

import os
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain_community.llms import HuggingFaceHub
from typing import Optional
import tempfile

# Configurações
EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
LLM_REPO_ID = "google/flan-t5-large"

def create_temporary_file(file_content: bytes) -> str:
    """Cria um arquivo temporário a partir dos bytes do arquivo."""
    try:
        temp_dir = tempfile.mkdtemp()
        temp_path = os.path.join(temp_dir, "temp.pdf")
        
        with open(temp_path, 'wb') as f:
            f.write(file_content)
        
        return temp_path
    except Exception as e:
        raise Exception(f"Erro ao criar arquivo temporário: {str(e)}")

def load_and_process_pdf(pdf_path: str) -> Optional[FAISS]:
    """
    Carrega e processa o PDF, com tratamento de erros adequado.
    """
    try:
        # Carrega o PDF
        loader = PyPDFLoader(pdf_path)
        documents = loader.load()
        
        if not documents:
            return None
            
        # Divide o texto em chunks
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len
        )
        texts = text_splitter.split_documents(documents)
        
        # Cria embeddings e armazena no vetor store
        embeddings = HuggingFaceEmbeddings(
            model_name=EMBEDDING_MODEL,
            model_kwargs={'device': 'cpu'}
        )
        db = FAISS.from_documents(texts, embeddings)
        
        return db
        
    except Exception as e:
        print(f"Erro ao processar o PDF: {str(e)}")
        return None

def generate_response(file_obj, query: str) -> str:
    """
    Gera resposta para a consulta usando RAG, com tratamento de erros.
    """
    if file_obj is None:
        return "Erro: Nenhum arquivo PDF foi carregado."
    
    if not query.strip():
        return "Erro: Por favor, insira uma pergunta."
    
    try:
        # Cria arquivo temporário e processa o PDF
        temp_path = create_temporary_file(file_obj)
        db = load_and_process_pdf(temp_path)
        
        if db is None:
            return "Erro: Não foi possível processar o PDF."
        
        # Configura o modelo de linguagem
        llm = HuggingFaceHub(
            repo_id=LLM_REPO_ID,
            huggingfacehub_api_token=os.environ.get("HUGGINGFACE_API_TOKEN"),
            model_kwargs={
                "temperature": 0.7,
                "max_length": 512,
                "top_p": 0.95
            }
        )
        
        # Cria a cadeia de RAG
        qa_chain = RetrievalQA.from_chain_type(
            llm=llm,
            chain_type="stuff",
            retriever=db.as_retriever(search_kwargs={"k": 3}),
            return_source_documents=True,
            verbose=True
        )
        
        # Executa a consulta
        result = qa_chain({"query": query})
        
        # Limpa arquivos temporários
        os.remove(temp_path)
        os.rmdir(os.path.dirname(temp_path))
        
        return result["result"]
        
    except Exception as e:
        return f"Erro ao gerar resposta: {str(e)}"

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Sistema de RAG com LangChain")
    gr.Markdown("Faça upload de um PDF e faça perguntas sobre o conteúdo.")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(
                label="Upload PDF",
                type="binary",
                file_types=[".pdf"]
            )
            query_input = gr.Textbox(
                label="Sua Pergunta",
                placeholder="Digite sua pergunta aqui...",
                lines=3
            )
            submit_btn = gr.Button("Enviar Pergunta")
        
        with gr.Column():
            output = gr.Textbox(
                label="Resposta Gerada",
                lines=10
            )
    
    submit_btn.click(
        fn=generate_response,
        inputs=[file_input, query_input],
        outputs=output
    )
    
    gr.Examples(
        examples=[
            [None, "Qual é o principal tema deste documento?"],
            [None, "Pode resumir os pontos principais?"]
        ],
        inputs=[file_input, query_input]
    )

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