app.py

import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
from transformers import T5ForConditionalGeneration, T5Tokenizer
from pypdf import PdfReader
import os

# --- 1. Carregamento dos Modelos (sem alterações) ---
print("Carregando o modelo de recuperação (Sentence Transformer)...")
retriever_model = SentenceTransformer('all-MiniLM-L6-v2')
print("Carregando o modelo de geração (Flan-T5)...")
generator_tokenizer = T5Tokenizer.from_pretrained('google/flan-t5-base')
generator_model = T5ForConditionalGeneration.from_pretrained('google/flan-t5-base')
print("Modelos carregados com sucesso!")

# --- 2. Função para Processar Arquivos Enviados ---
def process_files(files):
    if not files:
        return None, "Por favor, envie um ou mais arquivos."

    knowledge_text = ""
    for file in files:
        file_path = file.name
        if file_path.endswith(".pdf"):
            try:
                reader = PdfReader(file_path)
                for page in reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        knowledge_text += page_text + "\n"
            except Exception as e:
                return None, f"Erro ao ler o arquivo PDF {os.path.basename(file_path)}: {e}"
        elif file_path.endswith(".txt"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    knowledge_text += f.read() + "\n"
            except Exception as e:
                return None, f"Erro ao ler o arquivo TXT {os.path.basename(file_path)}: {e}"

    if not knowledge_text.strip():
        return None, "Não foi possível extrair texto dos arquivos fornecidos."

    # MUDANÇA 1: Melhorando o "Chunking" (Divisão do Texto)
    # Em vez de dividir por parágrafos (\n\n), dividimos por linha (\n).
    # Isso cria chunks menores e mais focados, o que é melhor para documentos técnicos.
    text_chunks = [chunk.strip() for chunk in knowledge_text.split('\n') if chunk.strip() and len(chunk) > 10]
    
    if not text_chunks:
        return None, "O texto extraído não continha blocos de texto válidos para processamento."

    print(f"Processando {len(text_chunks)} blocos de texto dos arquivos...")
    knowledge_base_embeddings = retriever_model.encode(text_chunks, convert_to_tensor=True, show_progress_bar=True)
    print("Base de conhecimento criada a partir dos arquivos.")
    
    return (text_chunks, knowledge_base_embeddings), f"✅ Sucesso! {len(files)} arquivo(s) processado(s), gerando {len(text_chunks)} blocos de texto."


# --- 3. A Função Principal do RAG (com melhorias) ---
def answer_question(question, knowledge_state):
    if not question:
        return "Por favor, insira uma pergunta."
        
    if not knowledge_state or not knowledge_state[0] or knowledge_state[1] is None:
        return "⚠️ A base de conhecimento está vazia. Por favor, processe alguns arquivos primeiro."

    knowledge_base, knowledge_base_embeddings = knowledge_state

    # Etapa de Recuperação (Retrieval)
    question_embedding = retriever_model.encode(question, convert_to_tensor=True)
    
    # MUDANÇA 2: Recuperando mais contexto (top_k=3)
    # Em vez de pegar apenas o melhor chunk, pegamos os 3 melhores.
    # Isso dá mais informação para a IA gerar uma resposta completa.
    cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
    top_k = min(3, len(knowledge_base)) # Garante que não tentamos pegar mais chunks do que existem
    top_results = torch.topk(cosine_scores, k=top_k, dim=-1)

    # Junta os 3 melhores chunks em um único contexto
    retrieved_context = "\n---\n".join([knowledge_base[i] for i in top_results.indices[0]])

    print(f"\n--- Nova Pergunta de Auditoria ---")
    print(f"Pergunta: {question}")
    print(f"Contexto Recuperado (Top {top_k}):\n{retrieved_context}")

    # Etapa de Geração (Generation)
    
    # MUDANÇA 3: Prompt de Geração Aprimorado
    # Damos instruções mais claras para que a IA elabore a resposta.
    prompt = f"""
    Você é um assistente de auditoria especialista. Sua tarefa é responder à pergunta do usuário de forma clara e concisa, baseando-se exclusivamente nos trechos de documentos fornecidos abaixo. Elabore uma resposta completa em vez de simplesmente copiar o texto.

    **Documentos Relevantes:**
    {retrieved_context}

    **Pergunta do Auditor:**
    {question}

    **Resposta Elaborada:**
    """

    input_ids = generator_tokenizer(prompt, return_tensors="pt").input_ids
    outputs = generator_model.generate(
        input_ids,
        max_length=256,
        num_beams=5,
        early_stopping=True
    )

    answer = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

# --- 4. Interface Gráfica (sem alterações na estrutura) ---
with gr.Blocks(theme=gr.themes.Soft()) as interface:
    knowledge_state = gr.State()
    gr.Markdown(
        """
        # 🤖 RAG - Auditor de Documentos (v3)
        **1. Carregue seus arquivos**: Envie um ou mais certificados ou documentos nos formatos `.pdf` ou `.txt`.
        **2. Processe os arquivos**: Clique no botão para criar a base de conhecimento.
        **3. Faça perguntas**: Após o processamento, faça perguntas sobre o conteúdo dos documentos.
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            file_uploader = gr.File(label="Carregar Certificados (.pdf, .txt)", file_count="multiple", file_types=[".pdf", ".txt"])
            process_button = gr.Button("Processar Arquivos", variant="primary")
            status_box = gr.Textbox(label="Status do Processamento", interactive=False)
        with gr.Column(scale=2):
            question_box = gr.Textbox(label="Faça sua pergunta aqui", placeholder="Ex: Qual o resultado da calibração do instrumento PI-101?")
            submit_button = gr.Button("Obter Resposta", variant="primary")
            answer_box = gr.Textbox(label="Resposta Baseada nos Documentos", interactive=False, lines=5)

    process_button.click(fn=process_files, inputs=[file_uploader], outputs=[knowledge_state, status_box])
    submit_button.click(fn=answer_question, inputs=[question_box, knowledge_state], outputs=[answer_box])

# --- 5. Lançamento do App ---
if __name__ == "__main__":
    interface.launch()