Create metrology_rag.py

metrology_rag.py

@@ -0,0 +1,254 @@
+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)