Como Implementar um Sistema RAG do Zero em Python

Disclaimer

Este texto foi inicialmente concebido pela IA Generativa em função da transcrição de um vídeo do canal de Daniel Romero(a pessoa que lidera nossa especialização em Engenharia de IA). Se preferir acompanhar por vídeo, é só dar o play.

Introdução

RAG (Retrieval-Augmented Generation) é uma técnica que combina busca por similaridade com geração de texto por LLMs para criar respostas mais precisas e contextualizadas. Neste post, vamos implementar um sistema RAG simples do zero, usando apenas três bibliotecas: NumPy, Sentence Transformers e Grok.

O que vamos construir

Nosso sistema RAG terá três componentes principais: geração de embeddings para os documentos, busca por similaridade usando similaridade do cosseno e geração de resposta usando um LLM.

Configuração inicial

Para este exemplo, precisamos instalar apenas três bibliotecas:

pip install numpy sentence-transformers grok

Vamos trabalhar com o básico, sem banco vetorial por enquanto. NumPy será usado para os cálculos, Sentence Transformers para os embeddings e Grok para o LLM.

Começamos com os imports:

import numpy as np
from sentence_transformers import SentenceTransformer
from grok import Grok

Carregando o modelo e gerando embeddings

O modelo que vamos utilizar é o all-MiniLM-L6-v2. Ele é pequeno, rápido e muito eficiente. Vale destacar que esse modelo funciona bem apenas para o idioma inglês, então para produção com conteúdo em português seria necessário trocar de modelo.

Um detalhe importante: esses modelos de embeddings são treinados olhando para milhões de pares de textos. Alguns pares são positivos, como duas frases que querem dizer a mesma coisa, ou uma pergunta com a resposta certa. Outros pares são negativos, com duas frases sem relação nenhuma. O treino funciona basicamente dizendo para o modelo: "esses dois aqui têm que ficar perto, e esses outros aqui têm que ficar distantes". Depois de repetir isso milhares de vezes, o modelo aprende a transformar qualquer texto em um vetor de números que preserva o sentido.

model = SentenceTransformer('all-MiniLM-L6-v2')
client = Grok()

# Seus documentos (fonte de conhecimento)
documents = [
    # Lista de documentos sobre machine learning
]

# Gerar embeddings de todos os documentos
doc_embeddings = model.encode(documents)

Com essa única linha, todos os documentos se transformam em vetores numéricos. No caso do all-MiniLM-L6-v2, cada documento tem 384 dimensões do tipo float32. Como não temos um banco vetorial neste momento, esses vetores ficam guardados na memória.

Entendendo a similaridade do cosseno

Uma das partes mais importantes é como medir a distância entre dois textos. Para resolver isso, implementamos a similaridade do cosseno.

A fórmula da similaridade do cosseno é expressa como:

similarity(A, B) = cos(θ) = (A · B) / (||A|| × ||B||)

O numerador A · B representa o produto escalar dos vetores A e B, calculado multiplicando cada elemento de A pelo elemento correspondente de B e somando todos esses produtos. No denominador, temos o produto das normas euclidianas dos dois vetores.

Geometricamente, essa fórmula calcula o cosseno do ângulo θ entre os dois vetores no espaço n-dimensional. Quando θ é igual a zero graus, os vetores apontam na mesma direção e o cosseno vale 1, indicando máxima similaridade. Quando θ é igual a 90 graus, os vetores são perpendiculares e o cosseno vale zero, indicando nenhuma similaridade.

Veja um exemplo prático:

v1 = np.array([1, 2, 3])
v2 = np.array([4, 5, 6])

# Produto escalar
dot_product = np.dot(v1, v2)  # Resultado: 32

# Magnitude (norma euclidiana)
magnitude = np.linalg.norm(v1) * np.linalg.norm(v2)

# Similaridade do cosseno
similarity = dot_product / magnitude  # Resultado: 0.97 (97% de similaridade)

Essa etapa final mede o quanto os vetores v1 e v2 estão na mesma direção.

Implementando as funções do RAG

Função de similaridade do cosseno

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

Função de recuperação de documentos

def retrieve(query, top_k=3):
    # Gera o embedding da query
    query_embedding = model.encode(query)

    # Calcula similaridade com todos os documentos
    similarities = []
    for i, doc_embedding in enumerate(doc_embeddings):
        sim = cosine_similarity(query_embedding, doc_embedding)
        similarities.append((i, sim))

    # Ordena pela similaridade (maior primeiro)
    similarities.sort(key=lambda x: x[1], reverse=True)

    # Retorna os top_k documentos mais relevantes
    return [(documents[idx], sim) for idx, sim in similarities[:top_k]]

Função de geração de resposta

def generate_answer(query, retrieved_docs):
    # Junta os documentos em um contexto único
    context = "\n".join([doc for doc, _ in retrieved_docs])

    response = client.chat.completions.create(
        model="llama-3.1-8b-instant",
        messages=[
            {
                "role": "system",
                "content": "Você é um especialista em machine learning. Use apenas o contexto fornecido para responder as perguntas."
            },
            {
                "role": "user",
                "content": f"Contexto fornecido:\n{context}\n\nPergunta do usuário: {query}"
            }
        ],
        temperature=0
    )

    return response.choices[0].message.content

A temperatura zero garante que as respostas sejam mais determinísticas e focadas no contexto.

Função principal RAG

def rag(query, top_k=3):
    # Recupera os documentos mais relevantes
    retrieved = retrieve(query, top_k)

    # Gera a resposta com o LLM
    answer = generate_answer(query, retrieved)

    return answer, retrieved

Testando o sistema

answer, docs = rag("O que é machine learning?")

print(answer)

# Mostra os documentos recuperados com suas similaridades
for doc, similarity in docs:
    print(f"- {similarity:.3f}: {doc}")

Próximos passos

Em projetos mais complexos, temos ainda mais desafios. Teríamos um banco de dados vetorial, estratégia de extração de conteúdo, chunking, filtros, metadados e muito mais. Mas essa é a base, e agora você entende como o RAG funciona na prática.

Especialização em Engenharia de IA

Este conteúdo é um trecho de uma das aulas da Especialização em Engenharia de IA, uma parceria com a Dev Mais Eficiente. O curso aborda RAG, Vector Search, Busca Híbrida, Agents, Tools e muito mais, sempre com aulas 100% práticas e com exemplos reais.

Faça sua inscrição em https://deveficiente.com/especializacao-engenharia-ia .