app.py

import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForCausalLM
from pypdf import PdfReader
import os
import re

# --- 1. Carregamento dos Modelos ---
print("A carregar o modelo de recuperação (Sentence Transformer)...")
retriever_model = SentenceTransformer('all-MiniLM-L6-v2')

print("A carregar o modelo de geração (DeepSeek)...")
generator_tokenizer = AutoTokenizer.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
generator_model = AutoModelForCausalLM.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
print("Modelos carregados com sucesso!")

# --- 2. Função de Processamento de Ficheiros (Chunking Estruturado) ---
def process_files(files):
    if not files:
        return None, "Por favor, envie um ou mais ficheiros."
    
    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\n"
            except Exception as e:
                return None, f"Erro ao ler o ficheiro 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\n"
            except Exception as e:
                return None, f"Erro ao ler o ficheiro TXT {os.path.basename(file_path)}: {e}"

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

    chunk_pattern = r"(?m)(^\d+\..*)"
    text_chunks = [chunk.strip() for chunk in re.split(chunk_pattern, knowledge_text) if chunk.strip()]
    
    structured_chunks = []
    i = 0
    while i < len(text_chunks):
        if re.match(chunk_pattern, text_chunks[i]) and i + 1 < len(text_chunks):
            structured_chunks.append(text_chunks[i] + "\n" + text_chunks[i+1])
            i += 2
        else:
            structured_chunks.append(text_chunks[i])
            i += 1

    if not structured_chunks:
        return None, "O texto extraído não continha blocos de texto válidos para processamento."

    print(f"A processar {len(structured_chunks)} chunks estruturados dos ficheiros...")
    knowledge_base_embeddings = retriever_model.encode(structured_chunks, convert_to_tensor=True, show_progress_bar=True)
    print("Base de conhecimento criada a partir dos ficheiros.")
    
    return (structured_chunks, knowledge_base_embeddings), f"✅ Sucesso! {len(files)} ficheiro(s) processado(s), gerando {len(structured_chunks)} chunks estruturados."


# --- 3. O CÉREBRO DA ANÁLISE: generate_compliance_report ---
def generate_compliance_report(task, knowledge_state, progress=gr.Progress(track_tqdm=True)):
    """
    Esta função orquestra todo o processo de RAG: recupera um contexto amplo e usa um único
    'super-prompt' para gerar o relatório de conformidade completo de uma só vez.
    """
    if not task:
        return "Por favor, forneça uma tarefa de análise."
    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 ficheiros primeiro."

    knowledge_base, knowledge_base_embeddings = knowledge_state
    
    progress(0, desc="A recuperar contexto relevante...")
    
    # Passo 1: Recuperação Ampla do Contexto
    # Usamos uma pergunta genérica para recuperar os chunks mais relevantes do documento inteiro.
    # Aumentamos o top_k para dar ao modelo uma visão mais completa.
    search_query = "Informações completas do certificado de calibração"
    question_embedding = retriever_model.encode(search_query, convert_to_tensor=True)
    cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
    top_k = min(15, len(knowledge_base)) # Aumentado para 15 para um contexto muito mais rico
    top_results = torch.topk(cosine_scores, k=top_k, dim=-1)
    retrieved_context = "\n\n---\n\n".join([knowledge_base[i] for i in top_results.indices[0]])

    progress(0.5, desc="A gerar o relatório de conformidade...")

    # Passo 2: Geração com "Super-Prompt"
    # Este prompt contém a checklist completa e instrui o modelo a preenchê-la.
    final_prompt = f"""### Instruction:
Você é um auditor de metrologia a preencher um relatório de conformidade. Com base no 'Contexto do Documento' fornecido, preencha cada item da 'Checklist de Análise' abaixo. Se uma informação não for encontrada no contexto, escreva 'Não encontrado'.

**Contexto do Documento:**
{retrieved_context}

**Checklist de Análise:**
# Relatório de Análise de Conformidade

## 1. Incerteza de Medição
   - **1.1 a 1.2 Casas Decimais / Compatibilidade:**
   - **1.3 Nível de Confiança, Fator de Abrangência (k) e Graus de Liberdade:**
   - **1.4 Declaração de Rastreabilidade dos Resultados:**

## 2. Resultados da Calibração
   - **2.1 a 2.3 Unidades SI, Casas Decimais e Algarismos Significativos:**

## 3. Conformidade da Faixa
   - **3.1 Faixa e Especificações Solicitadas:**

## 4. Condições Ambientais
   - **4.1 Registro das Condições e Incerteza Associada:**

## 5. Identificação do Item
   - **5.1 Descrição e Identificação do Item:**

## 6. Identificação do Método
   - **6.1 Método/Procedimento Utilizado:**

## 7. Identificação do Cliente
   - **7.1 Nome e Endereço do Cliente:**

## 8. Identificação do Laboratório
   - **8.1 Nome e Endereço do Laboratório:**

## 9. Identificação do Certificado
   - **9.1 Número do Certificado:**

## 10. Autorização
   - **10.1 Pessoas Autorizadas:**

### Response:
"""

    input_ids = generator_tokenizer(final_prompt, return_tensors="pt").input_ids
    outputs = generator_model.generate(
        input_ids, 
        max_new_tokens=1024, # Aumentado para relatórios detalhados
        do_sample=False, 
        pad_token_id=generator_tokenizer.eos_token_id
    )
    final_report = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Limpa a resposta para remover o prompt inicial
    if "### Response:" in final_report:
        final_report = final_report.split("### Response:")[1].strip()

    progress(1, desc="Análise concluída.")
    return final_report


# --- 4. Interface Gráfica (Simplificada para o novo fluxo) ---
with gr.Blocks(theme=gr.themes.Soft()) as interface:
    knowledge_state = gr.State()
    gr.Markdown(
        """
        # 🤖 Agente de Análise de Conformidade Metrológica (v12 - Robusto)
        **1. Carregue um documento**: Envie um certificado de calibração (`.pdf` ou `.txt`).
        **2. Processe o documento**: Clique no botão para criar a base de conhecimento.
        **3. Inicie a Análise**: Dê uma tarefa ao agente (ex: "Analisar conformidade deste certificado") e clique em "Iniciar Análise".
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_uploader = gr.File(label="Carregar Documento (.pdf, .txt)", file_count="multiple", file_types=[".pdf", ".txt"])
            process_button = gr.Button("Processar Documento", variant="primary")
            status_box = gr.Textbox(label="Status do Processamento", interactive=False)
        
        with gr.Column(scale=2):
            task_box = gr.Textbox(label="Tarefa de Análise", placeholder='Ex: Avaliar a conformidade deste certificado de calibração.')
            submit_button = gr.Button("Iniciar Análise", variant="primary")

    with gr.Row():
        report_box = gr.Markdown(label="Relatório Final de Análise")

    process_button.click(fn=process_files, inputs=[file_uploader], outputs=[knowledge_state, status_box])
    submit_button.click(fn=generate_compliance_report, inputs=[task_box, knowledge_state], outputs=[report_box])

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