Busca híbrida num único SQLite

Quando montei a busca de um chatbot de documentação, o caminho óbvio era subir um banco vetorial separado. Preferi deixar tudo num SQLite no Turso: documentos, vetores e índice de texto no mesmo banco, e cada pergunta custa uma ida só à rede.

O esquema

São três tabelas:

CREATE TABLE documents (
  doc_id  TEXT NOT NULL,
  locale  TEXT NOT NULL,
  public  INTEGER NOT NULL CHECK (public IN (0, 1)),
  roles   TEXT NOT NULL CHECK (json_type(roles) = 'array'),
  PRIMARY KEY (doc_id, locale)
);

CREATE TABLE chunks (
  chunk_id TEXT NOT NULL PRIMARY KEY,
  doc_id   TEXT NOT NULL,
  locale   TEXT NOT NULL,
  content  TEXT NOT NULL,
  emb      F32_BLOB(768) NOT NULL,
  FOREIGN KEY (doc_id, locale) REFERENCES documents (doc_id, locale) ON DELETE CASCADE
);

CREATE VIRTUAL TABLE chunks_fts USING fts5(
  chunk_id UNINDEXED,
  text,
  tokenize = "unicode61 remove_diacritics 2"
);

O controle de acesso (public, roles) fica só em documents. Numa versão anterior essa flag era copiada pra cada trecho, e uma cópia desatualizada acabou vazando um documento restrito nas duas buscas. Hoje as duas fazem JOIN em documents e filtram ali.

As duas buscas

A busca vetorial ordena pela distância de cosseno, sem índice ANN. Com 808 trechos, varrer tudo levou 3,8 ms com vetores de 768 dimensões (12,5 ms com 3072). Índice aproximado só compensa quando a varredura começa a pesar.

SELECT c.chunk_id, vector_distance_cos(c.emb, vector32(?)) AS score
  FROM chunks c
  JOIN documents d ON d.doc_id = c.doc_id AND d.locale = c.locale
 WHERE (d.public = 1 OR EXISTS (SELECT 1 FROM json_each(d.roles) WHERE value IN (?)))
 ORDER BY score, c.chunk_id
 LIMIT ?;

A busca por texto usa o FTS5 com bm25(), com o mesmo JOIN e o mesmo filtro:

SELECT c.chunk_id, bm25(chunks_fts) AS score
  FROM chunks_fts
  JOIN chunks c ON c.chunk_id = chunks_fts.chunk_id
  JOIN documents d ON d.doc_id = c.doc_id AND d.locale = c.locale
 WHERE chunks_fts MATCH ?
   AND (d.public = 1 OR EXISTS (SELECT 1 FROM json_each(d.roles) WHERE value IN (?)))
 ORDER BY score, c.chunk_id
 LIMIT ?;

Apanhei de duas coisas aqui. A primeira é que o filtro tem que vir antes do LIMIT, na mesma consulta. Do jeito errado (paginando numa subconsulta e filtrando depois), eu pedia 10 linhas e vinham 7, sem erro nenhum. A segunda é o chunk_id no ORDER BY: o SQLite pode devolver notas iguais em qualquer ordem, e a fusão trabalha com posição, então sem desempate o resultado pode mudar de uma execução pra outra.

Uma ida só à rede

As duas consultas vão no mesmo batch(), em modo leitura:

const [dense, lexical] = await client.batch(
  [
    { sql: denseSql, args: [vector, ...aclArgs, 40] },
    { sql: bm25Sql, args: [match, ...aclArgs, 40] },
  ],
  "read",
);

No banco real, o batch() levou 63,1 ms, contra 122,5 ms de duas chamadas execute(). Com o banco remoto, a ida à rede é metade do tempo da consulta. E o modo "read" ainda faz um DELETE colocado ali por engano falhar em vez de rodar.

Antes de mandar, o código confere se as duas consultas têm o filtro de acesso. Uma consulta sem filtro devolve linhas que parecem normais, e o vazamento passa batido.

Juntando os resultados

Não dá pra comparar as notas das duas buscas: o cosseno vai de 0 a 2, e o bm25() é um número negativo que depende do corpus. Por isso a fusão usa a posição de cada item, com Reciprocal Rank Fusion:

export function fuseByRrf(legs: string[][], k = 60) {
  const scores = new Map<string, number>();
  for (const leg of legs) {
    const seen = new Set<string>();
    leg.forEach((id, i) => {
      if (seen.has(id)) return; // conta uma vez, na melhor posição
      seen.add(id);
      scores.set(id, (scores.get(id) ?? 0) + 1 / (k + i + 1));
    });
  }
  return [...scores].sort((a, b) => b[1] - a[1]);
}

Cada busca traz 40 candidatos, e saem os 20 melhores da fusão.

O que o benchmark mostrou

Antes de sair defendendo o híbrido, fui medir. Rodei um benchmark público parecido com documentação técnica (CQADupStack, 13.145 perguntas). Com 100 resultados, a busca vetorial e a híbrida empataram dentro da margem de erro. Com 50, a vetorial foi melhor. O BM25 sozinho ficou de 12 a 20 pontos atrás.

Nesse corpus, a busca por texto não trouxe recall nenhum. Ela continua no sistema por causa de nome de API, código de erro e frase exata, que o vetor trata mal, e isso é assunto do post sobre Turso FTS no RAG. Pelo menos agora eu sei quanto ela vale.