Construindo um sistema RAG com Node.js e OpenAI

5 de novembro de 2024·Paulo de Paula

RAG (Retrieval-Augmented Generation) é a técnica mais prática para fazer LLMs responderem com base em seus próprios dados sem fine-tuning. A ideia é simples: antes de enviar a pergunta para o modelo, você busca os documentos relevantes e inclui no contexto.

Como funciona RAG

Usuário pergunta
    ↓
Gera embedding da pergunta
    ↓
Busca documentos similares (vetor search)
    ↓
Monta prompt com contexto + pergunta
    ↓
LLM responde baseado nos documentos
    ↓
Resposta para o usuário

Setup

npm init -y
npm install openai @supabase/supabase-js dotenv

Vamos usar o Supabase como banco vetorial (tem extensão pgvector nativa e é gratuito para começar).

1. Indexar documentos

// src/indexer.ts
import OpenAI from 'openai'
import { createClient } from '@supabase/supabase-js'

const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY })
const supabase = createClient(
  process.env.SUPABASE_URL!,
  process.env.SUPABASE_SERVICE_KEY!
)

async function gerarEmbedding(texto: string): Promise<number[]> {
  const response = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: texto,
  })
  return response.data[0].embedding
}

async function indexarDocumento(
  titulo: string,
  conteudo: string,
  metadados?: Record<string, any>
) {
  // Divide em chunks para documentos grandes
  const chunks = dividirEmChunks(conteudo, 500)

  for (const chunk of chunks) {
    const embedding = await gerarEmbedding(chunk)

    await supabase.from('documentos').insert({
      titulo,
      conteudo: chunk,
      embedding,
      metadados,
    })
  }
}

function dividirEmChunks(texto: string, tamanho: number): string[] {
  const palavras = texto.split(' ')
  const chunks: string[] = []

  for (let i = 0; i < palavras.length; i += tamanho) {
    chunks.push(palavras.slice(i, i + tamanho).join(' '))
  }

  return chunks
}

// Indexar seus documentos
await indexarDocumento(
  'Manual de produto',
  fs.readFileSync('docs/manual.txt', 'utf-8'),
  { fonte: 'manual', versao: '2.0' }
)

2. Tabela no Supabase

-- Ativa extensão pgvector
create extension if not exists vector;

-- Cria tabela de documentos
create table documentos (
  id uuid primary key default gen_random_uuid(),
  titulo text not null,
  conteudo text not null,
  embedding vector(1536),  -- dimensão do text-embedding-3-small
  metadados jsonb,
  criado_em timestamptz default now()
);

-- Índice para busca por similaridade (HNSW é mais rápido)
create index on documentos
  using hnsw (embedding vector_cosine_ops);

-- Função de busca
create or replace function buscar_documentos(
  query_embedding vector(1536),
  match_count int default 5,
  similarity_threshold float default 0.7
)
returns table (
  id uuid,
  titulo text,
  conteudo text,
  metadados jsonb,
  similarity float
)
language sql stable
as $$
  select
    id,
    titulo,
    conteudo,
    metadados,
    1 - (embedding <=> query_embedding) as similarity
  from documentos
  where 1 - (embedding <=> query_embedding) > similarity_threshold
  order by embedding <=> query_embedding
  limit match_count;
$$;

3. Query e resposta

// src/rag.ts
async function responder(pergunta: string): Promise<string> {
  // 1. Gera embedding da pergunta
  const perguntaEmbedding = await gerarEmbedding(pergunta)

  // 2. Busca documentos relevantes
  const { data: documentos } = await supabase.rpc('buscar_documentos', {
    query_embedding: perguntaEmbedding,
    match_count: 5,
    similarity_threshold: 0.7,
  })

  if (!documentos || documentos.length === 0) {
    return 'Não encontrei informações relevantes para responder sua pergunta.'
  }

  // 3. Monta contexto
  const contexto = documentos
    .map((doc, i) => `[Documento ${i + 1}] ${doc.titulo}\n${doc.conteudo}`)
    .join('\n\n')

  // 4. Chama o LLM com contexto
  const response = await openai.chat.completions.create({
    model: 'gpt-4o-mini',
    messages: [
      {
        role: 'system',
        content: `Você é um assistente útil. Responda APENAS com base nos documentos fornecidos.
Se a informação não estiver nos documentos, diga que não sabe.
Cite os documentos quando relevante.`,
      },
      {
        role: 'user',
        content: `Contexto:\n${contexto}\n\nPergunta: ${pergunta}`,
      },
    ],
    temperature: 0.2,  // mais baixo = mais factual
  })

  return response.choices[0].message.content ?? ''
}

4. API com Express

// src/server.ts
import express from 'express'

const app = express()
app.use(express.json())

app.post('/perguntar', async (req, res) => {
  const { pergunta } = req.body

  if (!pergunta) {
    return res.status(400).json({ erro: 'Pergunta é obrigatória' })
  }

  try {
    const resposta = await responder(pergunta)
    res.json({ resposta })
  } catch (error) {
    res.status(500).json({ erro: 'Erro ao processar pergunta' })
  }
})

app.listen(3000, () => console.log('RAG API rodando na porta 3000'))

Melhorias para produção

1. Cache de embeddings frequentes:

const cache = new Map<string, number[]>()

async function gerarEmbeddingComCache(texto: string) {
  if (cache.has(texto)) return cache.get(texto)!
  const embedding = await gerarEmbedding(texto)
  cache.set(texto, embedding)
  return embedding
}

2. Reranking dos resultados:

Depois de buscar por similaridade de vetor, use um modelo de reranking para refinar os resultados mais relevantes antes de passar ao LLM.

3. Streaming da resposta:

const stream = await openai.chat.completions.create({
  model: 'gpt-4o-mini',
  messages,
  stream: true,
})

for await (const chunk of stream) {
  const content = chunk.choices[0]?.delta?.content ?? ''
  res.write(content)
}
res.end()

Custo estimado

Para uso típico com text-embedding-3-small e gpt-4o-mini:

Indexar 1.000 documentos de 500 palavras: ~$0,02
1.000 perguntas por dia: ~$1-3/dia

RAG é uma das técnicas mais custo-eficientes para adicionar conhecimento específico ao domínio sem precisar de fine-tuning.

Fine-tuning de LLMs: Quando Vale a Pena (e Quando Não Vale)