Chunking e Embeddings: as decisões que fazem ou quebram seu RAG

# Para cada chunk, gerar um contexto curto usando LLM barato
from anthropic import Anthropic

client = Anthropic()

def contextualize(doc: str, chunk: str) -> str:
    prompt = f"""<document>{doc}</document>
Here is a chunk from the document:
<chunk>{chunk}</chunk>
Provide a short (50-100 tokens) context situating this chunk within the document.
Output only the context, nothing else."""
    r = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=120,
        messages=[{"role": "user", "content": prompt}],
    )
    return r.content[0].text.strip()

# Ingest
for chunk in chunks:
    ctx = contextualize(full_doc, chunk)
    embedded_text = f"{ctx}\n\n{chunk}"  # embeda o combinado
    vector = embedder.encode(embedded_text)
    db.insert(chunk=chunk, ctx=ctx, vector=vector)
# Prompt caching no doc full faz o custo cair ~10x

# OpenAI: pedir dimensão menor na API
embedding = client.embeddings.create(
    model="text-embedding-3-large",
    input="texto",
    dimensions=768,  # default 3072, mas 768 mantém ~95% da quality
).data[0].embedding

# Depois de obter, re-normalizar se for cosine
import numpy as np
v = np.array(embedding)
v_norm = v / np.linalg.norm(v)

from langchain_text_splitters import RecursiveCharacterTextSplitter
import psycopg
from openai import OpenAI

oai = OpenAI()
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,            # em tokens aproximados
    chunk_overlap=75,
    separators=["\n\n", "\n", ". ", " ", ""],
)

def embed(text: str) -> list[float]:
    return oai.embeddings.create(
        model="text-embedding-3-small",
        input=text,
        dimensions=768,
    ).data[0].embedding

conn = psycopg.connect("postgresql://...")
with conn.cursor() as cur:
    cur.execute("""
        CREATE TABLE IF NOT EXISTS chunks (
            id BIGSERIAL PRIMARY KEY,
            doc_id TEXT NOT NULL,
            section TEXT,
            chunk TEXT NOT NULL,
            embedding VECTOR(768) NOT NULL,
            created_at TIMESTAMPTZ DEFAULT now()
        );
        CREATE INDEX IF NOT EXISTS chunks_emb_idx
          ON chunks USING hnsw (embedding vector_cosine_ops);
    """)

def ingest(doc_id: str, section: str, text: str) -> None:
    for piece in splitter.split_text(text):
        vec = embed(piece)
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO chunks(doc_id,section,chunk,embedding) VALUES (%s,%s,%s,%s)",
                (doc_id, section, piece, vec),
            )
    conn.commit()