RAG: Retrieval Augmented Generation

from langchain.embeddings import OpenAIEmbeddings from langchain.vectorstores import Chroma from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.chains import RetrievalQA from langchain.llms import OpenAI from langchain.document_loaders import DirectoryLoader class RAGSystem: def __init__(self, docs_path, model_name="gpt-4"): """ Sistema RAG completo Args: docs_path: Diretório com documentos model_name: Modelo LLM (gpt-4, gpt-3.5-turbo) """ self.embeddings = OpenAIEmbeddings() self.llm = OpenAI(model_name=model_name, temperature=0) self.vectorstore = None self.qa_chain = None # Carrega e processa documentos self._load_documents(docs_path) def _load_documents(self, docs_path): """Carrega documentos e cria vector store""" # Carrega documentos de múltiplos formatos loader = DirectoryLoader( docs_path, glob="**/*.{txt,pdf,md}", show_progress=True ) documents = loader.load() # Split em chunks otimizados text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=200, separators=["\n\n", "\n", ". ", " ", ""] ) splits = text_splitter.split_documents(documents) # Cria vector store com ChromaDB self.vectorstore = Chroma.from_documents( documents=splits, embedding=self.embeddings, persist_directory="./chroma_db" ) # Configura retrieval chain self.qa_chain = RetrievalQA.from_chain_type( llm=self.llm, chain_type="stuff", retriever=self.vectorstore.as_retriever( search_type="mmr", # Maximum Marginal Relevance search_kwargs={"k": 5, "fetch_k": 20} ), return_source_documents=True ) def query(self, question): """ Faz pergunta ao sistema RAG Returns: dict com 'answer' e 'sources' """ result = self.qa_chain({"query": question}) return { "answer": result["result"], "sources": [ { "content": doc.page_content, "metadata": doc.metadata } for doc in result["source_documents"] ] } def add_documents(self, new_docs): """Adiciona novos documentos ao vector store""" self.vectorstore.add_documents(new_docs) self.vectorstore.persist() # GraphRAG: RAG com Knowledge Graph class GraphRAG: """Implementação simplificada de GraphRAG""" def __init__(self, documents): from langchain.graphs import Neo4jGraph from langchain.chains import GraphCypherQAChain # Conecta ao Neo4j self.graph = Neo4jGraph( url="bolt://localhost:7687", username="neo4j", password="password" ) # Extrai entidades e relações self._build_knowledge_graph(documents) def _build_knowledge_graph(self, documents): """Constrói knowledge graph de documentos""" # Usa NER para extrair entidades # Identifica relações entre entidades # Popula Neo4j com nodes e edges pass def query_graph(self, question): """Query usando Cypher sobre graph""" chain = GraphCypherQAChain.from_llm( llm=self.llm, graph=self.graph, verbose=True ) return chain.run(question) # HyDE: Hypothetical Document Embeddings class HyDERetriever: """Gera documentos hipotéticos para melhor retrieval""" def __init__(self, vectorstore, llm): self.vectorstore = vectorstore self.llm = llm def retrieve(self, query, n_docs=5): """ 1. LLM gera documento hipotético que responderia query 2. Usa embedding desse documento para buscar similares 3. Retorna documentos reais mais relevantes """ # Gera documento hipotético prompt = f"""Escreva um parágrafo que responderia: {query} Parágrafo:""" hypothetical_doc = self.llm(prompt) # Busca por similaridade com doc hipotético docs = self.vectorstore.similarity_search( hypothetical_doc, k=n_docs ) return docs # Uso if __name__ == "__main__": # Sistema RAG básico rag = RAGSystem("./knowledge_base") result = rag.query( "Quais são as políticas de reembolso da empresa?" ) print("Resposta:", result["answer"]) print("\nFontes:") for i, source in enumerate(result["sources"], 1): print(f"{i}. {source['metadata']['source']}") # HyDE para queries complexas hyde = HyDERetriever(rag.vectorstore, rag.llm) docs = hyde.retrieve("Como otimizar latência em produção?") print(f"\nEncontrados {len(docs)} documentos relevantes")