Embeddings e Busca Semântica: Como Funciona na Prática

import OpenAI from 'openai';

const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });

async function embed(text: string): Promise<number[]> {
  const response = await openai.embeddings.create({
    model: 'text-embedding-3-small',  // 1536 dims, barato (~$0.02/1M tokens)
    input: text,
    encoding_format: 'float',
  });
  return response.data[0].embedding;
}

// Batch para múltiplos textos (mais eficiente)
async function embedBatch(texts: string[]): Promise<number[][]> {
  const response = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: texts,
  });
  return response.data.map(d => d.embedding);
}

-- Habilita a extensão
CREATE EXTENSION IF NOT EXISTS vector;

-- Tabela de documentos com campo de embedding
CREATE TABLE documents (
  id          SERIAL PRIMARY KEY,
  title       TEXT NOT NULL,
  content     TEXT NOT NULL,
  embedding   vector(1536),   -- dimensão do modelo
  metadata    JSONB,
  created_at  TIMESTAMP DEFAULT NOW()
);

-- Índice HNSW para busca aproximada (mais rápido, levemente menos preciso)
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Ou IVFFlat para datasets maiores
-- CREATE INDEX ON documents USING ivfflat (embedding vector_cosine_ops)
-- WITH (lists = 100);

import { Pool } from 'pg';

const db = new Pool({ connectionString: process.env.DATABASE_URL });

async function indexDocument(doc: {
  title: string;
  content: string;
  metadata?: Record<string, unknown>;
}): Promise<void> {
  // Chunking: divide o conteúdo em partes menores
  const chunks = chunkText(doc.content, { maxTokens: 500, overlap: 50 });

  for (const chunk of chunks) {
    const embedding = await embed(`${doc.title}\n\n${chunk}`);

    await db.query(
      `INSERT INTO documents (title, content, embedding, metadata)
       VALUES ($1, $2, $3::vector, $4)`,
      [doc.title, chunk, JSON.stringify(embedding), doc.metadata ?? {}]
    );
  }
}

// Estratégia de chunking simples por tokens
function chunkText(text: string, { maxTokens = 500, overlap = 50 } = {}): string[] {
  const sentences = text.match(/[^.!?]+[.!?]+/g) ?? [text];
  const chunks: string[] = [];
  let current = '';
  let tokenCount = 0;

  for (const sentence of sentences) {
    const sentenceTokens = Math.ceil(sentence.length / 4); // aprox 4 chars/token
    if (tokenCount + sentenceTokens > maxTokens && current) {
      chunks.push(current.trim());
      // Overlap: mantém últimas frases para contexto
      const words = current.split(' ');
      current = words.slice(-overlap).join(' ') + ' ';
      tokenCount = overlap;
    }
    current += sentence;
    tokenCount += sentenceTokens;
  }

  if (current.trim()) chunks.push(current.trim());
  return chunks;
}

async function semanticSearch(
  query: string,
  { limit = 5, threshold = 0.7 }: { limit?: number; threshold?: number } = {}
): Promise<SearchResult[]> {
  const queryEmbedding = await embed(query);

  const results = await db.query<SearchResult>(
    `SELECT
       id,
       title,
       content,
       metadata,
       1 - (embedding <=> $1::vector) AS similarity
     FROM documents
     WHERE 1 - (embedding <=> $1::vector) > $2
     ORDER BY embedding <=> $1::vector  -- operador de distância cosseno
     LIMIT $3`,
    [JSON.stringify(queryEmbedding), threshold, limit]
  );

  return results.rows;
}

// Uso
const results = await semanticSearch('como otimizar performance em Node.js');
// Retorna documentos sobre profiling, bottlenecks, V8, etc.
// mesmo sem as palavras exatas na query

async function hybridSearch(query: string, limit = 5): Promise<SearchResult[]> {
  const queryEmbedding = await embed(query);

  const results = await db.query(
    `WITH semantic AS (
       SELECT id, 1 - (embedding <=> $1::vector) AS sem_score
       FROM documents
       ORDER BY embedding <=> $1::vector
       LIMIT 20
     ),
     keyword AS (
       SELECT id,
              ts_rank(to_tsvector('portuguese', content),
                      plainto_tsquery('portuguese', $2)) AS kw_score
       FROM documents
       WHERE to_tsvector('portuguese', content) @@ plainto_tsquery('portuguese', $2)
       LIMIT 20
     )
     SELECT
       d.id, d.title, d.content,
       COALESCE(s.sem_score, 0) * 0.7 + COALESCE(k.kw_score, 0) * 0.3 AS score
     FROM documents d
     LEFT JOIN semantic s ON s.id = d.id
     LEFT JOIN keyword k ON k.id = d.id
     WHERE s.id IS NOT NULL OR k.id IS NOT NULL
     ORDER BY score DESC
     LIMIT $3`,
    [JSON.stringify(queryEmbedding), query, limit]
  );

  return results.rows;
}