metrology_rag.py

import google.generativeai as genai
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Dict, Optional
import faiss
import pickle
import os
from datetime import datetime
import pdfplumber
from pathlib import Path

# Configuração inicial do Gemini
genai.configure(api_key="AIzaSyClWplmEF8_sDgmSbhg0h6xkAoFQcLU4p4")

class MetrologyGlossary:
    """Glossário interno de termos metrológicos para melhorar recuperação e respostas."""
    def __init__(self):
        self.terms = {
            "incerteza": "Medida da dispersão associada ao resultado de uma medição.",
            "calibração": "Comparação de um instrumento com um padrão de referência.",
            "traceability": "Propriedade de um resultado de medição que pode ser relacionado a um padrão nacional ou internacional.",
            "iso/iec 17025": "Norma internacional para laboratórios de ensaio e calibração.",
            # Adicione mais termos conforme necessário
        }

    def enhance_query(self, query: str) -> str:
        """Adiciona definições ou sinônimos à consulta para maior precisão."""
        for term, definition in self.terms.items():
            if term.lower() in query.lower():
                query += f" ({definition})"
        return query

class DocumentParser:
    """Extrai texto e tabelas de arquivos PDF, com foco em metrologia."""
    def parse_pdf(self, file_path: str) -> Dict:
        """Extrai texto e tabelas de um único PDF."""
        try:
            with pdfplumber.open(file_path) as pdf:
                text = ""
                tables = []
                for page in pdf.pages:
                    # Extrai texto
                    page_text = page.extract_text() or ""
                    text += page_text + "\n"
                    # Extrai tabelas
                    page_tables = page.extract_tables()
                    for table in page_tables:
                        tables.append(table)

                # Converte tabelas em texto estruturado
                table_text = ""
                for idx, table in enumerate(tables):
                    table_text += f"Tabela {idx + 1}:\n"
                    for row in table:
                        table_text += " | ".join([str(cell) or "" for cell in row]) + "\n"

                return {
                    "content": (text + "\n" + table_text).strip(),
                    "metadata": {
                        "file_name": os.path.basename(file_path),
                        "path": file_path,
                        "num_pages": len(pdf.pages),
                        "has_tables": len(tables) > 0
                    }
                }
        except Exception as e:
            print(f"Erro ao processar {file_path}: {str(e)}")
            return {"content": "", "metadata": {}}

    def parse_multiple_pdfs(self, pdf_paths: List[str]) -> List[Dict]:
        """Extrai texto e tabelas de múltiplos PDFs."""
        documents = []
        for path in pdf_paths:
            doc = self.parse_pdf(path)
            if doc["content"]:
                documents.append(doc)
        return documents

class KnowledgeBase:
    """Gerencia a base de conhecimento metrológico."""
    def __init__(self):
        self.documents: List[Dict] = []

    def add_document(self, content: str, metadata: Optional[Dict] = None):
        doc = {"content": content, "metadata": metadata or {}, "id": len(self.documents)}
        self.documents.append(doc)

    def add_documents(self, documents: List[Dict]):
        for doc in documents:
            self.add_document(doc["content"], doc["metadata"])

    def get_document(self, doc_id: int) -> Dict:
        return self.documents[doc_id]

    def get_all_contents(self) -> List[str]:
        return [doc["content"] for doc in self.documents]

class EmbeddingGenerator:
    """Gera embeddings para textos."""
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)

    def generate(self, texts: List[str]) -> np.ndarray:
        return self.model.encode(texts, convert_to_numpy=True)

class VectorStore:
    """Armazena e busca embeddings usando FAISS."""
    def __init__(self, dimension: int):
        self.index = faiss.IndexFlatL2(dimension)
        self.doc_ids = []

    def add_vectors(self, embeddings: np.ndarray, doc_ids: List[int]):
        self.index.add(embeddings)
        self.doc_ids.extend(doc_ids)

    def search(self, query_embedding: np.ndarray, k: int = 5) -> List[int]:
        distances, indices = self.index.search(query_embedding, k)
        return [self.doc_ids[idx] for idx in indices[0]]

class Retriever:
    """Recupera documentos relevantes para uma consulta."""
    def __init__(self, knowledge_base: KnowledgeBase, vector_store: VectorStore, embedding_generator: EmbeddingGenerator):
        self.knowledge_base = knowledge_base
        self.vector_store = vector_store
        self.embedding_generator = embedding_generator

    def retrieve(self, query: str, k: int = 5) -> List[Dict]:
        query_embedding = self.embedding_generator.generate([query])
        doc_ids = self.vector_store.search(query_embedding, k)
        return [self.knowledge_base.get_document(doc_id) for doc_id in doc_ids]

class ResponseGenerator:
    """Gera respostas técnicas para perguntas metrológicas."""
    def __init__(self, model_name: str = "gemini-2.0-flash-thinking-exp-1219"):
        self.model = genai.GenerativeModel(model_name)

    def generate(self, query: str, retrieved_docs: List[Dict]) -> str:
        context = "\n".join([doc["content"] for doc in retrieved_docs])
        prompt = (
            "Você é um especialista em metrologia, com conhecimento em normas como ISO/IEC 17025, incerteza de medição, "
            "calibração e rastreabilidade. Com base no contexto fornecido, responda à pergunta de forma técnica, precisa e clara:\n\n"
            f"Contexto:\n{context}\n\nPergunta: {query}\n\nResposta:"
        )
        try:
            response = self.model.generate_content(prompt)
            return response.text if response else "Desculpe, não consegui gerar uma resposta."
        except Exception as e:
            return f"Erro ao gerar resposta: {str(e)}"

class CacheManager:
    """Gerencia cache de respostas."""
    def __init__(self, cache_file: str = "metrology_cache.pkl"):
        self.cache_file = cache_file
        self.cache = self._load_cache()

    def _load_cache(self) -> Dict:
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return pickle.load(f)
        return {}

    def _save_cache(self):
        with open(self.cache_file, "wb") as f:
            pickle.dump(self.cache, f)

    def get(self, query: str) -> Optional[str]:
        return self.cache.get(query)

    def set(self, query: str, response: str):
        self.cache[query] = {"response": response, "timestamp": datetime.now()}
        self._save_cache()

class QueryProcessor:
    """Pré-processa consultas com foco em metrologia."""
    def __init__(self):
        self.glossary = MetrologyGlossary()

    def process(self, query: str) -> str:
        query = query.strip().lower()
        return self.glossary.enhance_query(query)

class MetrologyRAGPipeline:
    """Orquestra o agente de metrologia avançado."""
    def __init__(self):
        self.knowledge_base = KnowledgeBase()
        self.embedding_generator = EmbeddingGenerator()
        self.vector_store = VectorStore(dimension=384)
        self.retriever = Retriever(self.knowledge_base, self.vector_store, self.embedding_generator)
        self.response_generator = ResponseGenerator()
        self.cache_manager = CacheManager()
        self.query_processor = QueryProcessor()
        self.document_parser = DocumentParser()

    def load_pdfs(self, pdf_paths: List[str] = None, pdf_folder: Optional[str] = None):
        """Carrega N arquivos PDF de uma lista de caminhos ou pasta."""
        if pdf_paths and pdf_folder:
            raise ValueError("Forneça apenas pdf_paths ou pdf_folder, não ambos.")

        if pdf_folder:
            pdf_paths = [str(p) for p in Path(pdf_folder).glob("*.pdf")]

        if not pdf_paths:
            print("Nenhum arquivo PDF fornecido ou encontrado.")
            return

        print(f"Carregando {len(pdf_paths)} arquivos PDF...")
        documents = self.document_parser.parse_multiple_pdfs(pdf_paths)
        if documents:
            self.knowledge_base.add_documents(documents)
            self._index_documents()
            print(f"{len(documents)} documentos indexados com sucesso.")
        else:
            print("Nenhum documento válido foi extraído dos PDFs.")

    def _index_documents(self):
        contents = self.knowledge_base.get_all_contents()
        if not contents:
            return
        embeddings = self.embedding_generator.generate(contents)
        doc_ids = list(range(len(contents)))
        self.vector_store.add_vectors(embeddings, doc_ids)

    def query(self, query: str, k: int = 5) -> str:
        processed_query = self.query_processor.process(query)
        cached_response = self.cache_manager.get(processed_query)
        if cached_response:
            return f"[Resposta do cache] {cached_response}"

        retrieved_docs = self.retriever.retrieve(processed_query, k)
        response = self.response_generator.generate(processed_query, retrieved_docs)
        self.cache_manager.set(processed_query, response)
        return response

# # Exemplo de uso
if __name__ == "__main__":
    # Inicializa o pipeline
    rag = MetrologyRAGPipeline()

    # Carrega N arquivos PDF de uma pasta
    pdf_folder = "/content/"  # Substitua pelo caminho real
    rag.load_pdfs(pdf_folder=pdf_folder)

    # Alternativamente, carrega PDFs específicos
    pdf_paths = [
        # "caminho/para/manual_calibrador.pdf",
        # "caminho/para/iso_17025.pdf",
        # Adicione mais caminhos
    ]
    # rag.load_pdfs(pdf_paths=pdf_paths)

    # Faz uma consulta técnica
    pergunta = "faça uma avaliação sobre o documento CERTIFICADO DE CALIBRAÇÃO N RBC 25/0018"
    resposta = rag.query(pergunta)
    print("Agente de Metrologia:", resposta)