Configurar RAG

Guia completo para configurar Retrieval-Augmented Generation. Vector store, embeddings, indexação e busca híbrida.

Visão Geral

RAG (Retrieval-Augmented Generation) combina busca semântica com geração de texto para produzir respostas fundamentadas em documentos específicos. Este guia mostra como configurar um pipeline RAG completo usando a Polar.

Arquitetura

Documentos → Chunking → Embeddings (urso-embed) → Vector Store
                                                        ↓
Consulta → Embedding → Busca Híbrida (BM25 + Semântica) → Reranking → LLM (urso-mabe)

Passo 1: Escolher Vector Store

A Polar suporta dois vector stores:

Qdrant (Recomendado)

Otimizado para busca vetorial de alta performance.

from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance

client = QdrantClient(host="localhost", port=6333)

# Criar coleção
client.create_collection(
    collection_name="meus_documentos",
    vectors_config=VectorParams(
        size=1024,  # Dimensão do urso-embed
        distance=Distance.COSINE
    )
)

OpenSearch

Para quem já usa OpenSearch (ou Elasticsearch) ou precisa de busca por texto completo.

from opensearchpy import OpenSearch

es = OpenSearch("http://localhost:9200")

# Criar índice com mapeamento para vetores
es.indices.create(
    index="meus_documentos",
    body={
        "mappings": {
            "properties": {
                "content": {"type": "text", "analyzer": "portuguese"},
                "embedding": {
                    "type": "dense_vector",
                    "dims": 1024,
                    "index": True,
                    "similarity": "cosine"
                },
                "metadata": {"type": "object"}
            }
        }
    }
)

Passo 2: Criar Embeddings com urso-embed

from openai import OpenAI

client = OpenAI(
    base_url="https://api.polarai.com.br/v1",
    api_key="pk-sua-chave-aqui"
)

def get_embedding(text: str) -> list[float]:
    response = client.embeddings.create(
        model="urso-embed",
        input=text
    )
    return response.data[0].embedding

# Exemplo
embedding = get_embedding("Direito do consumidor brasileiro")
print(f"Dimensões: {len(embedding)}")  # 1024

Embeddings em Lote

def get_embeddings_batch(texts: list[str]) -> list[list[float]]:
    response = client.embeddings.create(
        model="urso-embed",
        input=texts
    )
    return [item.embedding for item in response.data]

# Processar em lotes de 100
texts = ["texto 1", "texto 2", ..., "texto 1000"]
batch_size = 100

all_embeddings = []
for i in range(0, len(texts), batch_size):
    batch = texts[i:i + batch_size]
    embeddings = get_embeddings_batch(batch)
    all_embeddings.extend(embeddings)

Passo 3: Indexar Documentos

Chunking

Divida documentos grandes em chunks gerenciáveis:

def chunk_text(text: str, chunk_size: int = 512, overlap: int = 50) -> list[str]:
    """Divide texto em chunks com sobreposição."""
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        if chunk:
            chunks.append(chunk)
    return chunks

Indexar no Qdrant

from qdrant_client.models import PointStruct
import uuid

def index_document(doc_text: str, metadata: dict):
    chunks = chunk_text(doc_text)
    embeddings = get_embeddings_batch(chunks)

    points = [
        PointStruct(
            id=str(uuid.uuid4()),
            vector=embedding,
            payload={
                "content": chunk,
                "metadata": metadata
            }
        )
        for chunk, embedding in zip(chunks, embeddings)
    ]

    client.upsert(
        collection_name="meus_documentos",
        points=points
    )

# Indexar um documento
index_document(
    doc_text="Conteúdo completo do documento aqui...",
    metadata={"source": "manual.pdf", "category": "tributário"}
)

Indexar no Elasticsearch

def index_document_es(doc_text: str, metadata: dict):
    chunks = chunk_text(doc_text)
    embeddings = get_embeddings_batch(chunks)

    for chunk, embedding in zip(chunks, embeddings):
        es.index(
            index="meus_documentos",
            body={
                "content": chunk,
                "embedding": embedding,
                "metadata": metadata
            }
        )

Passo 4: Busca Híbrida (BM25 + Semântica)

A busca híbrida combina busca por palavras-chave (BM25) com busca semântica vetorial para resultados superiores:

Com Qdrant

from qdrant_client.models import Filter, FieldCondition, MatchValue

def hybrid_search(query: str, limit: int = 10) -> list:
    query_embedding = get_embedding(query)

    # Busca semântica
    results = client.search(
        collection_name="meus_documentos",
        query_vector=query_embedding,
        limit=limit,
        score_threshold=0.5
    )

    return [
        {
            "content": hit.payload["content"],
            "score": hit.score,
            "metadata": hit.payload["metadata"]
        }
        for hit in results
    ]

Com Elasticsearch (Busca Híbrida Nativa)

def hybrid_search_es(query: str, limit: int = 10) -> list:
    query_embedding = get_embedding(query)

    results = es.search(
        index="meus_documentos",
        body={
            "size": limit,
            "query": {
                "bool": {
                    "should": [
                        # BM25 (palavras-chave)
                        {
                            "match": {
                                "content": {
                                    "query": query,
                                    "boost": 0.3
                                }
                            }
                        },
                        # KNN (semântico)
                        {
                            "knn": {
                                "field": "embedding",
                                "query_vector": query_embedding,
                                "num_candidates": 50,
                                "boost": 0.7
                            }
                        }
                    ]
                }
            }
        }
    )

    return [
        {
            "content": hit["_source"]["content"],
            "score": hit["_score"],
            "metadata": hit["_source"]["metadata"]
        }
        for hit in results["hits"]["hits"]
    ]

Passo 5: Reranking

Após a busca, aplique reranking para refinar a ordem dos resultados:

def rerank(query: str, documents: list, top_k: int = 5) -> list:
    """Rerank usando o modelo urso como cross-encoder."""
    scored_docs = []

    for doc in documents:
        response = client.chat.completions.create(
            model="urso-mabe",
            messages=[
                {
                    "role": "system",
                    "content": "Avalie a relevância do documento para a consulta. "
                               "Responda apenas com um número de 0 a 10."
                },
                {
                    "role": "user",
                    "content": f"Consulta: {query}\n\nDocumento: {doc['content']}"
                }
            ],
            max_tokens=5,
            temperature=0
        )

        try:
            score = float(response.choices[0].message.content.strip())
            scored_docs.append({**doc, "rerank_score": score})
        except ValueError:
            scored_docs.append({**doc, "rerank_score": 0})

    scored_docs.sort(key=lambda x: x["rerank_score"], reverse=True)
    return scored_docs[:top_k]

Pipeline Completo

def rag_query(question: str) -> str:
    """Pipeline RAG completo: busca → rerank → geração."""

    # 1. Busca híbrida
    results = hybrid_search(question, limit=20)

    # 2. Reranking
    top_docs = rerank(question, results, top_k=5)

    # 3. Montar contexto
    context = "\n\n---\n\n".join([
        f"[Fonte: {doc['metadata'].get('source', 'desconhecida')}]\n{doc['content']}"
        for doc in top_docs
    ])

    # 4. Gerar resposta
    response = client.chat.completions.create(
        model="urso-mabe",
        messages=[
            {
                "role": "system",
                "content": (
                    "Responda a pergunta baseando-se APENAS no contexto fornecido. "
                    "Se a resposta não estiver no contexto, diga que não encontrou "
                    "informação relevante. Cite as fontes."
                )
            },
            {
                "role": "user",
                "content": f"Contexto:\n{context}\n\nPergunta: {question}"
            }
        ],
        temperature=0.3,
        max_tokens=2000
    )

    return response.choices[0].message.content

# Usar
answer = rag_query("Qual o prazo prescricional para danos morais?")
print(answer)

Dicas de Otimização

Chunk size: Use chunks de 256-512 tokens para melhor precisão na busca
Overlap: 10-20% de sobreposição entre chunks para manter contexto
Metadata: Adicione metadados ricos (fonte, data, categoria) para filtros
Batch embeddings: Processe embeddings em lotes para melhor performance
Cache: Armazene embeddings de consultas frequentes
Avaliação: Teste com perguntas reais para calibrar pesos de busca híbrida

On this page